Spaces:
Sleeping
Sleeping
File size: 10,501 Bytes
1e6a9db |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 |
"""OAuth and authentication routes for Hugging Face integration."""
from __future__ import annotations
import logging
import secrets
import time
from datetime import datetime, timezone
from typing import Optional
from urllib.parse import urlencode
import httpx
from fastapi import APIRouter, Depends, HTTPException, Query, Request
from fastapi.responses import RedirectResponse
from ...models.auth import TokenResponse
from ...models.user import HFProfile, User
from ...services.auth import AuthError, AuthService
from ...services.config import get_config
from ...services.seed import ensure_welcome_note
from ...services.vault import VaultService
from ..middleware import AuthContext, get_auth_context
logger = logging.getLogger(__name__)
router = APIRouter()
OAUTH_STATE_TTL_SECONDS = 300
oauth_states: dict[str, float] = {}
auth_service = AuthService()
def _create_oauth_state() -> str:
"""Generate a state token and store it with a timestamp."""
now = time.time()
# Garbage collect expired states
expired = [
state
for state, ts in oauth_states.items()
if now - ts > OAUTH_STATE_TTL_SECONDS
]
for state in expired:
oauth_states.pop(state, None)
state = secrets.token_urlsafe(32)
oauth_states[state] = now
return state
def _consume_oauth_state(state: str | None) -> None:
"""Validate and remove the state token; raise if invalid."""
if not state or state not in oauth_states:
raise HTTPException(status_code=400, detail="Invalid or expired OAuth state.")
# Remove to prevent reuse
del oauth_states[state]
def get_base_url(request: Request) -> str:
"""
Get the base URL for OAuth redirects.
Uses the actual request URL scheme and hostname from FastAPI's request.url.
HF Spaces doesn't set X-Forwarded-Host, but the 'host' header is correct.
"""
# Get scheme from X-Forwarded-Proto or request
forwarded_proto = request.headers.get("x-forwarded-proto")
scheme = forwarded_proto if forwarded_proto else str(request.url.scheme)
# Get hostname from request URL (this comes from the 'host' header)
hostname = str(request.url.hostname)
# Check for port (but HF Spaces uses standard 443 for HTTPS)
port = request.url.port
if port and port not in (80, 443):
base_url = f"{scheme}://{hostname}:{port}"
else:
base_url = f"{scheme}://{hostname}"
logger.info(
f"OAuth base URL detected: {base_url}",
extra={
"scheme": scheme,
"hostname": hostname,
"port": port,
"request_url": str(request.url),
},
)
return base_url
@router.get("/auth/login")
async def login(request: Request):
"""Redirect to Hugging Face OAuth authorization page."""
config = get_config()
if not config.hf_oauth_client_id:
raise HTTPException(
status_code=501,
detail="OAuth not configured. Set HF_OAUTH_CLIENT_ID and HF_OAUTH_CLIENT_SECRET environment variables.",
)
# Get base URL from request (handles HF Spaces proxy)
base_url = get_base_url(request)
redirect_uri = f"{base_url}/auth/callback"
state = _create_oauth_state()
# Construct HF OAuth URL
oauth_base = "https://huggingface.co/oauth/authorize"
params = {
"client_id": config.hf_oauth_client_id,
"redirect_uri": redirect_uri,
"scope": "openid profile email",
"response_type": "code",
"state": state,
}
auth_url = f"{oauth_base}?{urlencode(params)}"
logger.info(
"Initiating OAuth flow",
extra={
"redirect_uri": redirect_uri,
"auth_url": auth_url,
"client_id": config.hf_oauth_client_id[:8] + "...",
"state": state,
},
)
return RedirectResponse(url=auth_url, status_code=302)
@router.get("/auth/callback")
async def callback(
request: Request,
code: str = Query(..., description="OAuth authorization code"),
state: Optional[str] = Query(
None, description="State parameter for CSRF protection"
),
):
"""Handle OAuth callback from Hugging Face."""
config = get_config()
if not config.hf_oauth_client_id or not config.hf_oauth_client_secret:
raise HTTPException(status_code=501, detail="OAuth not configured")
# Get base URL from request (must match the one sent to HF)
base_url = get_base_url(request)
redirect_uri = f"{base_url}/auth/callback"
# Validate state token to prevent CSRF and replay attacks
_consume_oauth_state(state)
logger.info(
"OAuth callback received",
extra={
"redirect_uri": redirect_uri,
"state": state,
"code_length": len(code) if code else 0,
},
)
try:
# Exchange authorization code for access token
async with httpx.AsyncClient() as client:
token_response = await client.post(
"https://huggingface.co/oauth/token",
data={
"grant_type": "authorization_code",
"code": code,
"redirect_uri": redirect_uri,
"client_id": config.hf_oauth_client_id,
"client_secret": config.hf_oauth_client_secret,
},
)
if token_response.status_code != 200:
logger.error(f"Token exchange failed: {token_response.text}")
raise HTTPException(
status_code=400,
detail="Failed to exchange authorization code for token",
)
token_data = token_response.json()
access_token = token_data.get("access_token")
if not access_token:
raise HTTPException(
status_code=400, detail="No access token in response"
)
# Get user profile from HF
user_response = await client.get(
"https://huggingface.co/api/whoami-v2",
headers={"Authorization": f"Bearer {access_token}"},
)
if user_response.status_code != 200:
logger.error(f"User profile fetch failed: {user_response.text}")
raise HTTPException(
status_code=400, detail="Failed to fetch user profile"
)
user_data = user_response.json()
username = user_data.get("name")
email = user_data.get("email")
hf_id = user_data.get("id") or user_data.get("uid")
if not username and not hf_id:
raise HTTPException(
status_code=400, detail="No user identity in user profile"
)
# Prefer stable HF account id; fall back to username if missing
user_id = str(hf_id or username)
# Create JWT for our application
import jwt
from datetime import datetime, timedelta, timezone
# Ensure the user has an initialized vault with a welcome note
try:
created = ensure_welcome_note(user_id)
logger.info(
"Ensured welcome note for user",
extra={"user_id": user_id, "created": created},
)
except Exception as seed_exc:
logger.exception(
"Failed to seed welcome note for user",
extra={"user_id": user_id},
)
payload = {
"sub": user_id,
"username": username,
"email": email,
"exp": datetime.now(timezone.utc) + timedelta(days=7),
"iat": datetime.now(timezone.utc),
}
try:
jwt_secret = auth_service._require_secret()
except AuthError as exc:
raise HTTPException(status_code=exc.status_code, detail=exc.message)
jwt_token = jwt.encode(payload, jwt_secret, algorithm="HS256")
logger.info(
"OAuth successful",
extra={
"username": username,
"user_id": user_id,
"email": email,
},
)
# Redirect to frontend with token in URL hash
frontend_url = base_url
redirect_url = f"{frontend_url}/#token={jwt_token}"
logger.info(f"Redirecting to frontend: {redirect_url}")
return RedirectResponse(url=redirect_url, status_code=302)
except httpx.HTTPError as e:
logger.exception(f"HTTP error during OAuth: {e}")
raise HTTPException(
status_code=500, detail="OAuth flow failed due to network error"
)
except Exception as e:
logger.exception(f"Unexpected error during OAuth: {e}")
raise HTTPException(status_code=500, detail="OAuth flow failed")
@router.post("/api/tokens", response_model=TokenResponse)
async def create_api_token(auth: AuthContext = Depends(get_auth_context)):
"""Issue a new JWT for the authenticated user."""
token, expires_at = auth_service.issue_token_response(auth.user_id)
return TokenResponse(token=token, token_type="bearer", expires_at=expires_at)
@router.get("/api/me", response_model=User)
async def get_current_user(auth: AuthContext = Depends(get_auth_context)):
"""Return profile metadata for the authenticated user."""
user_id = auth.user_id
vault_service = VaultService()
vault_path = vault_service.initialize_vault(user_id)
# Attempt to derive a stable "created" timestamp from the vault directory
try:
stat = vault_path.stat()
created_dt = datetime.fromtimestamp(stat.st_ctime, tz=timezone.utc)
except Exception:
created_dt = datetime.now(timezone.utc)
profile: Optional[HFProfile] = None
if user_id.startswith("hf-"):
username = user_id[len("hf-") :]
profile = HFProfile(
username=username,
name=username.replace("-", " ").title(),
avatar_url=f"https://api.dicebear.com/7.x/initials/svg?seed={username}",
)
elif user_id not in {"local-dev", "demo-user"}:
profile = HFProfile(username=user_id)
return User(
user_id=user_id,
hf_profile=profile,
vault_path=str(vault_path),
created=created_dt,
)
__all__ = ["router"]
|