feat: persist authentication sessions in postgres

src/pk_lab_server/api.py

@@ -2,13 +2,13 @@ import hashlib
 import secrets
 from datetime import UTC, datetime, timedelta
 
-from fastapi import APIRouter, Depends, HTTPException, Query
+from fastapi import APIRouter, Depends, HTTPException, Query, Response
 from sqlalchemy import func, select, update
 from sqlalchemy.orm import Session, selectinload
 
 from .config import get_settings
 from .db import get_db
-from .models import AppIdentity, Candidate, LinkCode, TestDefinition, TestRun, User, now_utc
+from .models import AppIdentity, AuthSession, LinkCode, TestDefinition, TestRun, User, now_utc
 from .schemas import (
     CompleteRunRequest,
     LinkConfirmIn,
@@ -20,10 +20,9 @@ from .schemas import (
     TestCard,
     TestDetail,
 )
-from .security import TokenClaims, current_claims, current_user, issue_token
+from .security import SessionClaims, current_claims, current_user, issue_session
 from .xhs_auth import XhsAuthError, exchange_code
 
-
 router = APIRouter(prefix="/api/v1")
 
 
@@ -54,15 +53,27 @@ async def login(body: LoginRequest, db: Session = Depends(get_db)) -> LoginRespo
         user = db.get(User, identity.user_id)
         if user is None:
             raise HTTPException(status_code=500, detail="identity references a missing user")
+    token = issue_session(db, user.id, body.app_key)
     db.commit()
     return LoginResponse(
-        token=issue_token(user.id, body.app_key),
+        token=token,
         user_id=user.id,
         app_key=body.app_key,
         linked_apps=_linked_apps(db, user.id),
     )
 
 
+@router.post("/auth/logout", status_code=204)
+def logout(
+    claims: SessionClaims = Depends(current_claims), db: Session = Depends(get_db)
+) -> Response:
+    session = db.get(AuthSession, claims.session_id)
+    if session is not None:
+        session.revoked_at = now_utc()
+        db.commit()
+    return Response(status_code=204)
+
+
 @router.get("/tests", response_model=list[TestCard])
 def list_tests(db: Session = Depends(get_db)) -> list[TestCard]:
     tests = db.scalars(
@@ -130,7 +141,7 @@ def complete_run(
     slug: str,
     body: CompleteRunRequest,
     user: User = Depends(current_user),
-    claims: TokenClaims = Depends(current_claims),
+    claims: SessionClaims = Depends(current_claims),
     db: Session = Depends(get_db),
 ) -> RunOut:
     if body.app_key != claims.app_key:
@@ -232,7 +243,7 @@ def create_link_code(user: User = Depends(current_user), db: Session = Depends(g
 def confirm_link_code(
     body: LinkConfirmIn,
     user: User = Depends(current_user),
-    claims: TokenClaims = Depends(current_claims),
+    claims: SessionClaims = Depends(current_claims),
     db: Session = Depends(get_db),
 ) -> LinkConfirmOut:
     link = db.scalar(
@@ -250,11 +261,18 @@ def confirm_link_code(
             raise HTTPException(status_code=409, detail="both users already contain the same app identity")
         db.execute(update(AppIdentity).where(AppIdentity.user_id == user.id).values(user_id=target_id))
         db.execute(update(TestRun).where(TestRun.user_id == user.id).values(user_id=target_id))
+        db.execute(update(AuthSession).where(AuthSession.user_id == user.id).values(user_id=target_id))
+        db.execute(
+            update(AuthSession)
+            .where(AuthSession.id == claims.session_id)
+            .values(revoked_at=now_utc())
+        )
         db.delete(user)
     link.used_at = now_utc()
+    token = issue_session(db, target_id, claims.app_key)
     db.commit()
     return LinkConfirmOut(
-        token=issue_token(target_id, claims.app_key),
+        token=token,
         user_id=target_id,
         linked_apps=_linked_apps(db, target_id),
     )

src/pk_lab_server/config.py

@@ -8,9 +8,8 @@ class Settings(BaseSettings):
     model_config = SettingsConfigDict(env_file=".env", extra="ignore")
 
     database_url: str = "sqlite:///./pk-lab.db"
-    token_secret: str = "development-only-change-me"
-    token_ttl_seconds: int = 60 * 60 * 24 * 30
-    cors_origins: str = "http://localhost:3001,http://localhost:3002"
+    session_ttl_seconds: int = 60 * 60 * 24 * 30
+    cors_origins: str = "http://localhost:3000,http://localhost:3002"
     public_base_url: str = "http://localhost:18082"
     test_auth_enabled: bool = False
     xhs_token_url: str = "https://miniapp.xiaohongshu.com/api/rmp/token"
@@ -35,4 +34,3 @@ class Settings(BaseSettings):
 @lru_cache
 def get_settings() -> Settings:
     return Settings()
-

src/pk_lab_server/models.py

@@ -1,7 +1,17 @@
 import uuid
 from datetime import UTC, datetime
 
-from sqlalchemy import JSON, Boolean, DateTime, ForeignKey, Index, Integer, String, Text, UniqueConstraint
+from sqlalchemy import (
+    JSON,
+    Boolean,
+    DateTime,
+    ForeignKey,
+    Index,
+    Integer,
+    String,
+    Text,
+    UniqueConstraint,
+)
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
 from .db import Base
@@ -39,6 +49,20 @@ class AppIdentity(Base):
     user: Mapped[User] = relationship()
 
 
+class AuthSession(Base):
+    __tablename__ = "auth_sessions"
+
+    id: Mapped[str] = mapped_column(String(36), primary_key=True, default=uuid_str)
+    token_hash: Mapped[str] = mapped_column(String(64), unique=True, index=True)
+    user_id: Mapped[str] = mapped_column(ForeignKey("users.id", ondelete="CASCADE"), index=True)
+    app_key: Mapped[str] = mapped_column(String(20), index=True)
+    expires_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), index=True)
+    revoked_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))
+    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=now_utc)
+
+    user: Mapped[User] = relationship()
+
+
 class TestDefinition(Base):
     __tablename__ = "tests"
 

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: