dr3s
547 lines
21 KiB
Python
547 lines
21 KiB
Python
"""
|
|
Tests for refactored OAuth client authentication implementation.
|
|
"""
|
|
|
|
import time
|
|
|
|
import httpx
|
|
import pytest
|
|
from pydantic import AnyHttpUrl, AnyUrl
|
|
|
|
from mcp.client.auth import OAuthClientProvider, PKCEParameters
|
|
from mcp.shared.auth import (
|
|
OAuthClientInformationFull,
|
|
OAuthClientMetadata,
|
|
OAuthToken,
|
|
ProtectedResourceMetadata,
|
|
)
|
|
|
|
|
|
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 valid_tokens():
|
|
return OAuthToken(
|
|
access_token="test_access_token",
|
|
token_type="Bearer",
|
|
expires_in=3600,
|
|
refresh_token="test_refresh_token",
|
|
scope="read write",
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def oauth_provider(client_metadata, mock_storage):
|
|
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 OAuthClientProvider(
|
|
server_url="https://api.example.com/v1/mcp",
|
|
client_metadata=client_metadata,
|
|
storage=mock_storage,
|
|
redirect_handler=redirect_handler,
|
|
callback_handler=callback_handler,
|
|
)
|
|
|
|
|
|
class TestPKCEParameters:
|
|
"""Test PKCE parameter generation."""
|
|
|
|
def test_pkce_generation(self):
|
|
"""Test PKCE parameter generation creates valid values."""
|
|
pkce = PKCEParameters.generate()
|
|
|
|
# Verify lengths
|
|
assert len(pkce.code_verifier) == 128
|
|
assert 43 <= len(pkce.code_challenge) <= 128
|
|
|
|
# Verify characters used in verifier
|
|
allowed_chars = set("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-._~")
|
|
assert all(c in allowed_chars for c in pkce.code_verifier)
|
|
|
|
# Verify base64url encoding in challenge (no padding)
|
|
assert "=" not in pkce.code_challenge
|
|
|
|
def test_pkce_uniqueness(self):
|
|
"""Test PKCE generates unique values each time."""
|
|
pkce1 = PKCEParameters.generate()
|
|
pkce2 = PKCEParameters.generate()
|
|
|
|
assert pkce1.code_verifier != pkce2.code_verifier
|
|
assert pkce1.code_challenge != pkce2.code_challenge
|
|
|
|
|
|
class TestOAuthContext:
|
|
"""Test OAuth context functionality."""
|
|
|
|
@pytest.mark.anyio
|
|
async def test_oauth_provider_initialization(self, oauth_provider, client_metadata, mock_storage):
|
|
"""Test OAuthClientProvider basic setup."""
|
|
assert oauth_provider.context.server_url == "https://api.example.com/v1/mcp"
|
|
assert oauth_provider.context.client_metadata == client_metadata
|
|
assert oauth_provider.context.storage == mock_storage
|
|
assert oauth_provider.context.timeout == 300.0
|
|
assert oauth_provider.context is not None
|
|
|
|
def test_context_url_parsing(self, oauth_provider):
|
|
"""Test get_authorization_base_url() extracts base URLs correctly."""
|
|
context = oauth_provider.context
|
|
|
|
# Test with path
|
|
assert context.get_authorization_base_url("https://api.example.com/v1/mcp") == "https://api.example.com"
|
|
|
|
# Test with no path
|
|
assert context.get_authorization_base_url("https://api.example.com") == "https://api.example.com"
|
|
|
|
# Test with port
|
|
assert (
|
|
context.get_authorization_base_url("https://api.example.com:8080/path/to/mcp")
|
|
== "https://api.example.com:8080"
|
|
)
|
|
|
|
# Test with query params
|
|
assert (
|
|
context.get_authorization_base_url("https://api.example.com/path?param=value") == "https://api.example.com"
|
|
)
|
|
|
|
@pytest.mark.anyio
|
|
async def test_token_validity_checking(self, oauth_provider, mock_storage, valid_tokens):
|
|
"""Test is_token_valid() and can_refresh_token() logic."""
|
|
context = oauth_provider.context
|
|
|
|
# No tokens - should be invalid
|
|
assert not context.is_token_valid()
|
|
assert not context.can_refresh_token()
|
|
|
|
# Set valid tokens and client info
|
|
context.current_tokens = valid_tokens
|
|
context.token_expiry_time = time.time() + 1800 # 30 minutes from now
|
|
context.client_info = OAuthClientInformationFull(
|
|
client_id="test_client_id",
|
|
client_secret="test_client_secret",
|
|
redirect_uris=[AnyUrl("http://localhost:3030/callback")],
|
|
)
|
|
|
|
# Should be valid
|
|
assert context.is_token_valid()
|
|
assert context.can_refresh_token() # Has refresh token and client info
|
|
|
|
# Expire the token
|
|
context.token_expiry_time = time.time() - 100 # Expired 100 seconds ago
|
|
assert not context.is_token_valid()
|
|
assert context.can_refresh_token() # Can still refresh
|
|
|
|
# Remove refresh token
|
|
context.current_tokens.refresh_token = None
|
|
assert not context.can_refresh_token()
|
|
|
|
# Remove client info
|
|
context.current_tokens.refresh_token = "test_refresh_token"
|
|
context.client_info = None
|
|
assert not context.can_refresh_token()
|
|
|
|
def test_clear_tokens(self, oauth_provider, valid_tokens):
|
|
"""Test clear_tokens() removes token data."""
|
|
context = oauth_provider.context
|
|
context.current_tokens = valid_tokens
|
|
context.token_expiry_time = time.time() + 1800
|
|
|
|
# Clear tokens
|
|
context.clear_tokens()
|
|
|
|
# Verify cleared
|
|
assert context.current_tokens is None
|
|
assert context.token_expiry_time is None
|
|
|
|
|
|
class TestOAuthFlow:
|
|
"""Test OAuth flow methods."""
|
|
|
|
@pytest.mark.anyio
|
|
async def test_discover_protected_resource_request(self, oauth_provider):
|
|
"""Test protected resource discovery request building."""
|
|
request = await oauth_provider._discover_protected_resource()
|
|
|
|
assert request.method == "GET"
|
|
assert str(request.url) == "https://api.example.com/.well-known/oauth-protected-resource"
|
|
assert "mcp-protocol-version" in request.headers
|
|
|
|
@pytest.mark.anyio
|
|
async def test_discover_oauth_metadata_request(self, oauth_provider):
|
|
"""Test OAuth metadata discovery request building."""
|
|
request = await oauth_provider._discover_oauth_metadata()
|
|
|
|
assert request.method == "GET"
|
|
assert str(request.url) == "https://api.example.com/.well-known/oauth-authorization-server/v1/mcp"
|
|
assert "mcp-protocol-version" in request.headers
|
|
|
|
@pytest.mark.anyio
|
|
async def test_discover_oauth_metadata_request_no_path(self, client_metadata, mock_storage):
|
|
"""Test OAuth metadata discovery request building when server has no path."""
|
|
|
|
async def redirect_handler(url: str) -> None:
|
|
pass
|
|
|
|
async def callback_handler() -> tuple[str, str | None]:
|
|
return "test_auth_code", "test_state"
|
|
|
|
provider = OAuthClientProvider(
|
|
server_url="https://api.example.com",
|
|
client_metadata=client_metadata,
|
|
storage=mock_storage,
|
|
redirect_handler=redirect_handler,
|
|
callback_handler=callback_handler,
|
|
)
|
|
|
|
request = await provider._discover_oauth_metadata()
|
|
|
|
assert request.method == "GET"
|
|
assert str(request.url) == "https://api.example.com/.well-known/oauth-authorization-server"
|
|
assert "mcp-protocol-version" in request.headers
|
|
|
|
@pytest.mark.anyio
|
|
async def test_discover_oauth_metadata_request_trailing_slash(self, client_metadata, mock_storage):
|
|
"""Test OAuth metadata discovery request building when server path has trailing slash."""
|
|
|
|
async def redirect_handler(url: str) -> None:
|
|
pass
|
|
|
|
async def callback_handler() -> tuple[str, str | None]:
|
|
return "test_auth_code", "test_state"
|
|
|
|
provider = OAuthClientProvider(
|
|
server_url="https://api.example.com/v1/mcp/",
|
|
client_metadata=client_metadata,
|
|
storage=mock_storage,
|
|
redirect_handler=redirect_handler,
|
|
callback_handler=callback_handler,
|
|
)
|
|
|
|
request = await provider._discover_oauth_metadata()
|
|
|
|
assert request.method == "GET"
|
|
assert str(request.url) == "https://api.example.com/.well-known/oauth-authorization-server/v1/mcp"
|
|
assert "mcp-protocol-version" in request.headers
|
|
|
|
|
|
class TestOAuthFallback:
|
|
"""Test OAuth discovery fallback behavior for legacy (act as AS not RS) servers."""
|
|
|
|
@pytest.mark.anyio
|
|
async def test_fallback_discovery_request(self, client_metadata, mock_storage):
|
|
"""Test fallback discovery request building."""
|
|
|
|
async def redirect_handler(url: str) -> None:
|
|
pass
|
|
|
|
async def callback_handler() -> tuple[str, str | None]:
|
|
return "test_auth_code", "test_state"
|
|
|
|
provider = OAuthClientProvider(
|
|
server_url="https://api.example.com/v1/mcp",
|
|
client_metadata=client_metadata,
|
|
storage=mock_storage,
|
|
redirect_handler=redirect_handler,
|
|
callback_handler=callback_handler,
|
|
)
|
|
|
|
# Set up discovery state manually as if path-aware discovery was attempted
|
|
provider.context.discovery_base_url = "https://api.example.com"
|
|
provider.context.discovery_pathname = "/v1/mcp"
|
|
|
|
# Test fallback request building
|
|
request = await provider._discover_oauth_metadata_fallback()
|
|
|
|
assert request.method == "GET"
|
|
assert str(request.url) == "https://api.example.com/.well-known/oauth-authorization-server"
|
|
assert "mcp-protocol-version" in request.headers
|
|
|
|
@pytest.mark.anyio
|
|
async def test_should_attempt_fallback(self, oauth_provider):
|
|
"""Test fallback decision logic."""
|
|
# Should attempt fallback on 404 with non-root path
|
|
assert oauth_provider._should_attempt_fallback(404, "/v1/mcp")
|
|
|
|
# Should NOT attempt fallback on 404 with root path
|
|
assert not oauth_provider._should_attempt_fallback(404, "/")
|
|
|
|
# Should NOT attempt fallback on other status codes
|
|
assert not oauth_provider._should_attempt_fallback(200, "/v1/mcp")
|
|
assert not oauth_provider._should_attempt_fallback(500, "/v1/mcp")
|
|
|
|
@pytest.mark.anyio
|
|
async def test_handle_metadata_response_success(self, oauth_provider):
|
|
"""Test successful metadata response handling."""
|
|
# Create minimal valid OAuth metadata
|
|
content = b"""{
|
|
"issuer": "https://auth.example.com",
|
|
"authorization_endpoint": "https://auth.example.com/authorize",
|
|
"token_endpoint": "https://auth.example.com/token"
|
|
}"""
|
|
response = httpx.Response(200, content=content)
|
|
|
|
# Should return True (success) and set metadata
|
|
result = await oauth_provider._handle_oauth_metadata_response(response, is_fallback=False)
|
|
assert result is True
|
|
assert oauth_provider.context.oauth_metadata is not None
|
|
assert str(oauth_provider.context.oauth_metadata.issuer) == "https://auth.example.com/"
|
|
|
|
@pytest.mark.anyio
|
|
async def test_handle_metadata_response_404_needs_fallback(self, oauth_provider):
|
|
"""Test 404 response handling that should trigger fallback."""
|
|
# Set up discovery state for non-root path
|
|
oauth_provider.context.discovery_base_url = "https://api.example.com"
|
|
oauth_provider.context.discovery_pathname = "/v1/mcp"
|
|
|
|
# Mock 404 response
|
|
response = httpx.Response(404)
|
|
|
|
# Should return False (needs fallback)
|
|
result = await oauth_provider._handle_oauth_metadata_response(response, is_fallback=False)
|
|
assert result is False
|
|
|
|
@pytest.mark.anyio
|
|
async def test_handle_metadata_response_404_no_fallback_needed(self, oauth_provider):
|
|
"""Test 404 response handling when no fallback is needed."""
|
|
# Set up discovery state for root path
|
|
oauth_provider.context.discovery_base_url = "https://api.example.com"
|
|
oauth_provider.context.discovery_pathname = "/"
|
|
|
|
# Mock 404 response
|
|
response = httpx.Response(404)
|
|
|
|
# Should return True (no fallback needed)
|
|
result = await oauth_provider._handle_oauth_metadata_response(response, is_fallback=False)
|
|
assert result is True
|
|
|
|
@pytest.mark.anyio
|
|
async def test_handle_metadata_response_404_fallback_attempt(self, oauth_provider):
|
|
"""Test 404 response handling during fallback attempt."""
|
|
# Mock 404 response during fallback
|
|
response = httpx.Response(404)
|
|
|
|
# Should return True (fallback attempt complete, no further action needed)
|
|
result = await oauth_provider._handle_oauth_metadata_response(response, is_fallback=True)
|
|
assert result is True
|
|
|
|
@pytest.mark.anyio
|
|
async def test_register_client_request(self, oauth_provider):
|
|
"""Test client registration request building."""
|
|
request = await oauth_provider._register_client()
|
|
|
|
assert request is not None
|
|
assert request.method == "POST"
|
|
assert str(request.url) == "https://api.example.com/register"
|
|
assert request.headers["Content-Type"] == "application/json"
|
|
|
|
@pytest.mark.anyio
|
|
async def test_register_client_skip_if_registered(self, oauth_provider, mock_storage):
|
|
"""Test client registration is skipped if already registered."""
|
|
# Set existing client info
|
|
client_info = OAuthClientInformationFull(
|
|
client_id="existing_client",
|
|
redirect_uris=[AnyUrl("http://localhost:3030/callback")],
|
|
)
|
|
oauth_provider.context.client_info = client_info
|
|
|
|
# Should return None (skip registration)
|
|
request = await oauth_provider._register_client()
|
|
assert request is None
|
|
|
|
@pytest.mark.anyio
|
|
async def test_token_exchange_request(self, oauth_provider):
|
|
"""Test token exchange request building."""
|
|
# Set up required context
|
|
oauth_provider.context.client_info = OAuthClientInformationFull(
|
|
client_id="test_client",
|
|
client_secret="test_secret",
|
|
redirect_uris=[AnyUrl("http://localhost:3030/callback")],
|
|
)
|
|
|
|
request = await oauth_provider._exchange_token("test_auth_code", "test_verifier")
|
|
|
|
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 = request.content.decode()
|
|
assert "grant_type=authorization_code" in content
|
|
assert "code=test_auth_code" in content
|
|
assert "code_verifier=test_verifier" in content
|
|
assert "client_id=test_client" in content
|
|
assert "client_secret=test_secret" in content
|
|
|
|
@pytest.mark.anyio
|
|
async def test_refresh_token_request(self, oauth_provider, valid_tokens):
|
|
"""Test refresh token request building."""
|
|
# Set up required context
|
|
oauth_provider.context.current_tokens = valid_tokens
|
|
oauth_provider.context.client_info = OAuthClientInformationFull(
|
|
client_id="test_client",
|
|
client_secret="test_secret",
|
|
redirect_uris=[AnyUrl("http://localhost:3030/callback")],
|
|
)
|
|
|
|
request = await oauth_provider._refresh_token()
|
|
|
|
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 = request.content.decode()
|
|
assert "grant_type=refresh_token" in content
|
|
assert "refresh_token=test_refresh_token" in content
|
|
assert "client_id=test_client" in content
|
|
assert "client_secret=test_secret" in content
|
|
|
|
|
|
class TestProtectedResourceMetadata:
|
|
"""Test protected resource handling."""
|
|
|
|
@pytest.mark.anyio
|
|
async def test_resource_param_included_with_recent_protocol_version(self, oauth_provider: OAuthClientProvider):
|
|
"""Test resource parameter is included for protocol version >= 2025-06-18."""
|
|
# Set protocol version to 2025-06-18
|
|
oauth_provider.context.protocol_version = "2025-06-18"
|
|
oauth_provider.context.client_info = OAuthClientInformationFull(
|
|
client_id="test_client",
|
|
client_secret="test_secret",
|
|
redirect_uris=[AnyUrl("http://localhost:3030/callback")],
|
|
)
|
|
|
|
# Test in token exchange
|
|
request = await oauth_provider._exchange_token("test_code", "test_verifier")
|
|
content = request.content.decode()
|
|
assert "resource=" in content
|
|
# Check URL-encoded resource parameter
|
|
from urllib.parse import quote
|
|
|
|
expected_resource = quote(oauth_provider.context.get_resource_url(), safe="")
|
|
assert f"resource={expected_resource}" in content
|
|
|
|
# Test in refresh token
|
|
oauth_provider.context.current_tokens = OAuthToken(
|
|
access_token="test_access",
|
|
token_type="Bearer",
|
|
refresh_token="test_refresh",
|
|
)
|
|
refresh_request = await oauth_provider._refresh_token()
|
|
refresh_content = refresh_request.content.decode()
|
|
assert "resource=" in refresh_content
|
|
|
|
@pytest.mark.anyio
|
|
async def test_resource_param_excluded_with_old_protocol_version(self, oauth_provider: OAuthClientProvider):
|
|
"""Test resource parameter is excluded for protocol version < 2025-06-18."""
|
|
# Set protocol version to older version
|
|
oauth_provider.context.protocol_version = "2025-03-26"
|
|
oauth_provider.context.client_info = OAuthClientInformationFull(
|
|
client_id="test_client",
|
|
client_secret="test_secret",
|
|
redirect_uris=[AnyUrl("http://localhost:3030/callback")],
|
|
)
|
|
|
|
# Test in token exchange
|
|
request = await oauth_provider._exchange_token("test_code", "test_verifier")
|
|
content = request.content.decode()
|
|
assert "resource=" not in content
|
|
|
|
# Test in refresh token
|
|
oauth_provider.context.current_tokens = OAuthToken(
|
|
access_token="test_access",
|
|
token_type="Bearer",
|
|
refresh_token="test_refresh",
|
|
)
|
|
refresh_request = await oauth_provider._refresh_token()
|
|
refresh_content = refresh_request.content.decode()
|
|
assert "resource=" not in refresh_content
|
|
|
|
@pytest.mark.anyio
|
|
async def test_resource_param_included_with_protected_resource_metadata(self, oauth_provider: OAuthClientProvider):
|
|
"""Test resource parameter is always included when protected resource metadata exists."""
|
|
# Set old protocol version but with protected resource metadata
|
|
oauth_provider.context.protocol_version = "2025-03-26"
|
|
oauth_provider.context.protected_resource_metadata = ProtectedResourceMetadata(
|
|
resource=AnyHttpUrl("https://api.example.com/v1/mcp"),
|
|
authorization_servers=[AnyHttpUrl("https://api.example.com")],
|
|
)
|
|
oauth_provider.context.client_info = OAuthClientInformationFull(
|
|
client_id="test_client",
|
|
client_secret="test_secret",
|
|
redirect_uris=[AnyUrl("http://localhost:3030/callback")],
|
|
)
|
|
|
|
# Test in token exchange
|
|
request = await oauth_provider._exchange_token("test_code", "test_verifier")
|
|
content = request.content.decode()
|
|
assert "resource=" in content
|
|
|
|
|
|
class TestAuthFlow:
|
|
"""Test the auth flow in httpx."""
|
|
|
|
@pytest.mark.anyio
|
|
async def test_auth_flow_with_valid_tokens(self, oauth_provider, mock_storage, valid_tokens):
|
|
"""Test auth flow when tokens are already valid."""
|
|
# Pre-store valid tokens
|
|
await mock_storage.set_tokens(valid_tokens)
|
|
oauth_provider.context.current_tokens = valid_tokens
|
|
oauth_provider.context.token_expiry_time = time.time() + 1800
|
|
oauth_provider._initialized = True
|
|
|
|
# Create a test request
|
|
test_request = httpx.Request("GET", "https://api.example.com/test")
|
|
|
|
# Mock the auth flow
|
|
auth_flow = oauth_provider.async_auth_flow(test_request)
|
|
|
|
# Should get the request with auth header added
|
|
request = await auth_flow.__anext__()
|
|
assert request.headers["Authorization"] == "Bearer test_access_token"
|
|
|
|
# Send a successful response
|
|
response = httpx.Response(200)
|
|
try:
|
|
await auth_flow.asend(response)
|
|
except StopAsyncIteration:
|
|
pass # Expected
|