Implement RFC 7523 JWT flows (#1247)
Co-authored-by: Yann Jouanin <yann.jouanin@valueandco.com>
This commit is contained in:
Luca Chang
GitHub
parent
db9e451551
commit
f161149680
3 changed files with 463 additions and 52 deletions
163
tests/client/auth/extensions/test_client_credentials.py
Normal file
163
tests/client/auth/extensions/test_client_credentials.py
Normal file
|
|
@ -0,0 +1,163 @@
|
|||
import urllib.parse
|
||||
|
||||
import jwt
|
||||
import pytest
|
||||
from pydantic import AnyHttpUrl, AnyUrl
|
||||
|
||||
from mcp.client.auth.extensions.client_credentials import JWTParameters, RFC7523OAuthClientProvider
|
||||
from mcp.shared.auth import OAuthClientInformationFull, OAuthClientMetadata, OAuthMetadata, OAuthToken
|
||||
|
||||
|
||||
class MockTokenStorage:
|
||||
"""Mock token storage for testing."""
|
||||
|
||||
def __init__(self):
|
||||
self._tokens: OAuthToken | None = None
|
||||
self._client_info: OAuthClientInformationFull | None = None
|
||||
|
||||
async def get_tokens(self) -> OAuthToken | None:
|
||||
return self._tokens
|
||||
|
||||
async def set_tokens(self, tokens: OAuthToken) -> None:
|
||||
self._tokens = tokens
|
||||
|
||||
async def get_client_info(self) -> OAuthClientInformationFull | None:
|
||||
return self._client_info
|
||||
|
||||
async def set_client_info(self, client_info: OAuthClientInformationFull) -> None:
|
||||
self._client_info = client_info
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_storage():
|
||||
return MockTokenStorage()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client_metadata():
|
||||
return OAuthClientMetadata(
|
||||
client_name="Test Client",
|
||||
client_uri=AnyHttpUrl("https://example.com"),
|
||||
redirect_uris=[AnyUrl("http://localhost:3030/callback")],
|
||||
scope="read write",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def rfc7523_oauth_provider(client_metadata: OAuthClientMetadata, mock_storage: MockTokenStorage):
|
||||
async def redirect_handler(url: str) -> None:
|
||||
"""Mock redirect handler."""
|
||||
pass
|
||||
|
||||
async def callback_handler() -> tuple[str, str | None]:
|
||||
"""Mock callback handler."""
|
||||
return "test_auth_code", "test_state"
|
||||
|
||||
return RFC7523OAuthClientProvider(
|
||||
server_url="https://api.example.com/v1/mcp",
|
||||
client_metadata=client_metadata,
|
||||
storage=mock_storage,
|
||||
redirect_handler=redirect_handler,
|
||||
callback_handler=callback_handler,
|
||||
)
|
||||
|
||||
|
||||
class TestOAuthFlowClientCredentials:
|
||||
"""Test OAuth flow behavior for client credentials flows."""
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_token_exchange_request_jwt_predefined(self, rfc7523_oauth_provider: RFC7523OAuthClientProvider):
|
||||
"""Test token exchange request building with a predefined JWT assertion."""
|
||||
# Set up required context
|
||||
rfc7523_oauth_provider.context.client_info = OAuthClientInformationFull(
|
||||
grant_types=["urn:ietf:params:oauth:grant-type:jwt-bearer"],
|
||||
token_endpoint_auth_method="private_key_jwt",
|
||||
redirect_uris=None,
|
||||
scope="read write",
|
||||
)
|
||||
rfc7523_oauth_provider.context.oauth_metadata = OAuthMetadata(
|
||||
issuer=AnyHttpUrl("https://api.example.com"),
|
||||
authorization_endpoint=AnyHttpUrl("https://api.example.com/authorize"),
|
||||
token_endpoint=AnyHttpUrl("https://api.example.com/token"),
|
||||
registration_endpoint=AnyHttpUrl("https://api.example.com/register"),
|
||||
)
|
||||
rfc7523_oauth_provider.context.client_metadata = rfc7523_oauth_provider.context.client_info
|
||||
rfc7523_oauth_provider.context.protocol_version = "2025-06-18"
|
||||
rfc7523_oauth_provider.jwt_parameters = JWTParameters(
|
||||
# https://www.jwt.io
|
||||
assertion="eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiYWRtaW4iOnRydWUsImlhdCI6MTUxNjIzOTAyMn0.KMUFsIDTnFmyG3nMiGM6H9FNFUROf3wh7SmqJp-QV30"
|
||||
)
|
||||
|
||||
request = await rfc7523_oauth_provider._exchange_token_jwt_bearer()
|
||||
|
||||
assert request.method == "POST"
|
||||
assert str(request.url) == "https://api.example.com/token"
|
||||
assert request.headers["Content-Type"] == "application/x-www-form-urlencoded"
|
||||
|
||||
# Check form data
|
||||
content = urllib.parse.unquote_plus(request.content.decode())
|
||||
assert "grant_type=urn:ietf:params:oauth:grant-type:jwt-bearer" in content
|
||||
assert "scope=read write" in content
|
||||
assert "resource=https://api.example.com/v1/mcp" in content
|
||||
assert (
|
||||
"assertion=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiYWRtaW4iOnRydWUsImlhdCI6MTUxNjIzOTAyMn0.KMUFsIDTnFmyG3nMiGM6H9FNFUROf3wh7SmqJp-QV30"
|
||||
in content
|
||||
)
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_token_exchange_request_jwt(self, rfc7523_oauth_provider: RFC7523OAuthClientProvider):
|
||||
"""Test token exchange request building wiith a generated JWT assertion."""
|
||||
# Set up required context
|
||||
rfc7523_oauth_provider.context.client_info = OAuthClientInformationFull(
|
||||
grant_types=["urn:ietf:params:oauth:grant-type:jwt-bearer"],
|
||||
token_endpoint_auth_method="private_key_jwt",
|
||||
redirect_uris=None,
|
||||
scope="read write",
|
||||
)
|
||||
rfc7523_oauth_provider.context.oauth_metadata = OAuthMetadata(
|
||||
issuer=AnyHttpUrl("https://api.example.com"),
|
||||
authorization_endpoint=AnyHttpUrl("https://api.example.com/authorize"),
|
||||
token_endpoint=AnyHttpUrl("https://api.example.com/token"),
|
||||
registration_endpoint=AnyHttpUrl("https://api.example.com/register"),
|
||||
)
|
||||
rfc7523_oauth_provider.context.client_metadata = rfc7523_oauth_provider.context.client_info
|
||||
rfc7523_oauth_provider.context.protocol_version = "2025-06-18"
|
||||
rfc7523_oauth_provider.jwt_parameters = JWTParameters(
|
||||
issuer="foo",
|
||||
subject="1234567890",
|
||||
claims={
|
||||
"name": "John Doe",
|
||||
"admin": True,
|
||||
"iat": 1516239022,
|
||||
},
|
||||
jwt_signing_algorithm="HS256",
|
||||
jwt_signing_key="a-string-secret-at-least-256-bits-long",
|
||||
jwt_lifetime_seconds=300,
|
||||
)
|
||||
|
||||
request = await rfc7523_oauth_provider._exchange_token_jwt_bearer()
|
||||
|
||||
assert request.method == "POST"
|
||||
assert str(request.url) == "https://api.example.com/token"
|
||||
assert request.headers["Content-Type"] == "application/x-www-form-urlencoded"
|
||||
|
||||
# Check form data
|
||||
content = urllib.parse.unquote_plus(request.content.decode()).split("&")
|
||||
assert "grant_type=urn:ietf:params:oauth:grant-type:jwt-bearer" in content
|
||||
assert "scope=read write" in content
|
||||
assert "resource=https://api.example.com/v1/mcp" in content
|
||||
|
||||
# Check assertion
|
||||
assertion = next(param for param in content if param.startswith("assertion="))[len("assertion=") :]
|
||||
claims = jwt.decode(
|
||||
assertion,
|
||||
key="a-string-secret-at-least-256-bits-long",
|
||||
algorithms=["HS256"],
|
||||
audience="https://api.example.com/",
|
||||
subject="1234567890",
|
||||
issuer="foo",
|
||||
verify=True,
|
||||
)
|
||||
assert claims["name"] == "John Doe"
|
||||
assert claims["admin"]
|
||||
assert claims["iat"] == 1516239022
|
||||
|
|
@ -391,7 +391,9 @@ class TestOAuthFallback:
|
|||
)
|
||||
|
||||
# Mock the authorization process to minimize unnecessary state in this test
|
||||
oauth_provider._perform_authorization = mock.AsyncMock(return_value=("test_auth_code", "test_code_verifier"))
|
||||
oauth_provider._perform_authorization_code_grant = mock.AsyncMock(
|
||||
return_value=("test_auth_code", "test_code_verifier")
|
||||
)
|
||||
|
||||
# Next request should fall back to legacy behavior and auth with the RS (mocked /authorize, next is /token)
|
||||
token_request = await auth_flow.asend(oauth_metadata_response_3)
|
||||
|
|
@ -513,7 +515,7 @@ class TestOAuthFallback:
|
|||
assert request is None
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_token_exchange_request(self, oauth_provider: OAuthClientProvider):
|
||||
async def test_token_exchange_request_authorization_code(self, oauth_provider: OAuthClientProvider):
|
||||
"""Test token exchange request building."""
|
||||
# Set up required context
|
||||
oauth_provider.context.client_info = OAuthClientInformationFull(
|
||||
|
|
@ -522,7 +524,7 @@ class TestOAuthFallback:
|
|||
redirect_uris=[AnyUrl("http://localhost:3030/callback")],
|
||||
)
|
||||
|
||||
request = await oauth_provider._exchange_token("test_auth_code", "test_verifier")
|
||||
request = await oauth_provider._exchange_token_authorization_code("test_auth_code", "test_verifier")
|
||||
|
||||
assert request.method == "POST"
|
||||
assert str(request.url) == "https://api.example.com/token"
|
||||
|
|
@ -576,7 +578,7 @@ class TestProtectedResourceMetadata:
|
|||
)
|
||||
|
||||
# Test in token exchange
|
||||
request = await oauth_provider._exchange_token("test_code", "test_verifier")
|
||||
request = await oauth_provider._exchange_token_authorization_code("test_code", "test_verifier")
|
||||
content = request.content.decode()
|
||||
assert "resource=" in content
|
||||
# Check URL-encoded resource parameter
|
||||
|
|
@ -607,7 +609,7 @@ class TestProtectedResourceMetadata:
|
|||
)
|
||||
|
||||
# Test in token exchange
|
||||
request = await oauth_provider._exchange_token("test_code", "test_verifier")
|
||||
request = await oauth_provider._exchange_token_authorization_code("test_code", "test_verifier")
|
||||
content = request.content.decode()
|
||||
assert "resource=" not in content
|
||||
|
||||
|
|
@ -637,7 +639,7 @@ class TestProtectedResourceMetadata:
|
|||
)
|
||||
|
||||
# Test in token exchange
|
||||
request = await oauth_provider._exchange_token("test_code", "test_verifier")
|
||||
request = await oauth_provider._exchange_token_authorization_code("test_code", "test_verifier")
|
||||
content = request.content.decode()
|
||||
assert "resource=" in content
|
||||
|
||||
|
|
@ -710,7 +712,7 @@ class TestAuthFlow:
|
|||
pass # Expected
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_auth_flow_with_no_tokens(self, oauth_provider: OAuthClientProvider):
|
||||
async def test_auth_flow_with_no_tokens(self, oauth_provider: OAuthClientProvider, mock_storage: MockTokenStorage):
|
||||
"""Test auth flow when no tokens are available, triggering the full OAuth flow."""
|
||||
# Ensure no tokens are stored
|
||||
oauth_provider.context.current_tokens = None
|
||||
|
|
@ -779,7 +781,9 @@ class TestAuthFlow:
|
|||
)
|
||||
|
||||
# Mock the authorization process
|
||||
oauth_provider._perform_authorization = mock.AsyncMock(return_value=("test_auth_code", "test_code_verifier"))
|
||||
oauth_provider._perform_authorization_code_grant = mock.AsyncMock(
|
||||
return_value=("test_auth_code", "test_code_verifier")
|
||||
)
|
||||
|
||||
# Next request should be to exchange token
|
||||
token_request = await auth_flow.asend(registration_response)
|
||||
|
|
|
|||
|
|
@ -39,11 +39,13 @@ class MockOAuthProvider(OAuthAuthorizationServerProvider[AuthorizationCode, Refr
|
|||
return self.clients.get(client_id)
|
||||
|
||||
async def register_client(self, client_info: OAuthClientInformationFull):
|
||||
assert client_info.client_id is not None
|
||||
self.clients[client_info.client_id] = client_info
|
||||
|
||||
async def authorize(self, client: OAuthClientInformationFull, params: AuthorizationParams) -> str:
|
||||
# toy authorize implementation which just immediately generates an authorization
|
||||
# code and completes the redirect
|
||||
assert client.client_id is not None
|
||||
code = AuthorizationCode(
|
||||
code=f"code_{int(time.time())}",
|
||||
client_id=client.client_id,
|
||||
|
|
@ -72,6 +74,7 @@ class MockOAuthProvider(OAuthAuthorizationServerProvider[AuthorizationCode, Refr
|
|||
refresh_token = f"refresh_{secrets.token_hex(32)}"
|
||||
|
||||
# Store the tokens
|
||||
assert client.client_id is not None
|
||||
self.tokens[access_token] = AccessToken(
|
||||
token=access_token,
|
||||
client_id=client.client_id,
|
||||
|
|
@ -133,6 +136,7 @@ class MockOAuthProvider(OAuthAuthorizationServerProvider[AuthorizationCode, Refr
|
|||
new_refresh_token = f"refresh_{secrets.token_hex(32)}"
|
||||
|
||||
# Store the new tokens
|
||||
assert client.client_id is not None
|
||||
self.tokens[new_access_token] = AccessToken(
|
||||
token=new_access_token,
|
||||
client_id=client.client_id,
|
||||
|
|
|
|||
Loading…
Reference in a new issue