mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-21 02:05:45 +00:00
Ryker_FengandGitHub
a028dfd5fb
* feat(auth): add keep me signed in login option * fix(auth): address remember-login review feedback
123 lines
4.4 KiB
Python
123 lines
4.4 KiB
Python
"""OIDC state management via signed HttpOnly cookies.
|
|
|
|
Stores OIDC state, nonce, and PKCE verifier in a short-lived signed cookie
|
|
instead of server-side storage. This keeps the implementation stateless and
|
|
compatible with multi-worker deployments without Redis.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import secrets
|
|
import time
|
|
|
|
import jwt
|
|
from fastapi import Request, Response
|
|
from pydantic import BaseModel, Field
|
|
|
|
from app.gateway.auth.config import get_auth_config
|
|
from app.gateway.csrf_middleware import is_secure_request
|
|
|
|
OIDC_STATE_COOKIE_PREFIX = "df_oidc_state_"
|
|
OIDC_STATE_MAX_AGE = 300 # 5 minutes
|
|
OIDC_STATE_BYTES = 32
|
|
OIDC_NONCE_BYTES = 16
|
|
OIDC_CODE_VERIFIER_BYTES = 32
|
|
|
|
|
|
class OIDCStatePayload(BaseModel):
|
|
"""Payload stored inside the signed OIDC state cookie."""
|
|
|
|
provider: str = Field(description="OIDC provider ID (must match the state cookie)") # noqa: E501
|
|
state: str = Field(description="Cryptographically random state value — compared in constant time with the query param") # noqa: E501
|
|
nonce: str | None = Field(default=None, description="OIDC nonce, verified against the ID token nonce claim")
|
|
code_verifier: str | None = Field(default=None, description="PKCE code verifier, sent during token exchange")
|
|
next_path: str = Field(default="/workspace", description="Redirect target after successful auth")
|
|
remember_me: bool = Field(default=True, description="Whether the resulting DeerFlow session should be persistent")
|
|
issued_at: float = Field(default_factory=time.time, description="Unix timestamp of cookie creation")
|
|
|
|
|
|
def _sign_state_payload(payload: OIDCStatePayload) -> str:
|
|
"""Sign the state payload with the JWT secret to prevent tampering."""
|
|
secret = get_auth_config().jwt_secret
|
|
return jwt.encode(payload.model_dump(), secret, algorithm="HS256")
|
|
|
|
|
|
def _verify_state_signed(signed: str, max_age: int = OIDC_STATE_MAX_AGE) -> OIDCStatePayload | None:
|
|
"""Verify a signed state payload and return it, or None if invalid/expired."""
|
|
secret = get_auth_config().jwt_secret
|
|
try:
|
|
decoded = jwt.decode(signed, secret, algorithms=["HS256"])
|
|
payload = OIDCStatePayload(**decoded)
|
|
if time.time() - payload.issued_at > max_age:
|
|
return None
|
|
return payload
|
|
except jwt.PyJWTError:
|
|
return None
|
|
|
|
|
|
def generate_oidc_state() -> str:
|
|
"""Generate a cryptographically random state string."""
|
|
return secrets.token_urlsafe(OIDC_STATE_BYTES)
|
|
|
|
|
|
def generate_nonce() -> str:
|
|
"""Generate a cryptographically random nonce for ID token validation."""
|
|
return secrets.token_urlsafe(OIDC_NONCE_BYTES)
|
|
|
|
|
|
def generate_code_verifier() -> str:
|
|
"""Generate a PKCE code verifier (plain random string)."""
|
|
return secrets.token_urlsafe(OIDC_CODE_VERIFIER_BYTES)
|
|
|
|
|
|
def compute_code_challenge(verifier: str) -> str:
|
|
"""Compute the S256 PKCE code challenge from a verifier."""
|
|
import hashlib
|
|
|
|
return _base64url_encode(hashlib.sha256(verifier.encode("ascii")).digest())
|
|
|
|
|
|
def _base64url_encode(data: bytes) -> str:
|
|
"""Base64url-encode without padding, as required by RFC 7636 and OIDC."""
|
|
import base64
|
|
|
|
return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")
|
|
|
|
|
|
def _cookie_name(provider: str) -> str:
|
|
return f"{OIDC_STATE_COOKIE_PREFIX}{provider}"
|
|
|
|
|
|
def set_state_cookie(response: Response, request: Request, payload: OIDCStatePayload) -> None:
|
|
"""Set the signed OIDC state cookie on the response."""
|
|
signed = _sign_state_payload(payload)
|
|
is_https = is_secure_request(request)
|
|
response.set_cookie(
|
|
key=_cookie_name(payload.provider),
|
|
value=signed,
|
|
httponly=True,
|
|
secure=is_https,
|
|
samesite="lax",
|
|
max_age=OIDC_STATE_MAX_AGE,
|
|
path=f"/api/v1/auth/callback/{payload.provider}",
|
|
)
|
|
|
|
|
|
def get_state_cookie(request: Request, provider: str) -> OIDCStatePayload | None:
|
|
"""Read and verify the signed OIDC state cookie for the given provider."""
|
|
signed = request.cookies.get(_cookie_name(provider))
|
|
if not signed:
|
|
return None
|
|
return _verify_state_signed(signed)
|
|
|
|
|
|
def delete_state_cookie(response: Response, request: Request, provider: str) -> None:
|
|
"""Delete the OIDC state cookie."""
|
|
is_https = is_secure_request(request)
|
|
response.delete_cookie(
|
|
key=_cookie_name(provider),
|
|
secure=is_https,
|
|
samesite="lax",
|
|
path=f"/api/v1/auth/callback/{provider}",
|
|
)
|