From 16f1b0e9a60510fecf2cf4ae7ff171c9e29e658f Mon Sep 17 00:00:00 2001 From: julianam-w <77252335+julianam-w@users.noreply.github.com> Date: Sun, 16 Aug 2026 09:26:37 +1000 Subject: [PATCH] refactor: source OAuth proxy from shared bes-mcp-oauth package Extract the Google OAuth proxy (oauth_proxy.py + auth.py) into the shared bes-mcp-oauth package so BES MCP servers (bes-codebases, tupaia-metadata-mcp, ...) share one implementation. Behaviour is unchanged: mcp_server imports GoogleOAuthProxyProvider from the package, and the two modules are removed here. Verified with uv sync + an import check. Co-Authored-By: Claude Opus 4.8 --- mcp_server.py | 2 +- pyproject.toml | 2 + rag/auth.py | 92 ----------- rag/oauth_proxy.py | 383 --------------------------------------------- uv.lock | 11 ++ 5 files changed, 14 insertions(+), 476 deletions(-) delete mode 100644 rag/auth.py delete mode 100644 rag/oauth_proxy.py diff --git a/mcp_server.py b/mcp_server.py index 84f926c..50916c7 100644 --- a/mcp_server.py +++ b/mcp_server.py @@ -39,7 +39,7 @@ def _build_mcp(host: str, port: int, with_auth: bool): RevocationOptions, ) - from rag.oauth_proxy import GoogleOAuthProxyProvider + from bes_mcp_oauth import GoogleOAuthProxyProvider allowed_domain = os.environ.get("GOOGLE_ALLOWED_DOMAIN", "") allowed_emails = [ diff --git a/pyproject.toml b/pyproject.toml index e2c903b..ef2080a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,6 +10,8 @@ dependencies = [ "langchain-text-splitters>=0.3.0", "python-dotenv>=1.0.0", "mcp>=1.0.0,<2", # 2.x removed mcp.server.fastmcp (breaking); stay on 1.x + # Shared Google OAuth proxy (one source of truth across BES MCP servers). + "bes-mcp-oauth @ git+https://github.com/beyondessential/bes-mcp-oauth@v0.1.0", ] [project.optional-dependencies] diff --git a/rag/auth.py b/rag/auth.py deleted file mode 100644 index 26b5d33..0000000 --- a/rag/auth.py +++ /dev/null @@ -1,92 +0,0 @@ -""" -Google OAuth token verification for the MCP server (HTTP transport). - -Implements MCP's TokenVerifier protocol — validates Google OAuth2 access tokens -via Google's tokeninfo endpoint and enforces domain/email restrictions. -""" - -import httpx -from mcp.server.auth.provider import AccessToken - -GOOGLE_TOKENINFO_URL = "https://www.googleapis.com/oauth2/v3/tokeninfo" - - -class GoogleTokenVerifier: - """Validates Google OAuth2 access tokens for MCP HTTP transport. - - Tokens are verified against Google's tokeninfo API. Callers must have - obtained a Google access token with at least the ``openid email`` scopes - (e.g. via ``gcloud auth print-access-token`` or any standard OAuth flow). - - Access is restricted by domain and/or email allowlist: - - ``allowed_domain``: only emails from this domain are accepted - (e.g. ``"bes.au"``) - - ``allowed_emails``: only these specific emails are accepted - - If both are empty, any verified Google account is accepted. - - ``allowed_client_ids`` binds tokens to our own OAuth client: the token's - ``aud``/``azp`` must match one of these. Without it, any valid Google token - for a permitted user — issued to *any* app the user has used — would be - accepted, allowing token replay. Leave empty only for local testing. - """ - - def __init__( - self, - allowed_domain: str = "", - allowed_emails: list[str] | None = None, - allowed_client_ids: list[str] | None = None, - ) -> None: - self.allowed_domain = allowed_domain - self.allowed_emails = set(allowed_emails or []) - self.allowed_client_ids = set(allowed_client_ids or []) - - async def verify_token(self, token: str) -> AccessToken | None: - """Validate a Google OAuth2 access token. - - Returns an AccessToken if valid and permitted, None otherwise. - """ - try: - async with httpx.AsyncClient() as client: - resp = await client.get( - GOOGLE_TOKENINFO_URL, - params={"access_token": token}, - ) - - if resp.status_code != 200: - return None - - info = resp.json() - if "error" in info: - return None - - email: str = info.get("email", "") - if not email: - return None # token missing email scope - if not info.get("email_verified"): - return None # unverified email - - # Audience binding: the token must have been issued to our OAuth - # client, not merely to some other app the user has authorised. - if self.allowed_client_ids: - aud = info.get("aud", "") - azp = info.get("azp", "") - if aud not in self.allowed_client_ids and azp not in self.allowed_client_ids: - return None - - # Domain restriction - if self.allowed_domain and not email.endswith(f"@{self.allowed_domain}"): - return None - - # Email allowlist - if self.allowed_emails and email not in self.allowed_emails: - return None - - return AccessToken( - token=token, - client_id=email, - scopes=info.get("scope", "").split() if info.get("scope") else [], - expires_at=int(info["exp"]) if info.get("exp") else None, - ) - - except Exception: - return None diff --git a/rag/oauth_proxy.py b/rag/oauth_proxy.py deleted file mode 100644 index 0ac257c..0000000 --- a/rag/oauth_proxy.py +++ /dev/null @@ -1,383 +0,0 @@ -""" -Google OAuth *proxy* authorization server for the MCP server (HTTP transport). - -Google does not support OAuth 2.0 Dynamic Client Registration (DCR, RFC 7591), -so MCP clients (Claude Code, the claude.ai connector) cannot register with it -automatically — they dead-end on "does not support dynamic client registration" -and require a hand-configured client id + secret. - -This module makes *this* server the authorization server the clients talk to, -and has it broker a second OAuth exchange with Google behind the scenes: - - +--------+ +------------+ +-------------------+ - | MCP | --> | this MCP | --> | Google (accounts. | - | client | | server | | google.com) | - +--------+ +------------+ +-------------------+ - -The MCP client speaks standard DCR + authorization-code + PKCE to us; we speak -a single pre-registered OAuth client to Google. The client never needs a Google -client id/secret — only our URL. - -Token model: we mint our own opaque access/refresh tokens and keep the Google -tokens server-side, so Google credentials are never handed to clients. Access is -still gated to the configured Google Workspace domain / email allowlist, reusing -``rag.auth.GoogleTokenVerifier`` to validate the Google token at login time. - -State is in-memory: a single server instance. A restart invalidates in-flight -logins and issued tokens, and clients re-authenticate — acceptable for a single -Railway instance. Move these stores to Postgres if the server is ever scaled out. -""" - -import secrets -import time - -import httpx -from mcp.server.auth.provider import ( - AccessToken, - AuthorizationCode, - AuthorizationParams, - RefreshToken, - TokenError, - construct_redirect_uri, -) -from mcp.shared.auth import OAuthClientInformationFull, OAuthToken -from pydantic import AnyUrl -from starlette.requests import Request -from starlette.responses import PlainTextResponse, RedirectResponse, Response - -from rag.auth import GoogleTokenVerifier - -GOOGLE_AUTH_URL = "https://accounts.google.com/o/oauth2/v2/auth" -GOOGLE_TOKEN_URL = "https://oauth2.googleapis.com/token" - -# Scopes we request from Google. "openid email" is the minimum the token -# verifier needs (verified email + domain gating). -GOOGLE_SCOPES = "openid email" - -AUTH_CODE_TTL_SECONDS = 300 # our authorization codes are short-lived - - -class GoogleOAuthProxyProvider: - """OAuthAuthorizationServerProvider that proxies to Google. - - Implements the MCP SDK's ``OAuthAuthorizationServerProvider`` protocol so it - can be passed to ``FastMCP(auth_server_provider=...)``. The paired Google - redirect handler is :meth:`handle_google_callback`, mounted as a custom route. - """ - - def __init__( - self, - *, - public_url: str, - google_client_id: str, - google_client_secret: str, - callback_path: str = "/auth/google/callback", - allowed_domain: str = "", - allowed_emails: list[str] | None = None, - ) -> None: - if not google_client_id or not google_client_secret: - raise ValueError( - "GoogleOAuthProxyProvider requires GOOGLE_OAUTH_CLIENT_ID and " - "GOOGLE_OAUTH_CLIENT_SECRET (the server-side Google web client " - "used to broker sign-in)." - ) - self.public_url = public_url.rstrip("/") - self.google_client_id = google_client_id - self.google_client_secret = google_client_secret - self.callback_path = callback_path - self.redirect_uri = f"{self.public_url}{callback_path}" - # Reuse the existing verifier for domain/email/audience gating. The - # Google token's audience is our own client, so bind to it. - self._verifier = GoogleTokenVerifier( - allowed_domain=allowed_domain, - allowed_emails=allowed_emails or None, - allowed_client_ids=[google_client_id], - ) - - # In-memory state (single instance; see module docstring). - self._clients: dict[str, OAuthClientInformationFull] = {} - self._transactions: dict[str, dict] = {} # google-state -> pending login - self._auth_codes: dict[str, AuthorizationCode] = {} # our code -> code - self._code_google: dict[str, dict] = {} # our code -> google token bundle - self._access_tokens: dict[str, AccessToken] = {} # our token -> AccessToken - self._refresh_tokens: dict[str, RefreshToken] = {} # our token -> RefreshToken - self._refresh_google: dict[str, str] = {} # our refresh -> google refresh - - # ── Dynamic client registration ───────────────────────────────────────── - - async def get_client(self, client_id: str) -> OAuthClientInformationFull | None: - return self._clients.get(client_id) - - async def register_client(self, client_info: OAuthClientInformationFull) -> None: - # The SDK's /register handler mints client_id/secret; we just persist. - self._clients[client_info.client_id] = client_info - - # ── Authorization: redirect the user to Google ────────────────────────── - - async def authorize( - self, client: OAuthClientInformationFull, params: AuthorizationParams - ) -> str: - # Open a login transaction keyed by the state we hand to Google, holding - # everything needed to complete the client's leg after Google returns. - google_state = secrets.token_urlsafe(32) - self._transactions[google_state] = { - "client_id": client.client_id, - "redirect_uri": str(params.redirect_uri), - "redirect_uri_provided_explicitly": params.redirect_uri_provided_explicitly, - "client_state": params.state, - "code_challenge": params.code_challenge, - "scopes": params.scopes or ["openid", "email"], - "resource": params.resource, - "created_at": time.time(), - } - - query = { - "client_id": self.google_client_id, - "redirect_uri": self.redirect_uri, - "response_type": "code", - "scope": GOOGLE_SCOPES, - "state": google_state, - "access_type": "offline", - "prompt": "consent", - } - return construct_redirect_uri(GOOGLE_AUTH_URL, **query) - - async def handle_google_callback(self, request: Request) -> Response: - """Google redirects here after the user signs in. - - Mounted as a custom GET route at ``self.callback_path``. Exchanges the - Google code for tokens, gates access, mints our own authorization code, - and redirects back to the MCP client's ``redirect_uri``. - """ - params = request.query_params - google_state = params.get("state", "") - txn = self._transactions.pop(google_state, None) - if txn is None: - return PlainTextResponse("Invalid or expired login state", status_code=400) - - if params.get("error"): - return self._fail_back( - txn, - params.get("error", "access_denied"), - params.get("error_description"), - ) - - code = params.get("code") - if not code: - return self._fail_back(txn, "invalid_request", "missing authorization code") - - # Exchange the Google code for tokens (server-side, confidential client). - try: - async with httpx.AsyncClient() as http: - resp = await http.post( - GOOGLE_TOKEN_URL, - data={ - "code": code, - "client_id": self.google_client_id, - "client_secret": self.google_client_secret, - "redirect_uri": self.redirect_uri, - "grant_type": "authorization_code", - }, - ) - except Exception: - return self._fail_back(txn, "server_error", "google token exchange failed") - - if resp.status_code != 200: - return self._fail_back( - txn, "access_denied", "google rejected the authorization" - ) - - google = resp.json() - google_access = google.get("access_token", "") - - # Gate access: verify the Google token (verified email, domain/email - # allowlist, audience bound to our client). Reuses rag.auth. - access = await self._verifier.verify_token(google_access) - if access is None: - return self._fail_back( - txn, "access_denied", "account not permitted for this server" - ) - - # Mint our own authorization code and stash the Google token bundle. - our_code = secrets.token_urlsafe(32) - self._auth_codes[our_code] = AuthorizationCode( - code=our_code, - scopes=txn["scopes"], - expires_at=time.time() + AUTH_CODE_TTL_SECONDS, - client_id=txn["client_id"], - code_challenge=txn["code_challenge"], - redirect_uri=AnyUrl(txn["redirect_uri"]), - redirect_uri_provided_explicitly=txn["redirect_uri_provided_explicitly"], - resource=txn["resource"], - ) - self._code_google[our_code] = { - "access_token": google_access, - "refresh_token": google.get("refresh_token"), - "expires_at": time.time() + int(google.get("expires_in", 3600)), - "email": access.client_id, # GoogleTokenVerifier sets client_id=email - } - - return RedirectResponse( - url=construct_redirect_uri( - txn["redirect_uri"], code=our_code, state=txn["client_state"] - ), - status_code=302, - ) - - def _fail_back(self, txn: dict, error: str, description: str | None) -> Response: - """Redirect the OAuth error back to the MCP client's redirect_uri.""" - return RedirectResponse( - url=construct_redirect_uri( - txn["redirect_uri"], - error=error, - error_description=description, - state=txn["client_state"], - ), - status_code=302, - ) - - # ── Token endpoint ────────────────────────────────────────────────────── - - async def load_authorization_code( - self, client: OAuthClientInformationFull, authorization_code: str - ) -> AuthorizationCode | None: - code = self._auth_codes.get(authorization_code) - if code is None or code.client_id != client.client_id: - return None - if code.expires_at < time.time(): - self._auth_codes.pop(authorization_code, None) - self._code_google.pop(authorization_code, None) - return None - return code - - async def exchange_authorization_code( - self, client: OAuthClientInformationFull, authorization_code: AuthorizationCode - ) -> OAuthToken: - # Single-use: consume the code and its Google bundle. - self._auth_codes.pop(authorization_code.code, None) - bundle = self._code_google.pop(authorization_code.code, None) - if bundle is None: - raise TokenError( - "invalid_grant", "authorization code already used or expired" - ) - - return self._issue_tokens( - client_id=client.client_id, - scopes=authorization_code.scopes, - resource=authorization_code.resource, - google_access=bundle["access_token"], - google_refresh=bundle.get("refresh_token"), - expires_at=bundle["expires_at"], - ) - - def _issue_tokens( - self, - *, - client_id: str, - scopes: list[str], - resource: str | None, - google_access: str, - google_refresh: str | None, - expires_at: float, - ) -> OAuthToken: - our_access = secrets.token_urlsafe(32) - self._access_tokens[our_access] = AccessToken( - token=our_access, - client_id=client_id, - scopes=scopes, - expires_at=int(expires_at), - resource=resource, - ) - - our_refresh: str | None = None - if google_refresh: - our_refresh = secrets.token_urlsafe(32) - self._refresh_tokens[our_refresh] = RefreshToken( - token=our_refresh, - client_id=client_id, - scopes=scopes, - ) - self._refresh_google[our_refresh] = google_refresh - - return OAuthToken( - access_token=our_access, - token_type="Bearer", - expires_in=max(1, int(expires_at - time.time())), - scope=" ".join(scopes), - refresh_token=our_refresh, - ) - - # ── Refresh ───────────────────────────────────────────────────────────── - - async def load_refresh_token( - self, client: OAuthClientInformationFull, refresh_token: str - ) -> RefreshToken | None: - token = self._refresh_tokens.get(refresh_token) - if token is None or token.client_id != client.client_id: - return None - return token - - async def exchange_refresh_token( - self, - client: OAuthClientInformationFull, - refresh_token: RefreshToken, - scopes: list[str], - ) -> OAuthToken: - google_refresh = self._refresh_google.get(refresh_token.token) - if google_refresh is None: - raise TokenError("invalid_grant", "unknown refresh token") - - try: - async with httpx.AsyncClient() as http: - resp = await http.post( - GOOGLE_TOKEN_URL, - data={ - "refresh_token": google_refresh, - "client_id": self.google_client_id, - "client_secret": self.google_client_secret, - "grant_type": "refresh_token", - }, - ) - except Exception as exc: - raise TokenError("server_error", "google refresh failed") from exc - - if resp.status_code != 200: - raise TokenError("invalid_grant", "google refused the refresh token") - - google = resp.json() - new_google_access = google.get("access_token", "") - # Re-gate on refresh so revoked/expired accounts lose access promptly. - if await self._verifier.verify_token(new_google_access) is None: - raise TokenError("access_denied", "account no longer permitted") - - # Rotate: retire the old refresh token, issue a fresh pair. Google may or - # may not return a new refresh token; keep the existing one if not. - self._refresh_tokens.pop(refresh_token.token, None) - self._refresh_google.pop(refresh_token.token, None) - - requested_scopes = scopes or refresh_token.scopes - return self._issue_tokens( - client_id=client.client_id, - scopes=requested_scopes, - resource=None, - google_access=new_google_access, - google_refresh=google.get("refresh_token") or google_refresh, - expires_at=time.time() + int(google.get("expires_in", 3600)), - ) - - # ── Verification / revocation ─────────────────────────────────────────── - - async def load_access_token(self, token: str) -> AccessToken | None: - access = self._access_tokens.get(token) - if access is None: - return None - if access.expires_at is not None and access.expires_at < time.time(): - self._access_tokens.pop(token, None) - return None - return access - - async def revoke_token(self, token: AccessToken | RefreshToken) -> None: - # Drop whichever kind we were handed; best-effort, idempotent. - self._access_tokens.pop(token.token, None) - if self._refresh_tokens.pop(token.token, None) is not None: - self._refresh_google.pop(token.token, None) diff --git a/uv.lock b/uv.lock index 7cf848c..a189697 100644 --- a/uv.lock +++ b/uv.lock @@ -207,6 +207,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/64/b4/17d4b0b2a2dc85a6df63d1157e028ed19f90d4cd97c36717afef2bc2f395/attrs-26.1.0-py3-none-any.whl", hash = "sha256:c647aa4a12dfbad9333ca4e71fe62ddc36f4e63b2d260a37a8b83d2f043ac309", size = 67548, upload-time = "2026-03-19T14:22:23.645Z" }, ] +[[package]] +name = "bes-mcp-oauth" +version = "0.1.0" +source = { git = "https://github.com/beyondessential/bes-mcp-oauth?rev=v0.1.0#df9cd055911d0c9320eb691fdae8aff5c1771caa" } +dependencies = [ + { name = "httpx" }, + { name = "mcp" }, +] + [[package]] name = "certifi" version = "2026.7.22" @@ -611,6 +620,7 @@ name = "github-repo-rag" version = "0.1.0" source = { editable = "." } dependencies = [ + { name = "bes-mcp-oauth" }, { name = "langchain-text-splitters" }, { name = "mcp" }, { name = "pgvector" }, @@ -627,6 +637,7 @@ ask = [ [package.metadata] requires-dist = [ { name = "anthropic", marker = "extra == 'ask'", specifier = ">=0.40.0" }, + { name = "bes-mcp-oauth", git = "https://github.com/beyondessential/bes-mcp-oauth?rev=v0.1.0" }, { name = "langchain-text-splitters", specifier = ">=0.3.0" }, { name = "mcp", specifier = ">=1.0.0,<2" }, { name = "pgvector", specifier = ">=0.3.0" },