AegisGitea-MCP/src/aegis_gitea_mcp/oauth/resolver.py

"""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__)