"""Identity resolution for per-user OAuth, with automatic upstream refresh. Static and OAuth identities never mix: * an access token carrying our OAuth marker (``claims["auth"] == "oauth"``) always resolves to that session's Gitea user, or fails; it never falls back to the admin token; * only the static bearer (``client_id == STATIC_CLIENT_ID`` and no OAuth marker) or no token at all reaches the admin resolver, and in an OAuth-only deployment (no ``GITEA_TOKEN``) that fails closed; * any other access token is refused. """ from __future__ import annotations import asyncio import dataclasses import logging import time from collections.abc import Callable from typing import TYPE_CHECKING from mcp.server.auth.provider import AccessToken from aegis_gitea_mcp.audit import AuditLogger, get_audit_logger from aegis_gitea_mcp.auth import STATIC_CLIENT_ID, IdentityResolver from aegis_gitea_mcp.errors import ToolError from aegis_gitea_mcp.identity import Identity from aegis_gitea_mcp.oauth.store import OAuthStore, OAuthStoreError, SessionRecord from aegis_gitea_mcp.oauth.upstream import GiteaUpstream, UpstreamError, UpstreamInvalidGrant if TYPE_CHECKING: pass logger = logging.getLogger(__name__) #: Refresh the user's Gitea token when it expires within this many seconds. REFRESH_MARGIN_SECONDS = 60.0 SIGN_IN_AGAIN = "your Gitea sign-in has ended; sign in again through your MCP client" OAUTH_ONLY = ( "this server accepts only per-user OAuth sign-in (no GITEA_TOKEN is configured for " "admin-token mode)" ) class OAuthIdentityResolver: """Maps a verified access token to the identity a tool call runs as. Args: store: OAuth store (sessions and Gitea tokens). upstream: Gitea OAuth client used to refresh the user's Gitea token. admin: Resolver for the static bearer (admin-token mode), or None in an OAuth-only deployment. clock: Wall clock (epoch seconds). audit: Audit logger for ended sessions (defaults to the process-wide one). """ def __init__( self, store: OAuthStore, upstream: GiteaUpstream, admin: IdentityResolver | None, *, clock: Callable[[], float] = time.time, audit: AuditLogger | None = None, ) -> None: self._store = store self._upstream = upstream self._admin = admin self._clock = clock self._audit_logger = audit self._locks: dict[str, asyncio.Lock] = {} async def resolve(self, access: AccessToken | None) -> Identity: """Return the user's identity for an OAuth token, the admin's for the static bearer. Raises: ToolError: Unknown token kind, ended session, failed upstream refresh, or admin-token mode requested while it is not configured. """ claims = (access.claims or {}) if access is not None else {} if access is not None and claims.get("auth") == "oauth": return await self._resolve_session(access, claims.get("session_id")) if access is None or access.client_id == STATIC_CLIENT_ID: if self._admin is None: raise ToolError(OAUTH_ONLY, status_code=401) return await self._admin.resolve(access) raise ToolError("unknown identity", status_code=401) async def _resolve_session(self, access: AccessToken, session_id: object) -> Identity: if not isinstance(session_id, str): raise ToolError("unknown identity", status_code=401) session = self._store.sessions.get(session_id) if session is None or session.login != access.subject: raise ToolError(SIGN_IN_AGAIN, status_code=401) if self._needs_refresh(session): session = await self._refresh(session_id) redact = (session.gitea_refresh_token,) if session.gitea_refresh_token else () # The executor redacts ``Identity.redact`` (the Gitea refresh token) from results. return Identity( login=session.login, mode="oauth", token=session.gitea_access_token, redact=redact ) def _needs_refresh(self, session: SessionRecord) -> bool: expires = session.gitea_expires_at return expires is not None and expires - self._clock() <= REFRESH_MARGIN_SECONDS async def _refresh(self, session_id: str) -> SessionRecord: lock = self._locks.setdefault(session_id, asyncio.Lock()) async with lock: session = self._store.sessions.get(session_id) if session is None: raise ToolError(SIGN_IN_AGAIN, status_code=401) if not self._needs_refresh(session): return session # refreshed by a concurrent call still_valid = (session.gitea_expires_at or 0.0) > self._clock() if session.gitea_refresh_token is None: if still_valid: return session await self._end(session, "Gitea token expired and no refresh token was issued") raise ToolError(SIGN_IN_AGAIN, status_code=401) try: tokens = await self._upstream.refresh(session.gitea_refresh_token) except UpstreamInvalidGrant: await self._end(session, "Gitea refused the refresh token") raise ToolError(SIGN_IN_AGAIN, status_code=401) from None except UpstreamError as exc: logger.warning("gitea token refresh failed for %s: %s", session.login, exc) if still_valid: return session raise ToolError( "could not refresh your Gitea sign-in; try again shortly", status_code=503 ) from None updated = dataclasses.replace( session, gitea_access_token=tokens.access_token, gitea_refresh_token=tokens.refresh_token or session.gitea_refresh_token, gitea_expires_at=tokens.expires_at, ) self._store.sessions[session_id] = updated try: await self._store.save() except (OSError, OAuthStoreError) as exc: logger.error("could not save the refreshed Gitea token (%s)", type(exc).__name__) logger.info("gitea token refreshed for %s", session.login) return updated async def _end(self, session: SessionRecord, reason: str) -> None: self._store.end_session(session.session_id) self._locks.pop(session.session_id, None) logger.warning("oauth session ended for %s: %s", session.login, reason) try: audit = self._audit_logger or get_audit_logger() audit.log_security_event( "oauth_session_ended", reason, severity="low", metadata={"login": session.login, "client_id": session.client_id}, ) except Exception as exc: logger.error("audit log write failed: %s", type(exc).__name__) try: await self._store.save() except (OSError, OAuthStoreError) as exc: logger.error("could not save the OAuth store (%s)", type(exc).__name__)