feat: persist authentication sessions in postgres

src/pk_lab_server/security.py

@@ -1,84 +1,68 @@
-import base64
-import binascii
 import hashlib
-import hmac
-import json
 import secrets
-import time
 from dataclasses import dataclass
+from datetime import UTC, datetime, timedelta
 
 from fastapi import Depends, HTTPException
 from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
+from sqlalchemy import select
 from sqlalchemy.orm import Session
 
 from .config import get_settings
 from .db import get_db
-from .models import User
-
+from .models import AuthSession, User
 
 bearer = HTTPBearer(auto_error=False)
 
 
-def _b64(data: bytes) -> str:
-    return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")
+def _token_hash(token: str) -> str:
+    return hashlib.sha256(token.encode()).hexdigest()
 
 
-def _unb64(value: str) -> bytes:
-    return base64.urlsafe_b64decode(value + "=" * (-len(value) % 4))
-
-
-def issue_token(user_id: str, app_key: str) -> str:
-    settings = get_settings()
-    payload = {
-        "sub": user_id,
-        "app": app_key,
-        "iat": int(time.time()),
-        "exp": int(time.time()) + settings.token_ttl_seconds,
-        "nonce": secrets.token_hex(8),
-    }
-    encoded = _b64(json.dumps(payload, separators=(",", ":"), sort_keys=True).encode())
-    signature = hmac.new(settings.token_secret.encode(), encoded.encode(), hashlib.sha256).digest()
-    return encoded + "." + _b64(signature)
+def issue_session(db: Session, user_id: str, app_key: str) -> str:
+    token = secrets.token_urlsafe(48)
+    db.add(
+        AuthSession(
+            token_hash=_token_hash(token),
+            user_id=user_id,
+            app_key=app_key,
+            expires_at=datetime.now(UTC) + timedelta(seconds=get_settings().session_ttl_seconds),
+        )
+    )
+    return token
 
 
 @dataclass(frozen=True)
-class TokenClaims:
+class SessionClaims:
+    session_id: str
     user_id: str
     app_key: str
 
 
-def parse_token(token: str) -> TokenClaims:
-    settings = get_settings()
-    try:
-        encoded, supplied = token.split(".", 1)
-        expected = hmac.new(settings.token_secret.encode(), encoded.encode(), hashlib.sha256).digest()
-        if not hmac.compare_digest(expected, _unb64(supplied)):
-            raise ValueError("signature")
-        payload = json.loads(_unb64(encoded))
-        if int(payload["exp"]) < int(time.time()):
-            raise ValueError("expired")
-        return TokenClaims(user_id=str(payload["sub"]), app_key=str(payload["app"]))
-    except (
-        KeyError,
-        ValueError,
-        TypeError,
-        UnicodeDecodeError,
-        binascii.Error,
-        json.JSONDecodeError,
-    ) as exc:
-        raise HTTPException(status_code=401, detail="invalid or expired token") from exc
-
-
 def current_claims(
     credentials: HTTPAuthorizationCredentials | None = Depends(bearer),
-) -> TokenClaims:
+    db: Session = Depends(get_db),
+) -> SessionClaims:
     if credentials is None or credentials.scheme.lower() != "bearer":
         raise HTTPException(status_code=401, detail="authentication required")
-    return parse_token(credentials.credentials)
+    session = db.scalar(
+        select(AuthSession).where(
+            AuthSession.token_hash == _token_hash(credentials.credentials),
+            AuthSession.revoked_at.is_(None),
+            AuthSession.expires_at > datetime.now(UTC),
+        )
+    )
+    if session is None:
+        raise HTTPException(status_code=401, detail="invalid or expired session")
+    return SessionClaims(
+        session_id=session.id,
+        user_id=session.user_id,
+        app_key=session.app_key,
+    )
 
 
 def current_user(
-    claims: TokenClaims = Depends(current_claims), db: Session = Depends(get_db)
+    claims: SessionClaims = Depends(current_claims), db: Session = Depends(get_db)
 ) -> User:
     user = db.get(User, claims.user_id)
     if user is None: