unsloth/studio/backend/auth/jwt.py

69 lines
2.1 KiB
Python

import secrets
from datetime import UTC, datetime, timedelta
from typing import Optional
from fastapi import Depends, HTTPException, status
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from jose import JWTError, jwt
# Ephemeral in-memory secret:
# - Generated fresh on each backend process start
# - Never written to disk
# - Not configurable by the user
SECRET_KEY = secrets.token_urlsafe(64)
ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE_MINUTES = 60
security = HTTPBearer() # Reads Authorization: Bearer <token>
def create_access_token(
subject: str,
expires_delta: Optional[timedelta] = None,
) -> str:
"""
Create a signed JWT for the given subject (e.g. "local-user").
Tokens are valid only for the lifetime of this process, because the
SECRET_KEY is regenerated each time the backend restarts.
"""
to_encode = {"sub": subject}
expire = datetime.now(UTC) + (
expires_delta or timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
)
to_encode.update({"exp": expire})
return jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
async def get_current_subject(
credentials: HTTPAuthorizationCredentials = Depends(security),
) -> str:
"""
FastAPI dependency to validate the JWT and return the subject.
Use this as a dependency on routes that should be protected, e.g.:
@router.get("/secure")
async def secure_endpoint(current_subject: str = Depends(get_current_subject)):
...
"""
token = credentials.credentials
try:
payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
subject: Optional[str] = payload.get("sub")
if subject is None:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Invalid token payload",
)
return subject
except JWTError:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Invalid or expired token",
)
token = create_access_token("local-user")
print(token)