Spaces:
Running
Running
File size: 4,710 Bytes
88e4a42 | 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 | """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
|