gpu_capacity/auth.py

"""Sign-in with Hugging Face (OAuth) and probe-rights gating.

The Space metadata sets `hf_oauth_authorized_org: huggingface`, so only org
members can complete the sign-in. The session is a signed cookie; probe
endpoints require it. Without OAUTH_CLIENT_ID (local dev), gating is disabled.
"""

import base64
import hashlib
import hmac
import json
import logging
import os
import secrets
import time
import urllib.parse
import urllib.request

from fastapi import APIRouter, HTTPException, Request
from fastapi.responses import RedirectResponse

log = logging.getLogger(__name__)
router = APIRouter()

CLIENT_ID = os.environ.get("OAUTH_CLIENT_ID", "")
CLIENT_SECRET = os.environ.get("OAUTH_CLIENT_SECRET", "")
PROVIDER = os.environ.get("OPENID_PROVIDER_URL", "https://huggingface.co")
SPACE_HOST = os.environ.get("SPACE_HOST", "")
SESSION_TTL = 8 * 3600
HF_ORG = "huggingface"

ENABLED = bool(CLIENT_ID)
_secret = (CLIENT_SECRET or secrets.token_hex(16)).encode()


def _sign(raw: str) -> str:
    return hmac.new(_secret, raw.encode(), hashlib.sha256).hexdigest()


def _make_cookie(payload: dict) -> str:
    raw = base64.urlsafe_b64encode(json.dumps(payload).encode()).decode()
    return raw + "." + _sign(raw)


def _read_cookie(value: str) -> dict | None:
    try:
        raw, sig = value.rsplit(".", 1)
        if not hmac.compare_digest(sig, _sign(raw)):
            return None
        payload = json.loads(base64.urlsafe_b64decode(raw))
        if payload.get("exp", 0) < time.time():
            return None
        return payload
    except Exception:
        return None


def current_user(request: Request) -> dict | None:
    value = request.cookies.get("session")
    return _read_cookie(value) if value else None


def require_probe_rights(request: Request):
    if not ENABLED:
        return
    user = current_user(request)
    if user is None:
        raise HTTPException(401, "sign in with Hugging Face to probe")
    if not user.get("hf"):
        raise HTTPException(403, "probing is restricted to Hugging Face members")


def _redirect_uri(request: Request) -> str:
    if SPACE_HOST:
        return f"https://{SPACE_HOST}/auth/callback"
    return str(request.base_url) + "auth/callback"


@router.get("/login/huggingface")
def login(request: Request):
    state = secrets.token_urlsafe(16)
    query = urllib.parse.urlencode({
        "redirect_uri": _redirect_uri(request),
        "scope": "openid profile email",
        "client_id": CLIENT_ID,
        "state": state,
        "response_type": "code",
    })
    resp = RedirectResponse(f"{PROVIDER}/oauth/authorize?{query}")
    resp.set_cookie("oauth_state", state, max_age=600,
                    httponly=True, secure=True, samesite="lax")
    return resp


@router.get("/auth/callback")
def callback(request: Request, code: str, state: str):
    if state != request.cookies.get("oauth_state"):
        raise HTTPException(400, "state mismatch")

    basic = base64.b64encode(f"{CLIENT_ID}:{CLIENT_SECRET}".encode()).decode()
    data = urllib.parse.urlencode({
        "client_id": CLIENT_ID,
        "code": code,
        "grant_type": "authorization_code",
        "redirect_uri": _redirect_uri(request),
    }).encode()
    req = urllib.request.Request(
        f"{PROVIDER}/oauth/token", data=data,
        headers={"Authorization": f"Basic {basic}",
                 "Content-Type": "application/x-www-form-urlencoded"},
    )
    token = json.load(urllib.request.urlopen(req, timeout=30))

    req = urllib.request.Request(
        f"{PROVIDER}/oauth/userinfo",
        headers={"Authorization": f"Bearer {token['access_token']}"},
    )
    info = json.load(urllib.request.urlopen(req, timeout=30))

    orgs = info.get("orgs") or info.get("organizations") or []
    org_names = {o.get("preferred_username") or o.get("name")
                 for o in orgs if isinstance(o, dict)}
    email = info.get("email") or ""
    # hf_oauth_authorized_org already gates the sign-in to org members; the
    # org/email checks are a second factor in case that metadata is removed.
    is_hf = (HF_ORG in org_names
             or bool(info.get("email_verified") and email.endswith("@huggingface.co")))
    username = info.get("preferred_username", "?")
    log.info("sign-in: %s (hf member: %s)", username, is_hf)

    resp = RedirectResponse("/")
    resp.set_cookie(
        "session",
        _make_cookie({"u": username, "hf": is_hf, "exp": time.time() + SESSION_TTL}),
        max_age=SESSION_TTL, httponly=True, secure=True, samesite="lax",
    )
    resp.delete_cookie("oauth_state")
    return resp


@router.get("/logout")
def logout():
    resp = RedirectResponse("/")
    resp.delete_cookie("session")
    return resp