diff --git a/server.py b/server.py index 68e35aa..442321d 100644 --- a/server.py +++ b/server.py @@ -85,6 +85,26 @@ redis_client = redis.Redis( ) RUNS_KEY_PREFIX = "artdag:run:" RECIPES_KEY_PREFIX = "artdag:recipe:" +REVOKED_KEY_PREFIX = "artdag:revoked:" + +# Token revocation (30 day expiry to match token lifetime) +TOKEN_EXPIRY_SECONDS = 60 * 60 * 24 * 30 + + +def revoke_token(token: str) -> bool: + """Add token to revocation set. Returns True if newly revoked.""" + token_hash = hashlib.sha256(token.encode()).hexdigest() + key = f"{REVOKED_KEY_PREFIX}{token_hash}" + result = redis_client.set(key, "1", ex=TOKEN_EXPIRY_SECONDS, nx=True) + return result is not None + + +def is_token_revoked(token: str) -> bool: + """Check if token has been revoked.""" + token_hash = hashlib.sha256(token.encode()).hexdigest() + key = f"{REVOKED_KEY_PREFIX}{token_hash}" + return redis_client.exists(key) > 0 + # Initialize L1 cache manager with Redis for shared state between workers cache_manager = L1CacheManager(cache_dir=CACHE_DIR, redis_client=redis_client) @@ -427,6 +447,10 @@ async def verify_token_with_l2(token: str, l2_server: str) -> Optional[str]: async def get_verified_user_context(token: str) -> Optional[UserContext]: """Get verified user context from token. Verifies with the L2 that issued it.""" + # Check if token has been revoked + if is_token_revoked(token): + return None + ctx = get_user_context_from_token(token) if not ctx: return None @@ -3877,6 +3901,27 @@ async def logout(): return response +@app.post("/auth/revoke") +async def auth_revoke(credentials: HTTPAuthorizationCredentials = Depends(security)): + """ + Revoke a token. Called by L2 when user logs out. + The token to revoke is passed in the Authorization header. + """ + if not credentials: + raise HTTPException(401, "No token provided") + + token = credentials.credentials + + # Verify token is valid before revoking (ensures caller has the token) + ctx = get_user_context_from_token(token) + if not ctx: + raise HTTPException(401, "Invalid token") + + # Revoke the token + newly_revoked = revoke_token(token) + + return {"revoked": True, "newly_revoked": newly_revoked} + @app.post("/ui/publish-run/{run_id}", response_class=HTMLResponse) async def ui_publish_run(run_id: str, request: Request):