fix: address code review findings for Telegram OIDC

- JWKS cache: add asyncio.Lock to prevent thundering herd, extract _build_public_keys helper, retry JWKS fetch on kid mismatch (key rotation) - Remove dead code: exchange_telegram_oidc_code (unused, popup sends id_token directly) - OIDC auth endpoint: add rate limiting, fix int() parse with try/except, extract last_name/photo_url/language from claims, narrow bare Exception to (ValueError, LookupError) - Schema: add max_length=4096 to id_token field - Branding: read TELEGRAM_OIDC_ENABLED from DB settings with env fallback
2026-03-07 02:54:29 +03:00
parent 000b0c0592
commit da1cc4fe5a
4 changed files with 65 additions and 57 deletions
@@ -1,5 +1,6 @@
 """Telegram authentication validation for cabinet."""

+import asyncio
 import hashlib
 import hmac
 import json
@@ -146,6 +147,18 @@ _JWKS_URL = 'https://oauth.telegram.org/.well-known/jwks.json'
 _OIDC_ISSUER = 'https://oauth.telegram.org'
 _OIDC_TOKEN_URL = 'https://oauth.telegram.org/token'

+_jwks_lock = asyncio.Lock()
+
+
+def _build_public_keys(jwks_data: dict[str, Any]) -> dict[str, Any]:
+    """Build public key mapping from JWKS data."""
+    public_keys: dict[str, Any] = {}
+    for key_data in jwks_data.get('keys', []):
+        kid = key_data.get('kid')
+        if kid:
+            public_keys[kid] = pyjwt.algorithms.RSAAlgorithm.from_jwk(key_data)
+    return public_keys
+

 async def _get_jwks() -> dict[str, Any]:
    """Fetch and cache Telegram OIDC JWKS keys."""
@@ -155,12 +168,18 @@ async def _get_jwks() -> dict[str, Any]:
    if _jwks_cache and _jwks_cache_expiry and now < _jwks_cache_expiry:
        return _jwks_cache

-    async with httpx.AsyncClient(timeout=10) as client:
-        response = await client.get(_JWKS_URL)
-        response.raise_for_status()
-        _jwks_cache = response.json()
-        _jwks_cache_expiry = now + timedelta(seconds=_JWKS_CACHE_TTL_SECONDS)
-        return _jwks_cache
+    async with _jwks_lock:
+        # Double-check after acquiring lock
+        now = datetime.now(UTC)
+        if _jwks_cache and _jwks_cache_expiry and now < _jwks_cache_expiry:
+            return _jwks_cache
+
+        async with httpx.AsyncClient(timeout=10) as client:
+            response = await client.get(_JWKS_URL)
+            response.raise_for_status()
+            _jwks_cache = response.json()
+            _jwks_cache_expiry = now + timedelta(seconds=_JWKS_CACHE_TTL_SECONDS)
+            return _jwks_cache


 async def validate_telegram_oidc_token(id_token: str, client_id: str) -> dict[str, Any] | None:
@@ -176,16 +195,21 @@ async def validate_telegram_oidc_token(id_token: str, client_id: str) -> dict[st
        Claims include: sub, id, name, preferred_username, picture, iss, aud, exp, iat
    """
    try:
+        # Build public keys from JWKS
        jwks_data = await _get_jwks()
-        public_keys = {}
-        for key_data in jwks_data.get('keys', []):
-            kid = key_data.get('kid')
-            if kid:
-                public_keys[kid] = pyjwt.algorithms.RSAAlgorithm.from_jwk(key_data)
+        public_keys = _build_public_keys(jwks_data)

        # Decode header to get kid
        unverified_header = pyjwt.get_unverified_header(id_token)
        kid = unverified_header.get('kid')
+
+        # If kid not found, force JWKS refresh (key rotation)
+        if kid and kid not in public_keys:
+            global _jwks_cache_expiry
+            _jwks_cache_expiry = None
+            jwks_data = await _get_jwks()
+            public_keys = _build_public_keys(jwks_data)
+
        if not kid or kid not in public_keys:
            logger.warning('Telegram OIDC: unknown kid in id_token', kid=kid)
            return None
@@ -211,43 +235,3 @@ async def validate_telegram_oidc_token(id_token: str, client_id: str) -> dict[st
        return None


-async def exchange_telegram_oidc_code(
-    code: str,
-    client_id: str,
-    client_secret: str,
-    redirect_uri: str,
-    code_verifier: str | None = None,
-) -> str | None:
-    """
-    Exchange authorization code for id_token at Telegram OIDC token endpoint.
-
-    Args:
-        code: Authorization code from Telegram
-        client_id: Bot numeric ID
-        client_secret: OIDC secret from BotFather
-        redirect_uri: Must match the one used in authorization request
-        code_verifier: PKCE code_verifier if S256 was used
-
-    Returns:
-        id_token string if successful, None otherwise
-    """
-    try:
-        data: dict[str, str] = {
-            'grant_type': 'authorization_code',
-            'code': code,
-            'redirect_uri': redirect_uri,
-            'client_id': client_id,
-            'client_secret': client_secret,
-        }
-        if code_verifier:
-            data['code_verifier'] = code_verifier
-
-        async with httpx.AsyncClient(timeout=30) as client:
-            response = await client.post(_OIDC_TOKEN_URL, data=data)
-            response.raise_for_status()
-            token_data = response.json()
-            return token_data.get('id_token')
-
-    except httpx.HTTPError as e:
-        logger.error('Telegram OIDC: token exchange failed', error=str(e))
-        return None
@@ -549,6 +549,7 @@ async def auth_telegram_widget(
@router.post('/telegram/oidc', response_model=AuthResponse)
 async def auth_telegram_oidc(
    request: TelegramOIDCAuthRequest,
+    raw_request: Request,
    db: AsyncSession = Depends(get_cabinet_db),
 ):
    """
@@ -557,6 +558,11 @@ async def auth_telegram_oidc(
    The frontend uses Telegram.Login.init() popup which returns an id_token.
    We validate it via JWKS and create/login the user.
    """
+    # Rate limit
+    client_ip = get_client_ip(raw_request)
+    if await RateLimitCache.is_ip_rate_limited(client_ip, 'telegram_oidc', limit=10, window=60, fail_closed=True):
+        raise HTTPException(status_code=status.HTTP_429_TOO_MANY_REQUESTS, detail='Too many requests')
+
    if not settings.TELEGRAM_OIDC_ENABLED or not settings.TELEGRAM_OIDC_CLIENT_ID:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
@@ -574,7 +580,13 @@ async def auth_telegram_oidc(
        )

    # Extract user info from OIDC claims
-    telegram_id = int(claims.get('id', claims.get('sub', 0)))
+    try:
+        telegram_id = int(claims.get('id', claims.get('sub', 0)))
+    except (ValueError, TypeError):
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail='Invalid user ID in OIDC claims',
+        )
    if not telegram_id:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
@@ -583,6 +595,9 @@ async def auth_telegram_oidc(

    first_name = claims.get('name', claims.get('given_name', ''))
    username = claims.get('preferred_username')
+    last_name = claims.get('family_name')
+    _photo_url = claims.get('picture')  # extracted for future use
+    language = claims.get('locale', 'ru')[:2] if claims.get('locale') else 'ru'

    user = await get_user_by_telegram_id(db, telegram_id)

@@ -593,8 +608,8 @@ async def auth_telegram_oidc(
            referrer = await get_user_by_referral_code(db, request.referral_code)
            if referrer:
                referrer_id = referrer.id
-        except Exception as e:
-            logger.warning('Failed to resolve referral code', referral_code=request.referral_code, error=e)
+        except (ValueError, LookupError) as e:
+            logger.warning('Failed to resolve referral code', referral_code=request.referral_code, error=str(e))

    if not user:
        logger.info('Creating new user from cabinet OIDC', telegram_id=telegram_id, username=username)
@@ -603,7 +618,8 @@ async def auth_telegram_oidc(
            telegram_id=telegram_id,
            username=username,
            first_name=first_name,
-            language='ru',
+            last_name=last_name,
+            language=language,
            referred_by_id=referrer_id,
        )
        logger.info('User created successfully', user_id=user.id, telegram_id=user.telegram_id)
@@ -619,6 +635,8 @@ async def auth_telegram_oidc(
        user.username = username
    if first_name and first_name != user.first_name:
        user.first_name = first_name
+    if last_name is not None and last_name != user.last_name:
+        user.last_name = last_name

    user.cabinet_last_login = datetime.now(UTC)
    await db.commit()
@@ -43,6 +43,7 @@ TELEGRAM_WIDGET_SIZE_KEY = 'TELEGRAM_WIDGET_SIZE'
 TELEGRAM_WIDGET_RADIUS_KEY = 'TELEGRAM_WIDGET_RADIUS'
 TELEGRAM_WIDGET_USERPIC_KEY = 'TELEGRAM_WIDGET_USERPIC'
 TELEGRAM_WIDGET_REQUEST_ACCESS_KEY = 'TELEGRAM_WIDGET_REQUEST_ACCESS'
+TELEGRAM_OIDC_ENABLED_KEY = 'TELEGRAM_OIDC_ENABLED'

 # Default animation config
 DEFAULT_ANIMATION_CONFIG = {
@@ -867,7 +868,12 @@ async def get_telegram_widget_config(
    userpic_val = await get_setting_value(db, TELEGRAM_WIDGET_USERPIC_KEY)
    request_access_val = await get_setting_value(db, TELEGRAM_WIDGET_REQUEST_ACCESS_KEY)

-    oidc_enabled = settings.TELEGRAM_OIDC_ENABLED and bool(settings.TELEGRAM_OIDC_CLIENT_ID)
+    oidc_enabled_val = await get_setting_value(db, TELEGRAM_OIDC_ENABLED_KEY)
+    oidc_enabled = (
+        oidc_enabled_val.lower() == 'true'
+        if oidc_enabled_val is not None
+        else settings.TELEGRAM_OIDC_ENABLED
+    ) and bool(settings.TELEGRAM_OIDC_CLIENT_ID)

    return TelegramWidgetConfigResponse(
        bot_username=bot_username,
@@ -34,7 +34,7 @@ class TelegramWidgetAuthRequest(BaseModel):
 class TelegramOIDCAuthRequest(BaseModel):
    """Request for Telegram OIDC authentication (popup flow)."""

-    id_token: str = Field(..., description='JWT id_token from Telegram OIDC popup')
+    id_token: str = Field(..., max_length=4096, description='JWT id_token from Telegram OIDC popup')
    campaign_slug: str | None = Field(
        None, min_length=1, max_length=64, pattern=r'^[a-zA-Z0-9_-]+$', description='Campaign slug from web link'
    )