From d0237df9a8731f27d249f52cabcdc497ab8ec431 Mon Sep 17 00:00:00 2001 From: Spencer Murray Date: Tue, 27 Oct 2026 14:41:46 -0400 Subject: [PATCH] Validate attributed Junior bot JWT subjects --- src/firetower/mcp_server/auth.py | 25 +++++- src/firetower/mcp_server/tests/test_auth.py | 82 +++++++++++++++++-- src/firetower/mcp_server/tests/test_server.py | 5 +- 3 files changed, 100 insertions(+), 12 deletions(-) diff --git a/src/firetower/mcp_server/auth.py b/src/firetower/mcp_server/auth.py index 3445435f..f49ce54e 100644 --- a/src/firetower/mcp_server/auth.py +++ b/src/firetower/mcp_server/auth.py @@ -9,7 +9,8 @@ Bots (Junior) authenticate without a Google account via the RFC 7523 jwt-bearer grant (SEP-990 ID-JAG): when configured, the token endpoint accepts short-lived assertions signed by a trusted issuer whose JWKS we pin, and mints a FastMCP -token marked ``fastmcp_grant == "id_jag"``. Those tokens skip the Google gates +token marked ``fastmcp_grant == "id_jag"``. Subjects must identify a Junior +actor as ``user:`` or ``system:``. Those tokens skip the Google gates and are audited as ``bot:``. Token refresh note: ``OAuthProxy`` re-calls ``_extract_upstream_claims`` on @@ -22,6 +23,7 @@ import asyncio import logging +import re import time from collections.abc import Callable from typing import Any, TypeGuard @@ -30,7 +32,10 @@ import jwt from fastmcp.exceptions import FastMCPError from fastmcp.server.auth import AccessToken, IdentityAssertion -from fastmcp.server.auth.identity_assertion import IdentityAssertionValidator +from fastmcp.server.auth.identity_assertion import ( + IdentityAssertionError, + IdentityAssertionValidator, +) from fastmcp.server.auth.providers.google import GoogleProvider from fastmcp.server.auth.providers.jwt import JWTVerifier from fastmcp.server.dependencies import get_access_token @@ -49,6 +54,7 @@ ACCESS_GROUP = "team@sentry.io" BOT_GRANT = "id_jag" BOT_JWKS_REFRESH_COOLDOWN_SECONDS = 30.0 +BOT_SUBJECT_PATTERN = re.compile(r"(?:user|system):[!-~]{1,200}") class GoogleGroupMembershipChecker: @@ -139,12 +145,25 @@ async def _get_jwks_key(self, kid: str | None) -> str: class CooldownIdentityAssertionValidator(IdentityAssertionValidator): - """IdentityAssertionValidator whose per-issuer verifiers rate-limit JWKS refreshes. + """Validate Junior subjects and rate-limit per-issuer JWKS refreshes. Mirrors FastMCP 4.0.10's private ``_get_verifier`` (lazy, cached per issuer in ``_verifiers``), swapping in ``CooldownJWTVerifier``. """ + async def validate( + self, assertion: str, *, client_id: str, resource_url: str | None + ) -> dict[str, Any]: + claims = await super().validate( + assertion, client_id=client_id, resource_url=resource_url + ) + subject = claims.get("sub") + if not isinstance(subject, str) or not BOT_SUBJECT_PATTERN.fullmatch(subject): + raise IdentityAssertionError( + "Assertion subject must identify a Junior user or system actor" + ) + return claims + async def _get_verifier(self, issuer: str) -> JWTVerifier: verifier = self._verifiers.get(issuer) if verifier is not None: diff --git a/src/firetower/mcp_server/tests/test_auth.py b/src/firetower/mcp_server/tests/test_auth.py index 2dab0234..bb8a0813 100644 --- a/src/firetower/mcp_server/tests/test_auth.py +++ b/src/firetower/mcp_server/tests/test_auth.py @@ -11,11 +11,14 @@ import pytest from cryptography.hazmat.primitives.asymmetric import rsa from fastmcp.exceptions import FastMCPError +from fastmcp.server.auth import IdentityAssertion +from fastmcp.server.auth.identity_assertion import IdentityAssertionError from jwt.algorithms import RSAAlgorithm from firetower.mcp_server import auth from firetower.mcp_server.auth import ( ACCESS_GROUP, + CooldownIdentityAssertionValidator, CooldownJWTVerifier, SentryGoogleProvider, requester_email, @@ -418,19 +421,30 @@ def test_fallback_rejects_when_no_token(monkeypatch): require_sentry_account() -def test_fallback_admits_bot_grant_and_audits_as_bot(monkeypatch): +@pytest.mark.parametrize( + ("subject", "audit_subject"), + [ + ("user:user-123", "bot:user:user-123"), + ("system:scheduled-automation", "bot:system:scheduled-automation"), + ], +) +def test_fallback_admits_bot_grant_and_audits_attributed_subject( + monkeypatch, subject, audit_subject +): monkeypatch.setattr( auth, "get_access_token", - lambda: _FakeToken({"fastmcp_grant": "id_jag"}, subject="firetower"), + lambda: _FakeToken({"fastmcp_grant": "id_jag"}, subject=subject), ) require_sentry_account() # no raise - assert requester_subject() == "bot:firetower" + assert requester_subject() == audit_subject assert requester_email() is None BOT_ISSUER = "https://junior.example" BOT_AUDIENCE = "https://mcp.example/" +BOT_CLIENT_ID = "bot-client" +BOT_RESOURCE = "https://mcp.example/mcp" def _jwk(key: rsa.RSAPrivateKey, kid: str) -> dict: @@ -441,20 +455,24 @@ def _jwk(key: rsa.RSAPrivateKey, kid: str) -> dict: } -def _bot_assertion(key: rsa.RSAPrivateKey, kid: str) -> str: +def _bot_assertion( + key: rsa.RSAPrivateKey, kid: str, subject: str = "user:user-123" +) -> str: now = int(time.time()) return jwt.encode( { "iss": BOT_ISSUER, - "sub": "firetower", + "sub": subject, "aud": BOT_AUDIENCE, "iat": now, "exp": now + 300, "jti": str(uuid.uuid4()), + "client_id": BOT_CLIENT_ID, + "resource": BOT_RESOURCE, }, key, algorithm="RS256", - headers={"kid": kid}, + headers={"kid": kid, "typ": "oauth-id-jag+jwt"}, ) @@ -494,6 +512,56 @@ def _cooldown_verifier(jwks: _JWKS, clock: _Clock) -> CooldownJWTVerifier: return verifier +def _identity_assertion_validator(jwks: _JWKS) -> CooldownIdentityAssertionValidator: + validator = CooldownIdentityAssertionValidator( + config=IdentityAssertion( + trusted_issuers=[BOT_ISSUER], + jwks_uris={BOT_ISSUER: f"{BOT_ISSUER}/jwks.json"}, + ), + audience=BOT_AUDIENCE, + ) + validator._verifiers[BOT_ISSUER] = _cooldown_verifier(jwks, _Clock()) + return validator + + +@pytest.mark.parametrize("subject", ["user:user-123", "system:scheduled-automation"]) +def test_identity_assertion_validator_admits_attributed_subjects(subject): + validator = _identity_assertion_validator(_JWKS(_jwk(_KEY, "current"))) + + claims = asyncio.run( + validator.validate( + _bot_assertion(_KEY, "current", subject), + client_id=BOT_CLIENT_ID, + resource_url=BOT_RESOURCE, + ) + ) + + assert claims["sub"] == subject + + +@pytest.mark.parametrize( + "subject", + [ + "firetower", + "user:", + "bot:user-123", + "user:user-123\nforged-log-line", + f"user:{'a' * 201}", + ], +) +def test_identity_assertion_validator_rejects_unattributed_subjects(subject): + validator = _identity_assertion_validator(_JWKS(_jwk(_KEY, "current"))) + + with pytest.raises(IdentityAssertionError, match="Junior user or system actor"): + asyncio.run( + validator.validate( + _bot_assertion(_KEY, "current", subject), + client_id=BOT_CLIENT_ID, + resource_url=BOT_RESOURCE, + ) + ) + + def test_cooldown_verifier_limits_unknown_kid_refreshes_and_keeps_known_keys_fast(): jwks = _JWKS(_jwk(_KEY, "current")) verifier = _cooldown_verifier(jwks, _Clock()) @@ -509,7 +577,7 @@ async def run() -> list: known, *unknown, known_again = asyncio.run(run()) - assert known is not None and known.claims["sub"] == "firetower" + assert known is not None and known.claims["sub"] == "user:user-123" assert unknown == [None] * 5 assert known_again is not None assert jwks.fetches == 1 diff --git a/src/firetower/mcp_server/tests/test_server.py b/src/firetower/mcp_server/tests/test_server.py index f5e18f71..0d60da2b 100644 --- a/src/firetower/mcp_server/tests/test_server.py +++ b/src/firetower/mcp_server/tests/test_server.py @@ -114,6 +114,7 @@ def _bot_token( client: TestClient, signing_key: rsa.RSAPrivateKey = BOT_KEY, kid: str = "junior-test", + subject: str = "user:user-123", ) -> httpx.Response: registration = client.post( "/register", @@ -128,7 +129,7 @@ def _bot_token( assertion = jwt.encode( { "iss": BOT_ISSUER, - "sub": "firetower", + "sub": subject, "aud": f"{BASE_URL}/", "iat": now, "exp": now + 300, @@ -356,7 +357,7 @@ def test_bot_assertion_mints_token_that_can_call_read_tools( [audit] = [ r for r in caplog.records if getattr(r, "event", None) == "mcp_tool_call" ] - assert audit.actor_sub == "bot:firetower" + assert audit.actor_sub == "bot:user:user-123" assert audit.actor_email is None