diff --git a/labelu/internal/common/config.py b/labelu/internal/common/config.py index fd696627..553b5d5a 100644 --- a/labelu/internal/common/config.py +++ b/labelu/internal/common/config.py @@ -60,6 +60,11 @@ class Settings(BaseSettings): TOKEN_GENERATE_ALGORITHM: str = "HS256" TOKEN_ACCESS_EXPIRE_MINUTES: int = 30 + # Sliding refresh: when an authenticated request arrives with a token whose + # remaining lifetime is below this many minutes, a freshly issued token is + # returned via the `X-New-Token` response header so active users never get + # logged out. Should be smaller than TOKEN_ACCESS_EXPIRE_MINUTES. + TOKEN_REFRESH_THRESHOLD_MINUTES: int = 15 TOKEN_TYPE: str = "Bearer" EXPOSE_INTERNAL_ERRORS: bool = False diff --git a/labelu/internal/common/security.py b/labelu/internal/common/security.py index 8414848d..ec2f40a1 100644 --- a/labelu/internal/common/security.py +++ b/labelu/internal/common/security.py @@ -50,3 +50,19 @@ def create_access_token( ) return encoded_jwt + + +def should_refresh_token(exp_timestamp: Union[int, float, None]) -> bool: + """Decide whether a still-valid token is close enough to expiry that it + should be transparently re-issued (sliding refresh). + + ``exp_timestamp`` is the raw ``exp`` claim (seconds since epoch) taken + straight from the decoded JWT payload, which avoids any timezone ambiguity + from re-parsing into datetime objects. + """ + if not exp_timestamp: + return False + now = datetime.now(timezone.utc).timestamp() + remaining_seconds = exp_timestamp - now + threshold_seconds = settings.TOKEN_REFRESH_THRESHOLD_MINUTES * 60 + return 0 < remaining_seconds <= threshold_seconds diff --git a/labelu/internal/dependencies/user.py b/labelu/internal/dependencies/user.py index c06df257..470fe7a9 100644 --- a/labelu/internal/dependencies/user.py +++ b/labelu/internal/dependencies/user.py @@ -1,15 +1,18 @@ from typing import Optional +from datetime import timedelta from jose import jwt from loguru import logger from sqlalchemy.orm import Session from pydantic import ValidationError -from fastapi import HTTPException, WebSocket, status, Depends, Request +from fastapi import HTTPException, WebSocket, status, Depends, Request, Response from labelu.internal.domain.models.user import User from labelu.internal.common import db as db_module from labelu.internal.common.config import settings from labelu.internal.common.security import AccessToken +from labelu.internal.common.security import create_access_token +from labelu.internal.common.security import should_refresh_token from labelu.internal.common.error_code import ErrorCode from labelu.internal.common.error_code import LabelUException from labelu.internal.adapter.persistence import crud_user @@ -31,7 +34,9 @@ def __call__(self, request: Request) -> Optional[str]: def get_current_user( - db: Session = Depends(db_module.get_db), token: str = Depends(reusable_oauth2) + response: Response, + db: Session = Depends(db_module.get_db), + token: str = Depends(reusable_oauth2), ) -> User: try: payload = jwt.decode( @@ -49,6 +54,17 @@ def get_current_user( user = crud_user.get(db, id=token_data.id) if not user: raise LabelUException(code=ErrorCode.CODE_40002_USER_NOT_FOUND, status_code=401) + + # Sliding refresh: hand back a freshly issued token (same `Bearer xxx` + # format the login endpoint returns) when the current one is close to + # expiry, so active users are never logged out mid-session. + if should_refresh_token(payload.get("exp")): + new_token = create_access_token( + token=AccessToken(id=user.id, username=user.username), + expires_delta=timedelta(minutes=settings.TOKEN_ACCESS_EXPIRE_MINUTES), + ) + response.headers["X-New-Token"] = f"{settings.TOKEN_TYPE} {new_token}" + return user diff --git a/labelu/tests/internal/adapter/routers/test_user.py b/labelu/tests/internal/adapter/routers/test_user.py index 844fcb08..fbb3ed91 100644 --- a/labelu/tests/internal/adapter/routers/test_user.py +++ b/labelu/tests/internal/adapter/routers/test_user.py @@ -11,6 +11,7 @@ from labelu.internal.common.security import create_access_token from labelu.internal.common.db import begin_transaction +from labelu.tests.conftest import TEST_USERNAME from labelu.tests.utils.utils import random_username from labelu.tests.utils.utils import random_lower_string @@ -132,3 +133,44 @@ def test_user_not_found(self, client: TestClient, db: Session) -> None: # check assert r.status_code == 401 assert r.json()["err_code"] == 40002 + + def test_token_sliding_refresh_when_near_expiry( + self, client: TestClient, db: Session + ) -> None: + # prepare: a still-valid token whose remaining lifetime is below the + # refresh threshold, so the request should trigger a sliding refresh. + user = crud_user.get_user_by_username(db, username=TEST_USERNAME) + remaining = settings.TOKEN_REFRESH_THRESHOLD_MINUTES - 1 + access_token = create_access_token( + token=AccessToken(id=user.id, username=user.username), + expires_delta=timedelta(minutes=remaining), + ) + headers = {"Authorization": f"Bearer {access_token}"} + + # run + r = client.get(f"{settings.API_V1_STR}/users/me", headers=headers) + + # check: a brand-new token is handed back via the response header + assert r.status_code == 200 + assert "X-New-Token" in r.headers + assert r.headers["X-New-Token"].startswith(settings.TOKEN_TYPE) + assert r.headers["X-New-Token"] != f"Bearer {access_token}" + + def test_token_not_refreshed_when_fresh( + self, client: TestClient, db: Session + ) -> None: + # prepare: a token whose remaining lifetime is above the threshold + user = crud_user.get_user_by_username(db, username=TEST_USERNAME) + remaining = settings.TOKEN_REFRESH_THRESHOLD_MINUTES + 5 + access_token = create_access_token( + token=AccessToken(id=user.id, username=user.username), + expires_delta=timedelta(minutes=remaining), + ) + headers = {"Authorization": f"Bearer {access_token}"} + + # run + r = client.get(f"{settings.API_V1_STR}/users/me", headers=headers) + + # check: no refresh header on a fresh token + assert r.status_code == 200 + assert "X-New-Token" not in r.headers