Spaces:
Running
Running
File size: 13,391 Bytes
2170658 bd469c1 2170658 722c296 2170658 722c296 2170658 722c296 2170658 722c296 2170658 bd469c1 2170658 | 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 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 | from __future__ import annotations
import asyncio
import time
from typing import Any, Dict, Optional
from urllib.parse import urlencode
import httpx
import jwt
from app.config import get_settings
from app.core.logger import get_logger
from app.core.thread_pool import run_in_executor
from app.models.schemas import GoogleOAuthUserInfo
from app.utils.http_utils import SharedAsyncClient
_logger = get_logger(__name__)
_settings = get_settings()
# Retryable transient HTTP status codes when talking to Google.
_RETRYABLE_STATUS = frozenset({429, 500, 502, 503, 504})
class GoogleOAuthError(Exception):
"""Raised for upstream Google OAuth failures that map to a client-facing error."""
def __init__(self, message: str, status_code: int = 400) -> None:
super().__init__(message)
self.message = message
self.status_code = status_code
class GoogleJWKSCache:
"""Asynchronously fetches and caches Google's public JWKS keys.
Keys are cached for a short TTL to avoid hammering Google's cert endpoint,
and re-fetched on demand if the requested ``kid`` is missing (key rotation).
"""
def __init__(self) -> None:
self._keys: Dict[str, Dict[str, Any]] = {}
self._last_fetched: float = 0.0
self._lock = asyncio.Lock()
def _is_expired(self) -> bool:
return (
not self._keys
or time.time() - self._last_fetched >= _settings.google_oauth_jwks_ttl_seconds
)
async def get_public_key(self, client: httpx.AsyncClient, kid: str) -> Any:
"""Return the RSA public key for the given key id (``kid``)."""
if self._is_expired():
await self._refresh(client)
jwk = self._keys.get(kid)
if jwk is None:
# Key may have rotated; force a refresh before giving up.
await self._refresh(client)
jwk = self._keys.get(kid)
if jwk is None:
raise GoogleOAuthError(
f"Google signing key not found for kid '{kid}'", status_code=401
)
return jwt.algorithms.RSAAlgorithm.from_jwk(jwk)
async def _refresh(self, client: httpx.AsyncClient) -> None:
async with self._lock:
if not self._is_expired():
return
try:
response = await client.get(
_settings.google_oauth_jwks_url, timeout=10.0
)
response.raise_for_status()
data = response.json()
self._keys = {
key["kid"]: key
for key in data.get("keys", [])
if isinstance(key, dict) and "kid" in key
}
self._last_fetched = time.time()
_logger.info("Google JWKS refreshed (%d keys)", len(self._keys))
except httpx.HTTPError as exc:
_logger.error("Failed to refresh Google JWKS: %s", exc)
raise GoogleOAuthError(
"Unable to fetch Google signing keys", status_code=502
) from exc
class GoogleOAuthService:
"""Abstraction over the Google OAuth 2.0 / OpenID Connect flow."""
def __init__(self) -> None:
self._http = SharedAsyncClient(
timeout=httpx.Timeout(_settings.google_oauth_timeout),
follow_redirects=False,
)
self._jwks = GoogleJWKSCache()
# ------------------------------------------------------------------
# HTTP client management
# ------------------------------------------------------------------
async def _get_client(self) -> httpx.AsyncClient:
return await self._http.get()
async def close(self) -> None:
await self._http.close()
# ------------------------------------------------------------------
# Authorization URL (Step 1)
# ------------------------------------------------------------------
def build_auth_url(
self,
*,
client_id: str,
redirect_uri: str,
state: str,
scope: str,
prompt: Optional[str] = None,
access_type: Optional[str] = None,
login_hint: Optional[str] = None,
include_granted_scopes: bool = False,
) -> str:
params: Dict[str, Any] = {
"client_id": client_id,
"redirect_uri": redirect_uri,
"response_type": "code",
"scope": scope,
"state": state,
}
if prompt:
params["prompt"] = prompt
if access_type:
params["access_type"] = access_type
if login_hint:
params["login_hint"] = login_hint
if include_granted_scopes:
params["include_granted_scopes"] = "true"
return f"{_settings.google_oauth_auth_url}?{urlencode(params)}"
# ------------------------------------------------------------------
# Token endpoint helpers
# ------------------------------------------------------------------
async def _post_token(
self, data: Dict[str, str], *, context: str
) -> Dict[str, Any]:
client = await self._get_client()
last_error: Optional[str] = None
for attempt in range(1 + _settings.google_oauth_max_retries):
try:
response = await client.post(_settings.google_oauth_token_url, data=data)
if response.status_code == 200:
return response.json()
if response.status_code in _RETRYABLE_STATUS:
last_error = (
f"Google OAuth upstream error HTTP {response.status_code}"
)
_logger.warning(
"%s: transient HTTP %s (attempt %d/%d)",
context,
response.status_code,
attempt + 1,
1 + _settings.google_oauth_max_retries,
)
else:
return self._parse_token_error(response, context=context)
except httpx.TimeoutException:
last_error = "Google OAuth request timed out"
_logger.warning(
"%s: timeout (attempt %d/%d)",
context,
attempt + 1,
1 + _settings.google_oauth_max_retries,
)
except httpx.RequestError as exc:
last_error = f"Google OAuth request failed: {exc}"
_logger.warning(
"%s: %s (attempt %d/%d)",
context,
last_error,
attempt + 1,
1 + _settings.google_oauth_max_retries,
)
if attempt < _settings.google_oauth_max_retries:
await asyncio.sleep(2 ** attempt)
raise GoogleOAuthError(last_error or "Unknown Google OAuth error", status_code=502)
@staticmethod
def _parse_token_error(
response: httpx.Response, *, context: str
) -> Dict[str, Any]:
try:
body = response.json()
error: str = body.get("error", "") or ""
error_description: str = body.get("error_description", "") or ""
except Exception:
error = ""
error_description = ""
body = {}
_logger.warning("%s: Google token error '%s': %s", context, error, error_description)
if error == "invalid_grant":
message = error_description or (
"The provided code or refresh token is invalid, expired, or has been revoked."
)
raise GoogleOAuthError(message, status_code=401)
if error == "invalid_client":
raise GoogleOAuthError(
error_description or "Invalid client_id or client_secret.", status_code=401
)
if error == "invalid_request":
raise GoogleOAuthError(
error_description or "Malformed OAuth request.", status_code=400
)
if error == "access_denied":
raise GoogleOAuthError(
error_description or "Access denied by the user.", status_code=403
)
# Google does not return a structured error body for some failures.
if response.status_code >= 500:
raise GoogleOAuthError(
"Google OAuth service is temporarily unavailable.", status_code=502
)
if response.status_code == 429:
raise GoogleOAuthError(
"Google OAuth rate limit exceeded. Please retry later.", status_code=429
)
message = error_description or f"Google OAuth error (HTTP {response.status_code})."
return {"error": True, "message": message, "raw": body}
# ------------------------------------------------------------------
# Code exchange (Step 2)
# ------------------------------------------------------------------
async def exchange_code(
self,
*,
client_id: str,
client_secret: str,
code: str,
redirect_uri: str,
) -> Dict[str, Any]:
data = {
"client_id": client_id,
"client_secret": client_secret,
"code": code,
"grant_type": "authorization_code",
"redirect_uri": redirect_uri,
}
tokens = await self._post_token(data, context="authorization code exchange")
self._raise_for_token_error(tokens)
return tokens
async def refresh_access_token(
self,
*,
client_id: str,
client_secret: str,
refresh_token: str,
) -> Dict[str, Any]:
data = {
"client_id": client_id,
"client_secret": client_secret,
"refresh_token": refresh_token,
"grant_type": "refresh_token",
}
tokens = await self._post_token(data, context="refresh token exchange")
self._raise_for_token_error(tokens)
return tokens
@staticmethod
def _raise_for_token_error(tokens: Dict[str, Any]) -> None:
if tokens.get("error"):
raise GoogleOAuthError(tokens.get("message", "Unknown OAuth error"), status_code=400)
if not tokens.get("access_token"):
raise GoogleOAuthError(
"Google did not return an access token.", status_code=502
)
# ------------------------------------------------------------------
# ID token verification
# ------------------------------------------------------------------
async def verify_id_token(
self, id_token: str, client_id: str
) -> GoogleOAuthUserInfo:
client = await self._get_client()
try:
unverified = jwt.get_unverified_header(id_token)
kid = unverified.get("kid")
alg = unverified.get("alg")
except jwt.DecodeError as exc:
raise GoogleOAuthError("Malformed ID token.", status_code=400) from exc
if not kid or alg != "RS256":
raise GoogleOAuthError(
"ID token does not use an RS256 signature.", status_code=401
)
try:
public_key = await self._jwks.get_public_key(client, kid)
except GoogleOAuthError:
raise
try:
payload = await run_in_executor(
jwt.decode,
id_token,
key=public_key,
algorithms=["RS256"],
audience=client_id,
issuer=["accounts.google.com", "https://accounts.google.com"],
options={
"verify_exp": True,
"verify_iat": True,
"verify_aud": True,
"verify_iss": True,
"require": ["exp", "iat", "sub", "email"],
},
)
except jwt.ExpiredSignatureError as exc:
raise GoogleOAuthError("ID token has expired.", status_code=401) from exc
except jwt.InvalidAudienceError as exc:
raise GoogleOAuthError(
"ID token audience does not match the provided client_id.", status_code=401
) from exc
except jwt.InvalidIssuerError as exc:
raise GoogleOAuthError(
"ID token issuer is not Google.", status_code=401
) from exc
except jwt.PyJWTError as exc:
_logger.warning("ID token verification failed: %s", exc)
raise GoogleOAuthError("Invalid ID token.", status_code=401) from exc
return self._user_info_from_claims(payload)
# ------------------------------------------------------------------
# Userinfo (fallback profile source)
# ------------------------------------------------------------------
@staticmethod
def _user_info_from_claims(payload: Dict[str, Any]) -> GoogleOAuthUserInfo:
return GoogleOAuthUserInfo(
sub=payload.get("sub", ""),
email=payload.get("email", ""),
email_verified=bool(payload.get("email_verified", False)),
name=payload.get("name"),
given_name=payload.get("given_name"),
family_name=payload.get("family_name"),
picture=payload.get("picture"),
locale=payload.get("locale"),
)
|