diff --git a/.idea/.gitignore b/.idea/.gitignore new file mode 100644 index 0000000..cc09cda --- /dev/null +++ b/.idea/.gitignore @@ -0,0 +1,9 @@ +# 默认忽略的文件 +/shelf/ +/workspace.xml +# 基于编辑器的 HTTP 客户端请求 +/httpRequests/ +# Datasource local storage ignored files +/dataSources/ +/dataSources.local.xml +.idea/* diff --git a/.idea/dictionaries/project.xml b/.idea/dictionaries/project.xml new file mode 100644 index 0000000..c298633 --- /dev/null +++ b/.idea/dictionaries/project.xml @@ -0,0 +1,7 @@ + + + + pkce + + + \ No newline at end of file diff --git a/.idea/modules.xml b/.idea/modules.xml new file mode 100644 index 0000000..6404f20 --- /dev/null +++ b/.idea/modules.xml @@ -0,0 +1,8 @@ + + + + + + + + \ No newline at end of file diff --git a/.idea/trpc-mcp-go.iml b/.idea/trpc-mcp-go.iml new file mode 100644 index 0000000..5e764c4 --- /dev/null +++ b/.idea/trpc-mcp-go.iml @@ -0,0 +1,9 @@ + + + + + + + + + \ No newline at end of file diff --git a/.idea/vcs.xml b/.idea/vcs.xml new file mode 100644 index 0000000..35eb1dd --- /dev/null +++ b/.idea/vcs.xml @@ -0,0 +1,6 @@ + + + + + + \ No newline at end of file diff --git a/README.md b/README.md index 08fac5c..f655171 100644 --- a/README.md +++ b/README.md @@ -769,6 +769,187 @@ weatherHandler := mcp.NewTypedToolHandler(func(ctx context.Context, req *mcp.Cal See [`examples/schema-generation/`](examples/schema-generation/) for a complete example. +## OAuth 2.1 Authentication Support +The library provides comprehensive OAuth 2.1 authentication support for securing MCP servers and enabling authenticated client access. + +### Feature + +- **OAuth 2.1 Authorization Code Flow:** Full RFC-compliant authorization code flow with PKCE support +- **JWT Token Handling:** HMAC-signed JWT access and refresh tokens with automatic validation +- **Token Introspection:** RFC 7662 compliant token introspection for real-time validation +- **Metadata Endpoints:** OAuth 2.1 and OpenID Connect discovery endpoints +- **Audit Logging:** Comprehensive security logging with sensitive data protection +- **Bearer Token Protection:** Automatic token verification for MCP endpoints + +### Server Configuration + +#### Basic OAuth Server Setup + +```go +// Create OAuth provider (or use built-in proxy provider) +provider := providers.NewProxyOAuthServerProvider(providers.ProxyOptions{ + Endpoints: providers.ProxyEndpoints{ + AuthorizationURL: "https://your-oauth-server.com/authorize", + TokenURL: "https://your-oauth-server.com/token", + RevocationURL: "https://your-oauth-server.com/revoke", + RegistrationURL: "https://your-oauth-server.com/register", + }, + VerifyAccessToken: yourTokenVerificationFunc, + GetClient: yourClientLookupFunc, +}) + +// Configure token verifier with introspection +verifier, err := server.NewTokenVerifier(ctx, server.TokenVerifierConfig{ + Introspection: &server.IntrospectionConfig{ + Endpoint: "https://your-oauth-server.com/introspect", + Timeout: 5 * time.Second, + CacheTTL: 30 * time.Second, + NegativeCacheTTL: 10 * time.Second, + UseOnJWTFail: true, + }, +}) + +// Create MCP server with OAuth protection +mcpServer := mcp.NewServer( + "Secure-MCP-Server", "1.0.0", + + // OAuth routes for client authorization + mcp.WithOAuthRoutes(mcp.OAuthRoutesConfig{ + Provider: provider, + IssuerURL: mustURL("https://your-oauth-server.com"), + BaseURL: mustURL("https://your-mcp-server.com"), + ScopesSupported: []string{"mcp.read", "mcp.write"}, + }), + + // OAuth metadata endpoints + mcp.WithOAuthMetadata(mcp.OAuthMetadataConfig{ + ResourceServerURL: mustURL("https://your-mcp-server.com"), + ScopesSupported: []string{"mcp.read", "mcp.write"}, + ResourceName: &"MCP Server", + }), + + // Bearer token authentication + mcp.WithBearerAuth(&mcp.BearerAuthConfig{ + Enabled: true, + RequiredScopes: []string{"mcp.read", "mcp.write"}, + Verifier: verifier, + }), + + // Security audit logging + mcp.WithAudit(&mcp.AuditConfig{ + Enabled: true, + Level: "basic", + HashSensitiveData: true, + IncludeRequestBody: false, + IncludeResponseBody: false, + EndpointPatterns: []string{"/mcp/", "/authorize", "/token"}, + }), +) +``` + +### Client Configuration + +#### OAuth-Enabled Client + +```go +// Configure OAuth authentication flow +authFlow := mcp.AuthFlowConfig{ + ServerURL: "https://your-mcp-server.com", + ClientMetadata: auth.OAuthClientMetadata{ + ClientName: &"my-mcp-client", + GrantTypes: []string{"authorization_code", "refresh_token"}, + TokenEndpointAuthMethod: "client_secret_post", + RedirectURIs: []string{"http://localhost:8080/callback"}, + Scope: &"mcp.read mcp.write", + }, + ResourceMetadataURL: &"https://your-mcp-server.com/.well-known/oauth-protected-resource", + RedirectURL: "http://localhost:8080/callback", + Scope: &"mcp.read mcp.write", + OnRedirect: func(u *url.URL) error { + fmt.Printf("Please authorize at: %s\n", u.String()) + // Open browser or handle redirect + return nil + }, +} + +// Create authenticated MCP client +client, err := mcp.NewClient( + "https://your-mcp-server.com/mcp", + mcp.Implementation{Name: "My-Client", Version: "1.0.0"}, + mcp.WithAuthFlow(authFlow), +) +if err != nil { + log.Fatal(err) +} + +// Complete authorization flow +ctx := context.Background() +_, err = client.Initialize(ctx, &mcp.InitializeRequest{}) +if err != nil { + log.Fatal(err) +} + +// Handle authorization callback (typically in HTTP handler) +authCode := "received_from_redirect" +if err := client.CompleteAuthFlow(ctx, authCode); err != nil { + log.Fatal(err) +} + +// Use authenticated client normally +result, err := client.CallTool(ctx, &mcp.CallToolRequest{ + Params: mcp.CallToolParams{Name: "secure_tool"}, +}) +``` + +### Security Features + +#### Token Verification + +The library supports multiple token verification methods: + +- **JWT Verification:** Direct HMAC/RSA signature validation +- **Token Introspection:** RFC 7662 compliant real-time validation +- **Hybrid Mode:** JWT first with introspection fallback +- **Caching:** Configurable positive/negative caching for performance + +#### Audit Logger + +Comprehensive security logging with configurable levels: + +```go +mcp.WithAudit(&mcp.AuditConfig{ + Enabled: true, + Level: "detailed", // "basic" or "detailed" + HashSensitiveData: true, // Hash tokens/secrets in logs + IncludeRequestBody: true, // Log request payloads + IncludeResponseBody: false, // Log response payloads + EndpointPatterns: []string{"/mcp/", "/oauth/"}, + ExcludePatterns: []string{"/health", "/metrics"}, +}) +``` + +### OAuth Endpoints + +When OAuth is enabled, the server automatically provides: + +| Endpoint | Description | +|-------------------------------------------|-------------| +| `/.well-known/oauth-protected-resource` | OAuth 2.1 Resource Server Metadata (RFC 8705) | +| `/.well-known/oauth-authorization-server` | OAuth 2.1 Authorization Server Metadata (RFC 8414) | +| `/authorize` | Authorization endpoint for client redirects | +| `/token` | Token endpoint for code/refresh exchange | +| `/register` | Dynamic client registration (RFC 7591) | +| `/revoke` | Token revocation endpoint (RFC 7009) | + +### Complete Example + +See the OAuth authentication example at [`examples/auth/`](examples/auth) for a complete working implementation including: + +- Mock OAuth 2.1 server setup +- MCP server with OAuth protection +- MCP client with authorization flow +- Token handling and refresh +- Security audit logging ## Example Patterns diff --git a/client.go b/client.go index 0f3d485..641874b 100644 --- a/client.go +++ b/client.go @@ -14,7 +14,8 @@ import ( "reflect" "sync" "sync/atomic" - + "trpc.group/trpc-go/trpc-mcp-go/internal/auth" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth/client" "trpc.group/trpc-go/trpc-mcp-go/internal/errors" "trpc.group/trpc-go/trpc-mcp-go/internal/retry" ) @@ -117,6 +118,15 @@ type Client struct { state State // State. transportOptions []transportOption + // OAuth provider + oauthProvider client.OAuthClientProvider + // OAuth token management + accessToken string + refreshToken string + + // OAuth flow configuration + authFlowConfig *AuthFlowConfig + // transport configuration. transportConfig *transportConfig @@ -159,6 +169,30 @@ func NewClient(serverURL string, clientInfo Implementation, options ...ClientOpt option(client) } + // Handle OAuth authentication if configured + if client.oauthProvider != nil { + if tokens, err := client.oauthProvider.Tokens(); err == nil && tokens != nil { + client.accessToken = tokens.AccessToken + if tokens.RefreshToken != nil { + client.refreshToken = *tokens.RefreshToken + } + if client.accessToken != "" { + // Ensure a headers map exists + if client.transportConfig.httpHeaders == nil { + client.transportConfig.httpHeaders = make(http.Header) + } + client.transportConfig.httpHeaders.Set("Authorization", "Bearer "+client.accessToken) + // Also push into transportOptions so the streamable transport sees it + client.transportOptions = append(client.transportOptions, withTransportHTTPHeaders(client.transportConfig.httpHeaders)) + } + } else if err != nil { + // Optional: Surface a warning via logger if available + if client.logger != nil { + client.logger.Warnf("OAuth provider returned no tokens at client initialization: %v", err) + } + } + } + // Create transport layer if not previously set via options. if client.transport == nil { client.transport = newStreamableHTTPClientTransport(client.transportConfig, client.transportOptions...) @@ -177,6 +211,38 @@ func NewClient(serverURL string, clientInfo Implementation, options ...ClientOpt return client, nil } +// AuthFlowConfig configures OAuth 2.0 authentication flow +type AuthFlowConfig struct { + // ServerURL is the OAuth authorization server URL (required) + ServerURL string + + // ClientMetadata contains OAuth client configuration + ClientMetadata auth.OAuthClientMetadata + + // RedirectURL is the OAuth redirect URI (required) + RedirectURL string + + // OnRedirect handles the authorization redirect (required) + // This function should redirect the user to the authorization URL + OnRedirect func(*url.URL) error + + // Scope defines the requested access permissions (optional) + Scope *string + + // CustomFetchFunc allows custom HTTP client behavior (optional) + CustomFetchFunc auth.FetchFunc + + // ResourceMetadataURL for discovering protected resource metadata (optional) + ResourceMetadataURL *string + + // AuthorizationCode can be provided if you already have one (optional) + // This is useful for handling the redirect callback + AuthorizationCode *string + + // State for CSRF protection (optional) + State *string +} + // transportConfig includes transport layer configuration. type transportConfig struct { serverURL *url.URL // server URL @@ -198,6 +264,8 @@ type transportConfig struct { // These options are typically not used by the default handler, but may be used by custom // implementations that replace the default NewHTTPReqHandler function for extensibility. httpReqHandlerOptions []HTTPReqHandlerOption + + oauthProvider client.OAuthClientProvider } // newDefaultTransportConfig creates a default transport configuration. @@ -210,6 +278,7 @@ func newDefaultTransportConfig() *transportConfig { httpReqHandlerOptions: []HTTPReqHandlerOption{}, enableGetSSE: true, path: "", + oauthProvider: nil, } } @@ -324,6 +393,18 @@ func (c *Client) Initialize(ctx context.Context, initReq *InitializeRequest) (*I return nil, errors.ErrAlreadyInitialized } + // If auth flow is configured, execute it first + if c.authFlowConfig != nil { + if err := c.executeAuthFlow(ctx); err != nil { + return nil, fmt.Errorf("authentication failed: %w", err) + } + + // Ensure transport uses the latest token + if err := c.updateClientTokens(); err != nil { + return nil, fmt.Errorf("failed to update tokens: %w", err) + } + } + // Create request. requestID := c.requestID.Add(1) req := newJSONRPCRequest(requestID, MethodInitialize, map[string]interface{}{ @@ -667,3 +748,362 @@ func (c *Client) SendRootsListChangedNotification(ctx context.Context) error { func isZeroStruct(x interface{}) bool { return reflect.ValueOf(x).IsZero() } + +// WithAuthFlow creates a client option that configures and executes the complete OAuth flow +func WithAuthFlow(config AuthFlowConfig) ClientOption { + return func(c *Client) { + // Validate required configuration + if config.ServerURL == "" { + panic("AuthFlowConfig.ServerURL is required") + } + if config.RedirectURL == "" { + panic("AuthFlowConfig.RedirectURL is required") + } + if config.OnRedirect == nil { + panic("AuthFlowConfig.OnRedirect is required") + } + + // Store config for later use + c.authFlowConfig = &config + + // Create internal OAuth provider + provider := client.NewInMemoryOAuthClientProvider( + config.RedirectURL, + config.ClientMetadata, + config.OnRedirect, + ) + + c.oauthProvider = provider + c.transportConfig.oauthProvider = provider + c.transportOptions = append(c.transportOptions, withTransportOAuthProvider(provider)) + + } +} + +// executeAuthFlow runs the complete OAuth authentication flow using internal methods +func (c *Client) executeAuthFlow(ctx context.Context) error { + if c.authFlowConfig == nil { + return fmt.Errorf("auth flow not configured") + } + + // Build auth options for internal flow + authOptions := auth.AuthOptions{ + ServerUrl: c.authFlowConfig.ServerURL, + Scope: c.authFlowConfig.Scope, + FetchFn: c.authFlowConfig.CustomFetchFunc, + ResourceMetadataUrl: c.authFlowConfig.ResourceMetadataURL, + } + + // Execute the complete internal authentication flow + result, err := c.authInternal(authOptions) + if err != nil { + return fmt.Errorf("OAuth flow failed: %w", err) + } + + // Handle authentication result + return c.handleAuthResult(result) +} + +// authInternal orchestrates the complete OAuth flow using internal methods +func (c *Client) authInternal(options auth.AuthOptions) (*client.AuthResult, error) { + // Discover protected resource metadata (if available) + var resourceMetadata *auth.OAuthProtectedResourceMetadata + var authorizationServerUrl string + + metadata, err := c.discoverProtectedResource(options.ServerUrl, options) + if err == nil { + resourceMetadata = metadata + if len(resourceMetadata.AuthorizationServers) > 0 { + authorizationServerUrl = resourceMetadata.AuthorizationServers[0] + } + } + + if authorizationServerUrl == "" { + authorizationServerUrl = options.ServerUrl + } + + // Select resource URL + resource, err := c.selectResourceURL(options.ServerUrl, resourceMetadata) + if err != nil { + return nil, fmt.Errorf("failed to select resource URL: %w", err) + } + + // Discover authorization server metadata + serverMetadata, err := c.discoverAuthServer(authorizationServerUrl) + if err != nil { + return nil, fmt.Errorf("failed to discover authorization server: %w", err) + } + + // Handle client registration if needed + clientInfo, err := c.handleClientRegistration(authorizationServerUrl, serverMetadata, options) + if err != nil { + return nil, fmt.Errorf("client registration failed: %w", err) + } + + // Try token refresh if refresh token exists + if result, err := c.tryTokenRefresh(authorizationServerUrl, serverMetadata, clientInfo, resource, options); err == nil { + return result, nil + } + + // Exchange the authorization code for a token + if options.AuthorizationCode != nil && *options.AuthorizationCode != "" { + cv, err := c.oauthProvider.CodeVerifier() + if err != nil || cv == "" { + return nil, fmt.Errorf("missing code_verifier: %w", err) + } + + var addClientAuth func(http.Header, url.Values, string) error + if ap, ok := c.oauthProvider.(client.OAuthClientAuthProvider); ok { + addClientAuth = ap.AddClientAuthentication + } + + tokens, err := client.ExchangeAuthorization(authorizationServerUrl, client.ExchangeAuthorizationOptions{ + Metadata: serverMetadata, + ClientInformation: clientInfo, + AuthorizationCode: *options.AuthorizationCode, + CodeVerifier: cv, + RedirectURI: c.oauthProvider.RedirectURL(), + Resource: resource, + AddClientAuthentication: addClientAuth, + FetchFn: options.FetchFn, + }) + if err != nil { + return nil, err + } + if err := c.oauthProvider.SaveTokens(*tokens); err != nil { + return nil, fmt.Errorf("failed to save tokens: %w", err) + } + res := client.AuthResultAuthorized + return &res, nil + } + + // Start authorization flow + return c.startAuthorizationFlow(authorizationServerUrl, serverMetadata, clientInfo, resource, options) +} + +// discoverProtectedResource discovers OAuth protected resource metadata +func (c *Client) discoverProtectedResource(serverUrl string, options auth.AuthOptions) (*auth.OAuthProtectedResourceMetadata, error) { + discoveryOptions := &auth.DiscoveryOptions{ + ResourceMetadataUrl: options.ResourceMetadataUrl, + } + + return client.DiscoverOAuthProtectedResourceMetadata(serverUrl, discoveryOptions, options.FetchFn) +} + +// discoverAuthServer discovers authorization server metadata +func (c *Client) discoverAuthServer(authServerUrl string) (auth.AuthorizationServerMetadata, error) { + return client.DiscoverAuthorizationServerMetadata(context.Background(), authServerUrl, nil) +} + +// handleClientRegistration handles dynamic client registration if needed +func (c *Client) handleClientRegistration(authServerUrl string, serverMetadata auth.AuthorizationServerMetadata, options auth.AuthOptions) (*auth.OAuthClientInformation, error) { + clientInfo := c.oauthProvider.ClientInformation() + + if clientInfo == nil { + // Need to register client + if _, ok := c.oauthProvider.(client.OAuthClientInfoProvider); !ok { + return nil, fmt.Errorf("OAuth client information must be saveable for dynamic registration") + } + + fullInfo, err := client.RegisterClient(context.Background(), authServerUrl, client.RegisterClientOptions{ + Metadata: serverMetadata, + ClientMetadata: c.oauthProvider.ClientMetadata(), + FetchFn: options.FetchFn, + }) + if err != nil { + return nil, fmt.Errorf("failed to register client: %w", err) + } + + if clientInfoProvider, ok := c.oauthProvider.(client.OAuthClientInfoProvider); ok { + if err := clientInfoProvider.SaveClientInformation(*fullInfo); err != nil { + return nil, fmt.Errorf("failed to save client information: %w", err) + } + } + + clientInfo = &auth.OAuthClientInformation{ + ClientID: fullInfo.ClientID, + ClientSecret: fullInfo.ClientSecret, + } + } + + return clientInfo, nil +} + +// selectResourceURL selects the appropriate resource URL +func (c *Client) selectResourceURL(serverUrl string, resourceMetadata *auth.OAuthProtectedResourceMetadata) (*url.URL, error) { + defaultResource, err := auth.ResourceURLFromServerURL(serverUrl) + if err != nil { + return nil, err + } + + // Use custom validator if available + if validator, ok := c.oauthProvider.(client.OAuthResourceValidator); ok { + return validator.ValidateResourceURL(defaultResource, resourceMetadata) + } + + // Include resource param only when metadata exists + if resourceMetadata == nil { + return nil, nil + } + + // Check metadata resource compatibility + allowed, err := auth.CheckResourceAllowed(auth.CheckResourceAllowedParams{ + RequestedResource: defaultResource, + ConfiguredResource: resourceMetadata.Resource, + }) + if err != nil { + return nil, fmt.Errorf("failed to validate resource: %w", err) + } + if !allowed { + return nil, fmt.Errorf("protected resource mismatch") + } + + return url.Parse(resourceMetadata.Resource) +} + +// tryTokenRefresh attempts to refresh existing tokens +func (c *Client) tryTokenRefresh(authServerUrl string, serverMetadata auth.AuthorizationServerMetadata, clientInfo *auth.OAuthClientInformation, resource *url.URL, options auth.AuthOptions) (*client.AuthResult, error) { + tokens, err := c.oauthProvider.Tokens() + if err != nil || tokens == nil || tokens.RefreshToken == nil || *tokens.RefreshToken == "" { + return nil, fmt.Errorf("no refresh token available") + } + + var addClientAuth func(http.Header, url.Values, string) error + if authProvider, ok := c.oauthProvider.(client.OAuthClientAuthProvider); ok { + addClientAuth = authProvider.AddClientAuthentication + } + + newTokens, err := client.RefreshAuthorization(authServerUrl, client.RefreshAuthorizationOptions{ + Metadata: serverMetadata, + ClientInformation: clientInfo, + RefreshToken: *tokens.RefreshToken, + Resource: resource, + AddClientAuthentication: addClientAuth, + FetchFn: options.FetchFn, + }) + if err != nil { + return nil, err + } + + if err := c.oauthProvider.SaveTokens(*newTokens); err != nil { + return nil, fmt.Errorf("failed to save refreshed tokens: %w", err) + } + + result := client.AuthResultAuthorized + return &result, nil +} + +// startAuthorizationFlow starts the authorization code flow +func (c *Client) startAuthorizationFlow(authServerUrl string, serverMetadata auth.AuthorizationServerMetadata, clientInfo *auth.OAuthClientInformation, resource *url.URL, options auth.AuthOptions) (*client.AuthResult, error) { + var state *string + if c.authFlowConfig.State != nil { + state = c.authFlowConfig.State + } else if stateProvider, ok := c.oauthProvider.(client.OAuthStateProvider); ok { + stateValue, err := stateProvider.State() + if err != nil { + return nil, fmt.Errorf("failed to get state: %w", err) + } + state = &stateValue + } + + scope := options.Scope + if scope == nil { + clientMetadata := c.oauthProvider.ClientMetadata() + if clientMetadata.Scope != nil { + scope = clientMetadata.Scope + } + } + + authResult, err := client.StartAuthorization(authServerUrl, client.StartAuthorizationOptions{ + Metadata: serverMetadata, + ClientInformation: *clientInfo, + State: state, + RedirectURL: c.oauthProvider.RedirectURL(), + Scope: scope, + Resource: resource, + }) + if err != nil { + return nil, fmt.Errorf("failed to start authorization: %w", err) + } + + if err := c.oauthProvider.SaveCodeVerifier(authResult.CodeVerifier); err != nil { + return nil, fmt.Errorf("failed to save code verifier: %w", err) + } + + if err := c.oauthProvider.RedirectToAuthorization(authResult.AuthorizationURL); err != nil { + return nil, fmt.Errorf("failed to redirect to authorization: %w", err) + } + + result := client.AuthResultRedirect + return &result, nil +} + +// handleAuthResult processes the authentication result +func (c *Client) handleAuthResult(result *client.AuthResult) error { + switch *result { + case client.AuthResultAuthorized: + return c.updateClientTokens() + case client.AuthResultRedirect: + // User needs to complete authorization, this is normal + return nil + default: + return fmt.Errorf("unknown authentication result: %s", *result) + } +} + +// updateClientTokens updates HTTP headers with new tokens +func (c *Client) updateClientTokens() error { + tokens, err := c.oauthProvider.Tokens() + if err != nil { + return fmt.Errorf("failed to get tokens: %w", err) + } + + if tokens != nil && tokens.AccessToken != "" { + c.accessToken = tokens.AccessToken + if tokens.RefreshToken != nil { + c.refreshToken = *tokens.RefreshToken + } + + // Update HTTP headers + if c.transportConfig.httpHeaders == nil { + c.transportConfig.httpHeaders = make(http.Header) + } + c.transportConfig.httpHeaders.Set("Authorization", "Bearer "+c.accessToken) + + // Update existing transport + if c.transport != nil { + if streamableTransport, ok := c.transport.(*streamableHTTPClientTransport); ok { + if streamableTransport.httpHeaders == nil { + streamableTransport.httpHeaders = make(http.Header) + } + streamableTransport.httpHeaders.Set("Authorization", "Bearer "+c.accessToken) + } + } + } + + return nil +} + +// CompleteAuthFlow completes the OAuth flow with authorization code +func (c *Client) CompleteAuthFlow(ctx context.Context, authorizationCode string) error { + if c.authFlowConfig == nil { + return fmt.Errorf("auth flow not configured") + } + + // Set the authorization code in options + authOptions := auth.AuthOptions{ + ServerUrl: c.authFlowConfig.ServerURL, + Scope: c.authFlowConfig.Scope, + FetchFn: c.authFlowConfig.CustomFetchFunc, + ResourceMetadataUrl: c.authFlowConfig.ResourceMetadataURL, + AuthorizationCode: &authorizationCode, + } + + // Execute the flow with the authorization code + result, err := c.authInternal(authOptions) + if err != nil { + return fmt.Errorf("failed to complete auth flow: %w", err) + } + + return c.handleAuthResult(result) +} diff --git a/e2e/oauth_integration_test.go b/e2e/oauth_integration_test.go new file mode 100644 index 0000000..08aca07 --- /dev/null +++ b/e2e/oauth_integration_test.go @@ -0,0 +1,693 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package e2e + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "testing" + "time" + + "github.com/golang-jwt/jwt/v4" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + mcp "trpc.group/trpc-go/trpc-mcp-go" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth/server" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth/server/providers" +) + +const ( + testHMACSecret = "test-oauth-secret" + testClientID = "test-client" + testScope = "mcp.read mcp.write" +) + +// TestOAuth2Integration tests the complete OAuth 2.1 flow with MCP server +func TestOAuth2Integration(t *testing.T) { + // Start mock OAuth authorization server + oauthServer := startMockOAuthServer(t) + defer oauthServer.Close() + + // Create OAuth Provider + provider := createTestOAuthProvider(oauthServer.URL) + + // Start MCP server with OAuth authentication + mcpServerURL, cleanup := startOAuthMCPServer(t, provider) + defer cleanup() + + // Test OAuth flows + t.Run("BearerTokenAuth", func(t *testing.T) { + testBearerTokenAuth(t, oauthServer.URL, mcpServerURL) + }) + + t.Run("InvalidToken", func(t *testing.T) { + testInvalidToken(t, mcpServerURL) + }) + + // Test authorization code flow + t.Run("AuthorizationCodeFlow", func(t *testing.T) { + testSimpleAuthorizationCodeFlow(t, oauthServer.URL, mcpServerURL) + }) + + t.Run("TokenRefresh", func(t *testing.T) { + testTokenRefresh(t, oauthServer.URL, mcpServerURL) + }) +} + +// startMockOAuthServer starts a mock OAuth authorization server for testing +func startMockOAuthServer(t *testing.T) *httptest.Server { + t.Helper() + + mux := http.NewServeMux() + + // Authorization endpoint + mux.HandleFunc("/authorize", func(w http.ResponseWriter, r *http.Request) { + if r.Method != "GET" { + http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) + return + } + + // Simulate user authorization by redirecting directly to callback URL + redirectURI := r.URL.Query().Get("redirect_uri") + state := r.URL.Query().Get("state") + code := "test-auth-code-" + fmt.Sprintf("%d", time.Now().Unix()) + + callbackURL := fmt.Sprintf("%s?code=%s&state=%s", redirectURI, code, state) + http.Redirect(w, r, callbackURL, http.StatusFound) + }) + + // Token endpoint + mux.HandleFunc("/token", func(w http.ResponseWriter, r *http.Request) { + if r.Method != "POST" { + http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) + return + } + + if err := r.ParseForm(); err != nil { + http.Error(w, "Invalid form", http.StatusBadRequest) + return + } + + grantType := r.FormValue("grant_type") + code := r.FormValue("code") + refreshToken := r.FormValue("refresh_token") + + var tokenResponse map[string]interface{} + + switch grantType { + case "authorization_code": + if code == "" { + http.Error(w, "Missing authorization code", http.StatusBadRequest) + return + } + tokenResponse = createTokenResponse(t, "access_token", "refresh_token") + + case "refresh_token": + if refreshToken == "" { + http.Error(w, "Missing refresh token", http.StatusBadRequest) + return + } + // Validate refresh token + if !strings.HasPrefix(refreshToken, "test-refresh-token-") { + http.Error(w, "Invalid refresh token", http.StatusBadRequest) + return + } + tokenResponse = createTokenResponse(t, "new_access_token", "new_refresh_token") + + default: + http.Error(w, "Unsupported grant type", http.StatusBadRequest) + return + } + + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(tokenResponse) + }) + + // Client registration endpoint + mux.HandleFunc("/register", func(w http.ResponseWriter, r *http.Request) { + if r.Method != "POST" { + http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) + return + } + + clientInfo := map[string]interface{}{ + "client_id": testClientID, + "client_secret": "", + "redirect_uris": []string{"http://localhost:5173/callback"}, + "grant_types": []string{"authorization_code", "refresh_token"}, + "scope": testScope, + } + + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(clientInfo) + }) + + // OAuth authorization server metadata endpoint + mux.HandleFunc("/.well-known/oauth-authorization-server", func(w http.ResponseWriter, r *http.Request) { + if r.Method != "GET" { + http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) + return + } + + // Dynamically get server URL + baseURL := "http://" + r.Host + metadata := map[string]interface{}{ + "issuer": baseURL, + "authorization_endpoint": baseURL + "/authorize", + "token_endpoint": baseURL + "/token", + "registration_endpoint": baseURL + "/register", + "response_types_supported": []string{"code"}, + "grant_types_supported": []string{"authorization_code", "refresh_token"}, + "code_challenge_methods_supported": []string{"S256"}, + "token_endpoint_auth_methods_supported": []string{"client_secret_post", "client_secret_basic"}, + "scopes_supported": []string{"mcp.read", "mcp.write"}, + } + + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(metadata) + }) + + // OpenID Connect configuration endpoint (for compatibility) + mux.HandleFunc("/.well-known/openid-configuration", func(w http.ResponseWriter, r *http.Request) { + if r.Method != "GET" { + http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) + return + } + + // Dynamically get server URL + baseURL := "http://" + r.Host + // Return same content as OAuth metadata + metadata := map[string]interface{}{ + "issuer": baseURL, + "authorization_endpoint": baseURL + "/authorize", + "token_endpoint": baseURL + "/token", + "registration_endpoint": baseURL + "/register", + "response_types_supported": []string{"code"}, + "grant_types_supported": []string{"authorization_code", "refresh_token"}, + "code_challenge_methods_supported": []string{"S256"}, + "token_endpoint_auth_methods_supported": []string{"client_secret_post", "client_secret_basic"}, + "scopes_supported": []string{"mcp.read", "mcp.write"}, + } + + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(metadata) + }) + + server := httptest.NewServer(mux) + t.Logf("Mock OAuth server started at: %s", server.URL) + return server +} + +// createTestOAuthProvider creates an OAuth provider for testing +func createTestOAuthProvider(oauthServerURL string) server.OAuthServerProvider { + return providers.NewProxyOAuthServerProvider(providers.ProxyOptions{ + Endpoints: providers.ProxyEndpoints{ + AuthorizationURL: oauthServerURL + "/authorize", + TokenURL: oauthServerURL + "/token", + RegistrationURL: oauthServerURL + "/register", + }, + VerifyAccessToken: func(token string) (*server.AuthInfo, error) { + return verifyTestJWT(token) + }, + GetClient: func(clientID string) (*auth.OAuthClientInformationFull, error) { + return &auth.OAuthClientInformationFull{ + OAuthClientMetadata: auth.OAuthClientMetadata{ + RedirectURIs: []string{"http://localhost:5173/callback"}, + ResponseTypes: []string{"code"}, + GrantTypes: []string{"authorization_code", "refresh_token"}, + ClientName: stringPtr("test-client"), + Scope: stringPtr(testScope), + }, + OAuthClientInformation: auth.OAuthClientInformation{ + ClientID: clientID, + ClientSecret: "", + }, + }, nil + }, + }) +} + +// startOAuthMCPServer starts an MCP server with OAuth authentication enabled +func startOAuthMCPServer(t *testing.T, provider server.OAuthServerProvider) (string, func()) { + t.Helper() + + // Create MCP server using standardized approach + server := mcp.NewServer( + "OAuth-Test-Server", + "1.0.0", + mcp.WithServerPath("/mcp"), + mcp.WithOAuthRoutes(mcp.OAuthRoutesConfig{ + Provider: provider, + IssuerURL: mustParseURL("http://localhost:3030"), + BaseURL: mustParseURL("http://localhost:3000"), + ScopesSupported: []string{"mcp.read", "mcp.write"}, + }), + mcp.WithBearerAuth(&mcp.BearerAuthConfig{ + Enabled: true, + RequiredScopes: []string{"mcp.read", "mcp.write"}, + Issuer: "http://localhost:3030", + Audience: []string{"http://localhost:3000"}, + Verifier: server.TokenVerifierFunc(func(ctx context.Context, token string) (server.AuthInfo, error) { + authInfo, err := verifyTestJWT(token) + if err != nil { + return server.AuthInfo{}, err + } + return *authInfo, nil + }), + }), + ) + + // Register test tools using standardized approach + RegisterTestTools(server) + + // Create HTTP test server + httpServer := httptest.NewServer(server.HTTPHandler()) + serverURL := httpServer.URL + "/mcp" + + t.Logf("OAuth MCP server started at: %s", serverURL) + + cleanup := func() { + t.Log("Closing OAuth MCP server") + httpServer.Close() + } + + return serverURL, cleanup +} + +// testBearerTokenAuth tests Bearer Token authentication using standardized approach +func testBearerTokenAuth(t *testing.T, oauthServerURL, mcpServerURL string) { + t.Helper() + + // Create client directly with a valid JWT Token + validToken := createTestJWT(t, "access_token") + + // Create HTTP headers with Bearer Token + headers := make(http.Header) + headers.Set("Authorization", "Bearer "+validToken) + + // Use standardized client creation + client := CreateTestClient(t, mcpServerURL, func(c *mcp.Client) { + // Apply OAuth headers + mcp.WithHTTPHeaders(headers)(c) + }) + defer CleanupClient(t, client) + + // Initialize client using standardized approach + InitializeClient(t, client) + + // Test tool invocation using standardized approach + content := ExecuteTestTool(t, client, "basic-greet", map[string]interface{}{ + "name": "bearer-test", + }) + + require.Len(t, content, 1) + textContent, ok := content[0].(mcp.TextContent) + assert.True(t, ok) + assert.Contains(t, textContent.Text, "Hello, bearer-test") +} + +// testInvalidToken tests invalid token handling using standardized approach +func testInvalidToken(t *testing.T, mcpServerURL string) { + t.Helper() + + // Create client with invalid token + invalidToken := "invalid.jwt.token" + + // Create HTTP headers with invalid Bearer Token + headers := make(http.Header) + headers.Set("Authorization", "Bearer "+invalidToken) + + // Use standardized client creation + client := CreateTestClient(t, mcpServerURL, func(c *mcp.Client) { + // Apply OAuth headers + mcp.WithHTTPHeaders(headers)(c) + }) + defer CleanupClient(t, client) + + // Try to initialize client, should fail + ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout) + defer cancel() + + _, err := client.Initialize(ctx, &mcp.InitializeRequest{ + Params: mcp.InitializeParams{ + ProtocolVersion: mcp.ProtocolVersion_2025_03_26, + ClientInfo: mcp.Implementation{ + Name: "Invalid-Token-Client", + Version: "1.0.0", + }, + }, + }) + + // Should return authentication error + assert.Error(t, err) + // Check if it contains authentication-related error message + errorMsg := err.Error() + assert.True(t, + strings.Contains(errorMsg, "unauthorized") || + strings.Contains(errorMsg, "401") || + strings.Contains(errorMsg, "authentication") || + strings.Contains(errorMsg, "auth"), + "Expected authentication error, got: %s", errorMsg) +} + +// testSimpleAuthorizationCodeFlow tests the simplified authorization code flow +func testSimpleAuthorizationCodeFlow(t *testing.T, oauthServerURL, mcpServerURL string) { + t.Helper() + + // Test authorization endpoint functionality + t.Run("AuthorizationEndpoint", func(t *testing.T) { + // Start a simple callback server + callbackServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + t.Logf("Callback received: %s", r.URL.RawQuery) + w.WriteHeader(http.StatusOK) + w.Write([]byte("OK")) + })) + defer callbackServer.Close() + + // Build authorization URL with URL encoding + params := url.Values{} + params.Set("client_id", testClientID) + params.Set("response_type", "code") + params.Set("redirect_uri", callbackServer.URL+"/callback") + params.Set("scope", testScope) + params.Set("state", "test-state") + + authURL := oauthServerURL + "/authorize?" + params.Encode() + t.Logf("Testing authorization URL: %s", authURL) + + // Create HTTP client that doesn't follow redirects + client := &http.Client{ + CheckRedirect: func(req *http.Request, via []*http.Request) error { + return http.ErrUseLastResponse // Don't follow redirects + }, + } + + // Access authorization endpoint + resp, err := client.Get(authURL) + require.NoError(t, err) + defer resp.Body.Close() + + // Should redirect to callback URL + assert.Equal(t, http.StatusFound, resp.StatusCode) + + // Check redirect URL + location := resp.Header.Get("Location") + t.Logf("Redirect location: %s", location) + assert.Contains(t, location, "code=") + assert.Contains(t, location, "state=test-state") + }) + + // Test token endpoint functionality + t.Run("TokenEndpoint", func(t *testing.T) { + // Simulate authorization code exchange for tokens + formData := url.Values{} + formData.Set("grant_type", "authorization_code") + formData.Set("code", "test-auth-code-123") + formData.Set("redirect_uri", "http://localhost:5173/callback") + + resp, err := http.PostForm(oauthServerURL+"/token", formData) + require.NoError(t, err) + defer resp.Body.Close() + + // Should return success + assert.Equal(t, http.StatusOK, resp.StatusCode) + + // Validate response content + var tokenResp map[string]interface{} + err = json.NewDecoder(resp.Body).Decode(&tokenResp) + require.NoError(t, err) + + assert.Contains(t, tokenResp, "access_token") + assert.Contains(t, tokenResp, "refresh_token") + assert.Equal(t, "Bearer", tokenResp["token_type"]) + assert.Equal(t, testScope, tokenResp["scope"]) + }) + + // Test OAuth metadata endpoint + t.Run("OAuthMetadata", func(t *testing.T) { + // Test OAuth authorization server metadata + resp, err := http.Get(oauthServerURL + "/.well-known/oauth-authorization-server") + require.NoError(t, err) + defer resp.Body.Close() + + assert.Equal(t, http.StatusOK, resp.StatusCode) + + var metadata map[string]interface{} + err = json.NewDecoder(resp.Body).Decode(&metadata) + require.NoError(t, err) + + assert.Equal(t, oauthServerURL, metadata["issuer"]) + assert.Contains(t, metadata, "authorization_endpoint") + assert.Contains(t, metadata, "token_endpoint") + assert.Contains(t, metadata, "scopes_supported") + }) +} + +// testTokenRefresh tests token refresh functionality +func testTokenRefresh(t *testing.T, oauthServerURL, mcpServerURL string) { + t.Helper() + + // Create a client with a refresh token + refreshToken := "test-refresh-token-" + fmt.Sprintf("%d", time.Now().Unix()) + + // Create OAuth Provider that supports token refresh + provider := providers.NewProxyOAuthServerProvider(providers.ProxyOptions{ + Endpoints: providers.ProxyEndpoints{ + AuthorizationURL: oauthServerURL + "/authorize", + TokenURL: oauthServerURL + "/token", + RegistrationURL: oauthServerURL + "/register", + }, + VerifyAccessToken: func(token string) (*server.AuthInfo, error) { + return verifyTestJWT(token) + }, + GetClient: func(clientID string) (*auth.OAuthClientInformationFull, error) { + return &auth.OAuthClientInformationFull{ + OAuthClientMetadata: auth.OAuthClientMetadata{ + RedirectURIs: []string{"http://localhost:5173/callback"}, + ResponseTypes: []string{"code"}, + GrantTypes: []string{"authorization_code", "refresh_token"}, + ClientName: stringPtr("test-client"), + Scope: stringPtr(testScope), + }, + OAuthClientInformation: auth.OAuthClientInformation{ + ClientID: clientID, + ClientSecret: "", + }, + }, nil + }, + }) + + // Start MCP server with OAuth + mcpServerURL, cleanup := startOAuthMCPServer(t, provider) + defer cleanup() + + // Test token refresh flow + t.Run("RefreshTokenFlow", func(t *testing.T) { + // Simulate refresh token request + refreshReq := map[string]string{ + "grant_type": "refresh_token", + "refresh_token": refreshToken, + } + + // Send refresh request to OAuth server + formData := url.Values{} + for k, v := range refreshReq { + formData.Set(k, v) + } + resp, err := http.PostForm(oauthServerURL+"/token", formData) + require.NoError(t, err) + defer resp.Body.Close() + + // Validate response + assert.Equal(t, http.StatusOK, resp.StatusCode) + + var tokenResp map[string]interface{} + err = json.NewDecoder(resp.Body).Decode(&tokenResp) + require.NoError(t, err) + + // Validate returned tokens + assert.Contains(t, tokenResp, "access_token") + assert.Contains(t, tokenResp, "refresh_token") + assert.Equal(t, "Bearer", tokenResp["token_type"]) + assert.Equal(t, testScope, tokenResp["scope"]) + + // Validate new access token is valid + newAccessToken, ok := tokenResp["access_token"].(string) + require.True(t, ok) + + // Create client with new access token using standardized approach + headers := make(http.Header) + headers.Set("Authorization", "Bearer "+newAccessToken) + + client := CreateTestClient(t, mcpServerURL, func(c *mcp.Client) { + mcp.WithHTTPHeaders(headers)(c) + }) + defer CleanupClient(t, client) + + // Initialize client using standardized approach + InitializeClient(t, client) + + // Test tool invocation using standardized approach + content := ExecuteTestTool(t, client, "basic-greet", map[string]interface{}{ + "name": "refresh-test", + }) + + require.Len(t, content, 1) + textContent, ok := content[0].(mcp.TextContent) + assert.True(t, ok) + assert.Contains(t, textContent.Text, "Hello, refresh-test") + }) + + // Test invalid refresh token + t.Run("InvalidRefreshToken", func(t *testing.T) { + invalidRefreshReq := map[string]string{ + "grant_type": "refresh_token", + "refresh_token": "invalid-refresh-token", + } + + formData := url.Values{} + for k, v := range invalidRefreshReq { + formData.Set(k, v) + } + resp, err := http.PostForm(oauthServerURL+"/token", formData) + require.NoError(t, err) + defer resp.Body.Close() + + // Should return error (400 Bad Request) + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) + }) +} + +// createTokenResponse creates a test token response with JWT access token +func createTokenResponse(t *testing.T, accessToken, refreshToken string) map[string]interface{} { + t.Helper() + + return map[string]interface{}{ + "access_token": createTestJWT(t, accessToken), + "refresh_token": refreshToken, + "token_type": "Bearer", + "expires_in": 3600, + "scope": testScope, + } +} + +// createTestJWT creates a test JWT token with specified token type +func createTestJWT(t *testing.T, tokenType string) string { + t.Helper() + + claims := jwt.MapClaims{ + "iss": "http://localhost:3030", + "aud": []string{"http://localhost:3000"}, + "sub": testClientID, + "scope": "mcp.read mcp.write", // Ensure correct scope + "iat": time.Now().Unix(), + "exp": time.Now().Add(time.Hour).Unix(), + "token_type": tokenType, + } + + token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims) + signedToken, err := token.SignedString([]byte(testHMACSecret)) + require.NoError(t, err) + return signedToken +} + +// verifyTestJWT verifies and parses a test JWT token +func verifyTestJWT(tokenString string) (*server.AuthInfo, error) { + token, err := jwt.Parse(tokenString, func(token *jwt.Token) (interface{}, error) { + if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok { + return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"]) + } + return []byte(testHMACSecret), nil + }) + + if err != nil { + return nil, err + } + + if claims, ok := token.Claims.(jwt.MapClaims); ok && token.Valid { + // Scopes + scopes := []string{} + if scope, ok := claims["scope"].(string); ok { + scopes = strings.Fields(scope) + } + + // ExpiresAt + var expiresAtPtr *int64 + if v, ok := claims["exp"].(float64); ok { + vv := int64(v) + expiresAtPtr = &vv + } + + // Audience -> Resource (first value) + var resourceURL *url.URL + if audVal, ok := claims["aud"]; ok { + switch v := audVal.(type) { + case string: + if u, err := url.Parse(v); err == nil { + resourceURL = u + } + case []interface{}: + if len(v) > 0 { + if s, ok := v[0].(string); ok { + if u, err := url.Parse(s); err == nil { + resourceURL = u + } + } + } + case []string: + if len(v) > 0 { + if u, err := url.Parse(v[0]); err == nil { + resourceURL = u + } + } + } + } + + // ClientID + clientID, _ := claims["sub"].(string) + + // Extra (include iss and client_id) + extra := map[string]interface{}{} + if iss, ok := claims["iss"].(string); ok { + extra["iss"] = iss + } + if clientID != "" { + extra["client_id"] = clientID + } + + return &server.AuthInfo{ + ClientID: clientID, + Scopes: scopes, + ExpiresAt: expiresAtPtr, + Resource: resourceURL, + Extra: extra, + }, nil + } + + return nil, fmt.Errorf("invalid token") +} + +// stringPtr returns a pointer to the given string +func stringPtr(s string) *string { + return &s +} + +// mustParseURL parses a URL string and panics if parsing fails +func mustParseURL(s string) *url.URL { + u, err := url.Parse(s) + if err != nil { + panic(err) + } + return u +} diff --git a/examples/auth/README.md b/examples/auth/README.md new file mode 100644 index 0000000..00e95ee --- /dev/null +++ b/examples/auth/README.md @@ -0,0 +1,48 @@ +# OAuth 2.1 Authentication Example + +This example demonstrates how to use **trpc-mcp-go** to implement an OAuth 2.1 authentication flow between a client and a server. + +## Features + +- **OAuth 2.1 Authorization Code Flow** + - Client registration and redirect handling + - Token exchange (authorization code & refresh token) + - Token verification with HMAC-signed JWTs + +- **MCP Integration** + - MCP server with OAuth-protected routes + - MCP client with integrated authorization flow + - Example of authenticated tool invocation + +- **Infrastructure** + - Mock OAuth 2.1 server for local testing + - Support for access & refresh tokens + - Token introspection and metadata endpoints + +## Quick Start + +### 1. Start the OAuth Authentication Server +```bash +cd server +go run main.go +``` + +- Runs a mock OAuth 2.1 server on `http://localhost:3030` +- Starts an MCP server with OAuth-protected endpoints on `http://localhost:3000/mcp` + +### 2. Start the OAuth Client +```bash +cd client +go run main.go +``` +- Launches an MCP client with OAuth support +- Opens a browser redirect for user authorization +- Completes the authorization code exchange automatically + +## What it demonstrates + +1. **Client Authorization Flow**: How an MCP client performs OAuth 2.1 authorization using redirect URIs. +2. **Server Protection**: How an MCP server integrates OAuth 2.1 for protecting tools and resources. +3. **JWT Token Handling**: Issuing and verifying HMAC-signed JWT access and refresh tokens. +4. **Metadata & Introspection**: Provides `.well-known` OAuth server metadata and introspection endpoints for compatibility. +5. **Audit Logging**: Examples of secure server-side logging with sensitive data hashing and reduced verbosity. \ No newline at end of file diff --git a/examples/auth/client/main.go b/examples/auth/client/main.go new file mode 100644 index 0000000..9a610a7 --- /dev/null +++ b/examples/auth/client/main.go @@ -0,0 +1,191 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package main + +import ( + "context" + "fmt" + "log" + "net/http" + "net/url" + "time" + + mcp "trpc.group/trpc-go/trpc-mcp-go" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth" +) + +const ( + // Base origin of the MCP resource server + serverURL = "http://localhost:3000" + + // Well-known OAuth protected resource metadata endpoint + resourceMetadataURL = "http://localhost:3000/.well-known/oauth-protected-resource" + + // Local redirect URI that receives the authorization code + redirectURL = "http://localhost:5173/callback" + + // Requested scopes for this demo + scope = "mcp.read mcp.write" + + // HTTP listen address for the local callback server + callbackListenAddr = ":5173" + + // MCP entry endpoint used by the SDK client + mcpEndpoint = "http://localhost:3000/mcp/" +) + +func main() { + log.Println("🖥️ Starting OAuth Client Demo") + log.Println(" Target Server: http://localhost:3000") + log.Println(" Callback URL: http://localhost:5173/callback") + log.Println(" Required Scopes: mcp.read mcp.write") + log.Println() + + // Configure the auth flow used by the MCP SDK + authFlow := mcp.AuthFlowConfig{ + ServerURL: serverURL, + ClientMetadata: auth.OAuthClientMetadata{ + ClientName: strPtr("demo-client"), + GrantTypes: []string{"authorization_code", "refresh_token"}, + TokenEndpointAuthMethod: "client_secret_post", + RedirectURIs: []string{redirectURL}, + Scope: strPtr(scope), + }, + ResourceMetadataURL: strPtr(resourceMetadataURL), + RedirectURL: redirectURL, + Scope: strPtr(scope), + OnRedirect: func(u *url.URL) error { + log.Printf("🌐 Authorization Required\n") + log.Printf(" Please open this URL in your browser:\n") + log.Printf(" %s\n\n", u.String()) + log.Printf(" Waiting for authorization...\n") + return nil + }, + } + + // Create the MCP client with auth flow enabled + client, err := mcp.NewClient( + mcpEndpoint, + mcp.Implementation{Name: "Auth-Example-Client", Version: "0.1.0"}, + mcp.WithAuthFlow(authFlow), + ) + if err != nil { + log.Printf("❌ Failed to create client: %v\n", err) + return + } + + // Start the local HTTP callback server to capture the authorization code + authDone := make(chan struct{}, 1) + cbServer := startCallbackServer(client, authDone) + defer shutdownServer(cbServer) + + // First initialize will typically request user authorization + log.Println("🔄 Step 1: Initializing client (triggering OAuth flow)...") + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + _, _ = client.Initialize(ctx, &mcp.InitializeRequest{}) + + // Wait for the browser redirect to complete the code exchange + select { + case <-authDone: + log.Println("✅ Step 2: Authorization flow completed successfully") + case <-time.After(3 * time.Minute): + log.Println("❌ Authorization timeout after 3 minutes") + return + } + + // Small delay to ensure token persistence + time.Sleep(2 * time.Second) + + // Second initialize should succeed using the stored tokens + log.Println("🔄 Step 3: Testing authenticated connection...") + ctx2, cancel2 := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel2() + + initResp, err := client.Initialize(ctx2, &mcp.InitializeRequest{}) + if err != nil { + log.Printf("❌ Authenticated connection failed: %v\n", err) + return + } + + log.Println("📋 OAuth Flow Summary:") + log.Println(" 1. ✅ Client registration") + log.Println(" 2. ✅ User authorization") + log.Println(" 3. ✅ Token exchange") + log.Println(" 4. ✅ Authenticated API access") + log.Println() + log.Printf("🎉 Success! Connected to MCP Server\n") + log.Printf(" Server: %s v%s\n", initResp.ServerInfo.Name, initResp.ServerInfo.Version) + log.Printf(" Authentication: OAuth 2.0 with Bearer Token\n") +} + +// startCallbackServer runs an HTTP server that handles /callback and completes the OAuth flow via the SDK +func startCallbackServer(c *mcp.Client, done chan<- struct{}) *http.Server { + mux := http.NewServeMux() + mux.HandleFunc("/callback", func(w http.ResponseWriter, r *http.Request) { + code := r.URL.Query().Get("code") + state := r.URL.Query().Get("state") + + log.Printf("🔄 Callback received\n") + log.Printf(" Authorization Code: %s\n", code[:20]+"...") + if state != "" { + log.Printf(" State: %s\n", state[:20]+"...") + } + + if code == "" { + log.Println("❌ Missing authorization code") + http.Error(w, "missing code parameter", http.StatusBadRequest) + return + } + + log.Printf("🎫 Exchanging authorization code for tokens...\n") + + ctx, cancel := context.WithTimeout(r.Context(), 30*time.Second) + defer cancel() + + if err := c.CompleteAuthFlow(ctx, code); err != nil { + log.Printf("❌ Token exchange failed: %v\n", err) + http.Error(w, fmt.Sprintf("Authorization failed: %v", err), http.StatusBadRequest) + return + } + + log.Println("✅ Token exchange successful") + + // Send a nice response page + w.Header().Set("Content-Type", "text/html") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`Authorization Complete!`)) + + // Notify the main goroutine + select { + case done <- struct{}{}: + default: + } + }) + + srv := &http.Server{ + Addr: callbackListenAddr, + Handler: mux, + } + go func() { + log.Printf("🌐 Callback server listening on %s\n", callbackListenAddr) + srv.ListenAndServe() + }() + return srv +} + +// shutdownServer gracefully stops the HTTP server within a short timeout +func shutdownServer(srv *http.Server) { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + _ = srv.Shutdown(ctx) +} + +// strPtr returns a pointer to s +func strPtr(s string) *string { + return &s +} diff --git a/examples/auth/server/main.go b/examples/auth/server/main.go new file mode 100644 index 0000000..d0c21aa --- /dev/null +++ b/examples/auth/server/main.go @@ -0,0 +1,545 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package main + +import ( + "context" + "encoding/json" + "fmt" + "log" + "net/http" + "net/url" + "os" + "os/signal" + "strings" + "syscall" + "time" + + "github.com/golang-jwt/jwt/v4" + mcp "trpc.group/trpc-go/trpc-mcp-go" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth/server" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth/server/providers" +) + +const hmacSecret = "demo-shared-secret" + +// strPtr returns a pointer to the given string +func strPtr(s string) *string { + return &s +} + +// mustURL parses the given string as a URL and panics if invalid +func mustURL(s string) *url.URL { + u, err := url.Parse(s) + if err != nil { + panic(err) + } + return u +} + +func main() { + log.Println("🚀 Starting OAuth Authentication Server...") + log.Println(" Mock OAuth Server: http://localhost:3030") + log.Println(" MCP Server: http://localhost:3000/mcp") + log.Println() + + // Start the mock OAuth server first + go startMockOAuthServer() + time.Sleep(2 * time.Second) + + // Test the mock server + resp, err := http.Get("http://localhost:3030/authorize?test=1") + if err != nil { + log.Fatalf("Mock OAuth server not ready: %v", err) + } + resp.Body.Close() + log.Println("✅ OAuth infrastructure ready") + + // Create OAuth Provider + provider := providers.NewProxyOAuthServerProvider(providers.ProxyOptions{ + Endpoints: providers.ProxyEndpoints{ + AuthorizationURL: "http://localhost:3030/authorize", + TokenURL: "http://localhost:3030/token", + RevocationURL: "http://localhost:3030/revoke", + RegistrationURL: "http://localhost:3030/register", + }, + + VerifyAccessToken: func(token string) (*server.AuthInfo, error) { + ai, err := mockVerifyJWT(token) + if err != nil { + fmt.Printf("❌ Token verification failed: %v\n", err) + return nil, err + } + return &ai, nil + }, + + GetClient: func(clientID string) (*auth.OAuthClientInformationFull, error) { + return &auth.OAuthClientInformationFull{ + OAuthClientMetadata: auth.OAuthClientMetadata{ + RedirectURIs: []string{"http://localhost:5173/callback"}, + ResponseTypes: []string{"code"}, + GrantTypes: []string{"authorization_code", "refresh_token"}, + ClientName: strPtr("demo-client"), + Scope: strPtr("mcp.read mcp.write"), + }, + OAuthClientInformation: auth.OAuthClientInformation{ + ClientID: clientID, + ClientSecret: "", // Public client, no secret + }, + }, nil + }, + }) + + // Create and start the MCP server + // Build a TokenVerifier (use introspection for the demo) + ctx := context.Background() + v, err := server.NewTokenVerifier(ctx, server.TokenVerifierConfig{ + Introspection: &server.IntrospectionConfig{ + Endpoint: "http://localhost:3030/introspect", + Timeout: 5 * time.Second, + CacheTTL: 30 * time.Second, + NegativeCacheTTL: 10 * time.Second, + UseOnJWTFail: true, + }, + }) + if err != nil { + log.Fatalf("failed to create TokenVerifier: %v", err) + } + + mcpServer := mcp.NewServer( + "Auth-Example-Server", + "1.0.0", + mcp.WithServerAddress(":3000"), + mcp.WithServerPath("/mcp"), + mcp.WithOAuthRoutes(mcp.OAuthRoutesConfig{ + Provider: provider, + IssuerURL: mustURL("http://localhost:3030"), + BaseURL: mustURL("http://localhost:3000"), + ScopesSupported: []string{"mcp.read", "mcp.write"}, + }), + mcp.WithOAuthMetadata(mcp.OAuthMetadataConfig{ + ResourceServerURL: mustURL("http://localhost:3000"), + ScopesSupported: []string{"mcp.read", "mcp.write"}, + ResourceName: strPtr("MCP Server"), + }), + mcp.WithBearerAuth(&mcp.BearerAuthConfig{ + Enabled: true, + RequiredScopes: []string{"mcp.read", "mcp.write"}, + Verifier: v, // directly use TokenVerifier implementation + }), + mcp.WithAudit(&mcp.AuditConfig{ + Enabled: true, + Level: "basic", // Reduced from "detailed" + HashSensitiveData: true, + IncludeRequestBody: false, // Disabled to reduce noise + IncludeResponseBody: false, // Disabled to reduce noise + EndpointPatterns: []string{"/mcp/", "/authorize", "/token"}, + ExcludePatterns: []string{"/healthz"}, + }), + ) + + // Set up a graceful shutdown. + stop := make(chan os.Signal, 1) + signal.Notify(stop, os.Interrupt, syscall.SIGTERM) + + // Start server (run in goroutine). + go func() { + log.Println("🔐 MCP Auth Server started successfully") + log.Println(" Waiting for authentication requests...") + fmt.Println() + if err := mcpServer.Start(); err != nil { + log.Fatalf("Server failed to start: %v", err) + } + }() + // Wait for termination signal. + <-stop + log.Println("🛑 Shutting down server...") +} + +// startMockOAuthServer starts a simple mock OAuth server on port 3030 +func startMockOAuthServer() { + mux := http.NewServeMux() + + // Store the authorization code + var authCode = "mock_auth_code_12345" + + // Authorize endpoint + mux.HandleFunc("/authorize", func(w http.ResponseWriter, r *http.Request) { + // Handle test requests silently + if r.URL.Query().Get("test") != "" { + w.WriteHeader(http.StatusOK) + w.Write([]byte("OK")) + return + } + + redirectURI := r.URL.Query().Get("redirect_uri") + state := r.URL.Query().Get("state") + clientID := r.URL.Query().Get("client_id") + + log.Printf("🔐 OAuth Authorization Request\n") + log.Printf(" Client ID: %s\n", clientID) + log.Printf(" Scopes: %s\n", r.URL.Query().Get("scope")) + + if redirectURI == "" { + http.Error(w, "Missing redirect_uri", http.StatusBadRequest) + return + } + + // Construct redirect URLs + redirectURL := redirectURI + "?code=" + authCode + if state != "" { + redirectURL += "&state=" + state + } + + log.Printf(" Redirecting to client callback\n\n") + http.Redirect(w, r, redirectURL, http.StatusFound) + }) + + // Token endpoint: supports authorization_code and refresh_token, and issues HS256 JWT + mux.HandleFunc("/token", func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) + return + } + + // Unified parsing form + if err := r.ParseForm(); err != nil { + http.Error(w, "Invalid form", http.StatusBadRequest) + return + } + + grantType := r.FormValue("grant_type") + + // Issuing HS256 JWT + signJWT := func(claims jwt.MapClaims) (string, error) { + tok := jwt.NewWithClaims(jwt.SigningMethodHS256, claims) + signed, err := tok.SignedString([]byte(hmacSecret)) + if err != nil { + return "", err + } + return signed, nil + } + + switch grantType { + case "authorization_code": + clientID := r.FormValue("client_id") + code := r.FormValue("code") + + log.Printf("🎫 Token Exchange (Authorization Code)\n") + log.Printf(" Client ID: %s\n", clientID) + log.Printf(" Code: %s\n", code) + + // Basic parameter verification + if clientID == "" || code == "" { + http.Error(w, "Missing required parameters", http.StatusBadRequest) + return + } + + now := time.Now() + // Issue access_token + accessToken, err := signJWT(jwt.MapClaims{ + "iss": "http://localhost:3030", + "aud": "http://localhost:3000", + "iat": now.Unix(), + "exp": now.Add(1 * time.Hour).Unix(), + "client_id": clientID, + "sub": clientID, + "scope": "mcp.read mcp.write", + }) + if err != nil { + http.Error(w, "failed to sign access token", http.StatusInternalServerError) + return + } + + // Issue refresh token + refreshToken, err := signJWT(jwt.MapClaims{ + "iss": "http://localhost:3030", + "aud": "http://localhost:3000", + "iat": now.Unix(), + "exp": now.Add(24 * time.Hour).Unix(), + "client_id": clientID, + "sub": clientID, + "typ": "refresh", + }) + if err != nil { + http.Error(w, "failed to sign refresh token", http.StatusInternalServerError) + return + } + + log.Printf(" ✅ Tokens issued successfully\n\n") + + resp := map[string]any{ + "access_token": accessToken, + "token_type": "Bearer", + "expires_in": 3600, + "scope": "mcp.read mcp.write", + "refresh_token": refreshToken, + } + + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(resp) + + case "refresh_token": + rt := r.FormValue("refresh_token") + if rt == "" { + http.Error(w, "Missing refresh_token", http.StatusBadRequest) + return + } + + log.Printf("🔄 Token Refresh Request\n") + + // Parse and verify RT (HS256) + parsed, err := jwt.Parse(rt, func(t *jwt.Token) (interface{}, error) { + if _, ok := t.Method.(*jwt.SigningMethodHMAC); !ok { + return nil, fmt.Errorf("unexpected signing method: %v", t.Header["alg"]) + } + return []byte(hmacSecret), nil + }) + if err != nil || !parsed.Valid { + log.Printf(" ❌ Invalid refresh token\n\n") + http.Error(w, "invalid refresh_token", http.StatusUnauthorized) + return + } + claims, ok := parsed.Claims.(jwt.MapClaims) + if !ok { + http.Error(w, "invalid refresh_token claims", http.StatusUnauthorized) + return + } + + // Extract client_id from RT claims + clientID, _ := claims["client_id"].(string) + if clientID == "" { + clientID = "public-client" + } + + fmt.Printf(" Client ID: %s\n", clientID) + + now := time.Now() + // New access_token + newAT, err := signJWT(jwt.MapClaims{ + "iss": "http://localhost:3030", + "aud": "http://localhost:3000", + "iat": now.Unix(), + "exp": now.Add(1 * time.Hour).Unix(), + "client_id": clientID, + "sub": clientID, + "scope": "mcp.read mcp.write", + }) + if err != nil { + http.Error(w, "failed to sign access token", http.StatusInternalServerError) + return + } + + // New refresh_token + newRT, err := signJWT(jwt.MapClaims{ + "iss": "http://localhost:3030", + "aud": "http://localhost:3000", + "iat": now.Unix(), + "exp": now.Add(24 * time.Hour).Unix(), + "client_id": clientID, + "sub": clientID, + "typ": "refresh", + }) + if err != nil { + http.Error(w, "failed to sign refresh token", http.StatusInternalServerError) + return + } + + fmt.Printf(" ✅ New tokens issued\n\n") + + resp := map[string]any{ + "access_token": newAT, + "token_type": "Bearer", + "expires_in": 3600, + "scope": "mcp.read mcp.write", + "refresh_token": newRT, + } + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(resp) + + default: + http.Error(w, "unsupported_grant_type", http.StatusBadRequest) + return + } + }) + + // Revocation endpoint (optional) - silent + mux.HandleFunc("/revoke", func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + }) + + // Registration endpoint (optional) + mux.HandleFunc("/register", func(w http.ResponseWriter, r *http.Request) { + log.Printf("📝 Client Registration Request\n") + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]interface{}{ + "client_id": "test-client-id", + "client_secret": "test-secret", + "client_name": "demo-client", + "scope": "mcp.read mcp.write", + "redirect_uris": []string{"http://localhost:5173/callback"}, + "grant_types": []string{"authorization_code", "refresh_token"}, + "response_types": []string{"code"}, + }) + log.Printf(" ✅ Client registered: test-client-id\n\n") + }) + + // Introspection endpoint (RFC7662 simplified for demo) + mux.HandleFunc("/introspect", func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) + return + } + if err := r.ParseForm(); err != nil { + http.Error(w, "invalid form", http.StatusBadRequest) + return + } + token := r.FormValue("token") + resp := map[string]any{"active": false} + if token != "" { + parsed, err := jwt.Parse(token, func(t *jwt.Token) (interface{}, error) { + if _, ok := t.Method.(*jwt.SigningMethodHMAC); !ok { + return nil, fmt.Errorf("unexpected signing method: %v", t.Header["alg"]) + } + return []byte(hmacSecret), nil + }) + if err == nil && parsed != nil && parsed.Valid { + if claims, ok := parsed.Claims.(jwt.MapClaims); ok { + var exp int64 + if v, ok := claims["exp"].(float64); ok { + exp = int64(v) + } + scope, _ := claims["scope"].(string) + clientID, _ := claims["client_id"].(string) + if clientID == "" { + if sub, _ := claims["sub"].(string); sub != "" { + clientID = sub + } + } + resp = map[string]any{ + "active": true, + "exp": exp, + "scope": scope, + "client_id": clientID, + "aud": "http://localhost:3000", + "iss": "http://localhost:3030", + } + } + } + } + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(resp) + }) + + // Authorization Server Metadata (RFC 8414) - silent + mux.HandleFunc("/.well-known/oauth-authorization-server", func(w http.ResponseWriter, r *http.Request) { + meta := map[string]any{ + "issuer": "http://localhost:3030", + "authorization_endpoint": "http://localhost:3030/authorize", + "token_endpoint": "http://localhost:3030/token", + "registration_endpoint": "http://localhost:3030/register", + "revocation_endpoint": "http://localhost:3030/revoke", + "response_types_supported": []string{"code"}, + "grant_types_supported": []string{"authorization_code", "refresh_token"}, + "code_challenge_methods_supported": []string{"S256"}, + "token_endpoint_auth_methods_supported": []string{"client_secret_post"}, + "scopes_supported": []string{"mcp.read", "mcp.write"}, + } + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(meta) + }) + + // Compatible with OIDC discovery - silent + mux.HandleFunc("/.well-known/openid-configuration", func(w http.ResponseWriter, r *http.Request) { + cfg := map[string]any{ + "issuer": "http://localhost:3030", + "authorization_endpoint": "http://localhost:3030/authorize", + "token_endpoint": "http://localhost:3030/token", + "registration_endpoint": "http://localhost:3030/register", + "revocation_endpoint": "http://localhost:3030/revoke", + "response_types_supported": []string{"code"}, + "grant_types_supported": []string{"authorization_code", "refresh_token"}, + "code_challenge_methods_supported": []string{"S256"}, + "token_endpoint_auth_methods_supported": []string{"client_secret_post"}, + "scopes_supported": []string{"mcp.read", "mcp.write"}, + } + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(cfg) + }) + + server := &http.Server{ + Addr: ":3030", + Handler: mux, + } + + if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed { + log.Printf("Mock OAuth server error: %v", err) + } +} + +// mockVerifyJWT verifies a JWT using HMAC and extracts AuthInfo +func mockVerifyJWT(token string) (server.AuthInfo, error) { + parsed, err := jwt.Parse(token, func(t *jwt.Token) (interface{}, error) { + if _, ok := t.Method.(*jwt.SigningMethodHMAC); !ok { + return nil, fmt.Errorf("unexpected signing method: %v", t.Header["alg"]) + } + return []byte(hmacSecret), nil + }) + if err != nil || !parsed.Valid { + return server.AuthInfo{}, fmt.Errorf("invalid token: %w", err) + } + + claims, ok := parsed.Claims.(jwt.MapClaims) + if !ok { + return server.AuthInfo{}, fmt.Errorf("invalid claims") + } + + // Parse client_id or sub + var clientID string + if cid, _ := claims["client_id"].(string); cid != "" { + clientID = cid + } + if sub, _ := claims["sub"].(string); sub != "" { + clientID = sub + } + + // Parse scope + scopeStr, _ := claims["scope"].(string) + var scopes []string + if scopeStr != "" { + scopes = strings.Split(scopeStr, " ") + } + + // Parse exp + var expPtr *int64 + if v, ok := claims["exp"].(float64); ok { + vv := int64(v) + expPtr = &vv + } + + // Make sure that Extra contains sub + client_id + if claims["client_id"] == nil && clientID != "" { + claims["client_id"] = clientID + } + if claims["sub"] == nil && clientID != "" { + claims["sub"] = clientID + } + + return server.AuthInfo{ + Token: token, + ClientID: clientID, + Scopes: scopes, + ExpiresAt: expPtr, + Extra: map[string]any{ + "client_id": clientID, + "sub": clientID, + "scope": strings.Join(scopes, " "), + "exp": expPtr, + }, + }, nil +} diff --git a/go.mod b/go.mod index 911a3fe..f1cec73 100644 --- a/go.mod +++ b/go.mod @@ -4,21 +4,43 @@ go 1.20 require ( github.com/getkin/kin-openapi v0.124.0 + github.com/go-playground/validator/v10 v10.22.0 + github.com/golang-jwt/jwt/v4 v4.5.2 + github.com/google/uuid v1.6.0 + github.com/lestrrat-go/jwx/v2 v2.0.21 github.com/stretchr/testify v1.10.0 github.com/yosida95/uritemplate/v3 v3.0.2 go.uber.org/zap v1.27.0 + golang.org/x/time v0.10.0 ) require ( github.com/davecgh/go-spew v1.1.1 // indirect - github.com/go-openapi/jsonpointer v0.20.2 // indirect - github.com/go-openapi/swag v0.22.8 // indirect + github.com/decred/dcrd/dcrec/secp256k1/v4 v4.4.0 // indirect + github.com/gabriel-vasile/mimetype v1.4.8 // indirect + github.com/go-openapi/jsonpointer v0.21.1 // indirect + github.com/go-openapi/swag v0.23.1 // indirect + github.com/go-playground/locales v0.14.1 // indirect + github.com/go-playground/universal-translator v0.18.1 // indirect + github.com/goccy/go-json v0.10.3 // indirect github.com/invopop/yaml v0.2.0 // indirect github.com/josharian/intern v1.0.0 // indirect - github.com/mailru/easyjson v0.7.7 // indirect + github.com/leodido/go-urn v1.4.0 // indirect + github.com/lestrrat-go/blackmagic v1.0.2 // indirect + github.com/lestrrat-go/httpcc v1.0.1 // indirect + github.com/lestrrat-go/httprc v1.0.6 // indirect + github.com/lestrrat-go/iter v1.0.2 // indirect + github.com/lestrrat-go/option v1.0.1 // indirect + github.com/mailru/easyjson v0.9.0 // indirect github.com/mohae/deepcopy v0.0.0-20170929034955-c48cc78d4826 // indirect github.com/perimeterx/marshmallow v1.1.5 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect + github.com/segmentio/asm v1.2.0 // indirect + github.com/ugorji/go/codec v1.2.12 // indirect go.uber.org/multierr v1.10.0 // indirect + golang.org/x/crypto v0.31.0 // indirect + golang.org/x/net v0.33.0 // indirect + golang.org/x/sys v0.28.0 // indirect + golang.org/x/text v0.21.0 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect ) diff --git a/go.sum b/go.sum index c15603c..301073b 100644 --- a/go.sum +++ b/go.sum @@ -1,20 +1,71 @@ +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/decred/dcrd/crypto/blake256 v1.0.0/go.mod h1:sQl2p6Y26YV+ZOcSTP6thNdn47hh8kt6rqSlvmrXFAc= +github.com/decred/dcrd/dcrec/secp256k1/v4 v4.0.1/go.mod h1:hyedUtir6IdtD/7lIxGeCxkaw7y45JueMRL4DIyJDKs= +github.com/decred/dcrd/dcrec/secp256k1/v4 v4.4.0 h1:NMZiJj8QnKe1LgsbDayM4UoHwbvwDRwnI3hwNaAHRnc= +github.com/decred/dcrd/dcrec/secp256k1/v4 v4.4.0/go.mod h1:ZXNYxsqcloTdSy/rNShjYzMhyjf0LaoftYK0p+A3h40= +github.com/gabriel-vasile/mimetype v1.4.8 h1:FfZ3gj38NjllZIeJAmMhr+qKL8Wu+nOoI3GqacKw1NM= +github.com/gabriel-vasile/mimetype v1.4.8/go.mod h1:ByKUIKGjh1ODkGM1asKUbQZOLGrPjydw3hYPU2YU9t8= github.com/getkin/kin-openapi v0.124.0 h1:VSFNMB9C9rTKBnQ/fpyDU8ytMTr4dWI9QovSKj9kz/M= github.com/getkin/kin-openapi v0.124.0/go.mod h1:wb1aSZA/iWmorQP9KTAS/phLj/t17B5jT7+fS8ed9NM= -github.com/go-openapi/jsonpointer v0.20.2 h1:mQc3nmndL8ZBzStEo3JYF8wzmeWffDH4VbXz58sAx6Q= -github.com/go-openapi/jsonpointer v0.20.2/go.mod h1:bHen+N0u1KEO3YlmqOjTT9Adn1RfD91Ar825/PuiRVs= -github.com/go-openapi/swag v0.22.8 h1:/9RjDSQ0vbFR+NyjGMkFTsA1IA0fmhKSThmfGZjicbw= -github.com/go-openapi/swag v0.22.8/go.mod h1:6QT22icPLEqAM/z/TChgb4WAveCHF92+2gF0CNjHpPI= +github.com/getkin/kin-openapi v0.131.0 h1:NO2UeHnFKRYhZ8wg6Nyh5Cq7dHk4suQQr72a4pMrDxE= +github.com/getkin/kin-openapi v0.131.0/go.mod h1:3OlG51PCYNsPByuiMB0t4fjnNlIDnaEDsjiKUV8nL58= +github.com/go-openapi/jsonpointer v0.21.1 h1:whnzv/pNXtK2FbX/W9yJfRmE2gsmkfahjMKB0fZvcic= +github.com/go-openapi/jsonpointer v0.21.1/go.mod h1:50I1STOfbY1ycR8jGz8DaMeLCdXiI6aDteEdRNNzpdk= +github.com/go-openapi/swag v0.23.1 h1:lpsStH0n2ittzTnbaSloVZLuB5+fvSY/+hnagBjSNZU= +github.com/go-openapi/swag v0.23.1/go.mod h1:STZs8TbRvEQQKUA+JZNAm3EWlgaOBGpyFDqQnDHMef0= +github.com/go-playground/assert/v2 v2.2.0 h1:JvknZsQTYeFEAhQwI4qEt9cyV5ONwRHC+lYKSsYSR8s= +github.com/go-playground/assert/v2 v2.2.0/go.mod h1:VDjEfimB/XKnb+ZQfWdccd7VUvScMdVu0Titje2rxJ4= +github.com/go-playground/locales v0.14.1 h1:EWaQ/wswjilfKLTECiXz7Rh+3BjFhfDFKv/oXslEjJA= +github.com/go-playground/locales v0.14.1/go.mod h1:hxrqLVvrK65+Rwrd5Fc6F2O76J/NuW9t0sjnWqG1slY= +github.com/go-playground/universal-translator v0.18.1 h1:Bcnm0ZwsGyWbCzImXv+pAJnYK9S473LQFuzCbDbfSFY= +github.com/go-playground/universal-translator v0.18.1/go.mod h1:xekY+UJKNuX9WP91TpwSH2VMlDf28Uj24BCp08ZFTUY= +github.com/go-playground/validator/v10 v10.22.0 h1:k6HsTZ0sTnROkhS//R0O+55JgM8C4Bx7ia+JlgcnOao= +github.com/go-playground/validator/v10 v10.22.0/go.mod h1:dbuPbCMFw/DrkbEynArYaCwl3amGuJotoKCe95atGMM= github.com/go-test/deep v1.0.8 h1:TDsG77qcSprGbC6vTN8OuXp5g+J+b5Pcguhf7Zt61VM= +github.com/go-test/deep v1.0.8/go.mod h1:5C2ZWiW0ErCdrYzpqxLbTX7MG14M9iiw8DgHncVwcsE= +github.com/goccy/go-json v0.9.7/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I= +github.com/goccy/go-json v0.10.3 h1:KZ5WoDbxAIgm2HNbYckL0se1fHD6rz5j4ywS6ebzDqA= +github.com/goccy/go-json v0.10.3/go.mod h1:oq7eo15ShAhp70Anwd5lgX2pLfOS3QCiwU/PULtXL6M= +github.com/golang-jwt/jwt/v4 v4.5.2 h1:YtQM7lnr8iZ+j5q71MGKkNw9Mn7AjHM68uc9g5fXeUI= +github.com/golang-jwt/jwt/v4 v4.5.2/go.mod h1:m21LjoU+eqJr34lmDMbreY2eSTRJ1cv77w39/MY0Ch0= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/invopop/yaml v0.2.0 h1:7zky/qH+O0DwAyoobXUqvVBwgBFRxKoQ/3FjcVpjTMY= github.com/invopop/yaml v0.2.0/go.mod h1:2XuRLgs/ouIrW3XNzuNj7J3Nvu/Dig5MXvbCEdiBN3Q= github.com/josharian/intern v1.0.0 h1:vlS4z54oSdjm0bgjRigI+G1HpF+tI+9rE5LLzOg8HmY= github.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFFd8Hwg//Y= github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= +github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= -github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0= -github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc= +github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= +github.com/leodido/go-urn v1.4.0 h1:WT9HwE9SGECu3lg4d/dIA+jxlljEa1/ffXKmRjqdmIQ= +github.com/leodido/go-urn v1.4.0/go.mod h1:bvxc+MVxLKB4z00jd1z+Dvzr47oO32F/QSNjSBOlFxI= +github.com/lestrrat-go/blackmagic v1.0.1/go.mod h1:UrEqBzIR2U6CnzVyUtfM6oZNMt/7O7Vohk2J0OGSAtU= +github.com/lestrrat-go/blackmagic v1.0.2 h1:Cg2gVSc9h7sz9NOByczrbUvLopQmXrfFx//N+AkAr5k= +github.com/lestrrat-go/blackmagic v1.0.2/go.mod h1:UrEqBzIR2U6CnzVyUtfM6oZNMt/7O7Vohk2J0OGSAtU= +github.com/lestrrat-go/blackmagic v1.0.3/go.mod h1:6AWFyKNNj0zEXQYfTMPfZrAXUWUfTIZ5ECEUEJaijtw= +github.com/lestrrat-go/httpcc v1.0.1 h1:ydWCStUeJLkpYyjLDHihupbn2tYmZ7m22BGkcvZZrIE= +github.com/lestrrat-go/httpcc v1.0.1/go.mod h1:qiltp3Mt56+55GPVCbTdM9MlqhvzyuL6W/NMDA8vA5E= +github.com/lestrrat-go/httprc v1.0.1/go.mod h1:5Ml+nB++j6IC0e6LzefJnrpMQDKgDwDCaIQQzhbqhJM= +github.com/lestrrat-go/httprc v1.0.6 h1:qgmgIRhpvBqexMJjA/PmwSvhNk679oqD1RbovdCGW8k= +github.com/lestrrat-go/httprc v1.0.6/go.mod h1:mwwz3JMTPBjHUkkDv/IGJ39aALInZLrhBp0X7KGUZlo= +github.com/lestrrat-go/httprc/v2 v2.0.0 h1:dZia9gCSXkYYZN9YUe4U3KU4rvpKXzmGB4QTYDDrOU0= +github.com/lestrrat-go/httprc/v2 v2.0.0/go.mod h1:smhwnjMK58yn+xnN/hxtdSRW2PCi9vNTZDzB85bxj24= +github.com/lestrrat-go/iter v1.0.2 h1:gMXo1q4c2pHmC3dn8LzRhJfP1ceCbgSiT9lUydIzltI= +github.com/lestrrat-go/iter v1.0.2/go.mod h1:Momfcq3AnRlRjI5b5O8/G5/BvpzrhoFTZcn06fEOPt4= +github.com/lestrrat-go/jwx/v2 v2.0.2 h1:wkq9jwCkF3xrykISzn0Eksd7NEMOZ9yvCdnEpovIJX8= +github.com/lestrrat-go/jwx/v2 v2.0.2/go.mod h1:xV8+xRcrKbmnScV8adOzUuuTrL8aAZJoY4q2JAqIYU8= +github.com/lestrrat-go/jwx/v2 v2.0.21 h1:jAPKupy4uHgrHFEdjVjNkUgoBKtVDgrQPB/h55FHrR0= +github.com/lestrrat-go/jwx/v2 v2.0.21/go.mod h1:09mLW8zto6bWL9GbwnqAli+ArLf+5M33QLQPDggkUWM= +github.com/lestrrat-go/jwx/v2 v2.1.6 h1:hxM1gfDILk/l5ylers6BX/Eq1m/pnxe9NBwW6lVfecA= +github.com/lestrrat-go/jwx/v2 v2.1.6/go.mod h1:Y722kU5r/8mV7fYDifjug0r8FK8mZdw0K0GpJw/l8pU= +github.com/lestrrat-go/option v1.0.0/go.mod h1:5ZHFbivi4xwXxhxY9XHDe2FHo6/Z7WWmtT7T5nBBp3I= +github.com/lestrrat-go/option v1.0.1 h1:oAzP2fvZGQKWkvHa1/SAcFolBEca1oN+mQ7eooNBEYU= +github.com/lestrrat-go/option v1.0.1/go.mod h1:5ZHFbivi4xwXxhxY9XHDe2FHo6/Z7WWmtT7T5nBBp3I= +github.com/mailru/easyjson v0.9.0 h1:PrnmzHw7262yW8sTBwxi1PdJA3Iw/EKBa8psRf7d9a4= +github.com/mailru/easyjson v0.9.0/go.mod h1:1+xMtQp2MRNVL/V1bOzuP3aP8VNwRW55fQUto+XFtTU= github.com/mohae/deepcopy v0.0.0-20170929034955-c48cc78d4826 h1:RWengNIwukTxcDr9M+97sNutRR1RKhG96O6jWumTTnw= github.com/mohae/deepcopy v0.0.0-20170929034955-c48cc78d4826/go.mod h1:TaXosZuwdSHYgviHp1DAtfrULt5eUgsSMsZf+YrPgl8= github.com/perimeterx/marshmallow v1.1.5 h1:a2LALqQ1BlHM8PZblsDdidgv1mWi1DgC2UmX50IvK2s= @@ -22,18 +73,57 @@ github.com/perimeterx/marshmallow v1.1.5/go.mod h1:dsXbUu8CRzfYP5a87xpp0xq9S3u0V github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/rogpeppe/go-internal v1.12.0 h1:exVL4IDcn6na9z1rAb56Vxr+CgyK3nn3O+epU5NdKM8= +github.com/rogpeppe/go-internal v1.12.0/go.mod h1:E+RYuTGaKKdloAfM02xzb0FW3Paa99yedzYV+kq4uf4= +github.com/segmentio/asm v1.2.0 h1:9BQrFxC+YOHJlTlHGkTrFWf59nbL3XnCoFLTwDCI7ys= +github.com/segmentio/asm v1.2.0/go.mod h1:BqMnlJP91P8d+4ibuonYZw9mfnzI9HfxselHZr5aAcs= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= -github.com/ugorji/go/codec v1.2.7 h1:YPXUKf7fYbp/y8xloBqZOw2qaVggbfwMlI8WM3wZUJ0= +github.com/ugorji/go/codec v1.2.12 h1:9LC83zGrHhuUA9l16C9AHXAqEV/2wBQ4nkvumAE65EE= +github.com/ugorji/go/codec v1.2.12/go.mod h1:UNopzCgEMSXjBc6AOMqYvWC1ktqTAfzJZUZgYf6w6lg= github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4= github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4= go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= +go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= go.uber.org/multierr v1.10.0 h1:S0h4aNzvfcFsC3dRF1jLoaov7oRaKqRGC/pUEJ2yvPQ= go.uber.org/multierr v1.10.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN80Y= go.uber.org/zap v1.27.0 h1:aJMhYGrd5QSmlpLMr2MftRKl7t8J8PTZPA732ud/XR8= go.uber.org/zap v1.27.0/go.mod h1:GB2qFLM7cTU87MWRP2mPIjqfIDnGu+VIO4V/SdhGo2E= +golang.org/x/crypto v0.0.0-20220427172511-eb4f295cb31f/go.mod h1:IxCIyHEi3zRg3s0A5j5BB6A9Jmi73HwBIUl50j+osU4= +golang.org/x/crypto v0.31.0 h1:ihbySMvVjLAeSH1IbfcRTkD/iNscyz8rGzjF/E5hV6U= +golang.org/x/crypto v0.31.0/go.mod h1:kDsLvtWBEx7MV9tJOj9bnXsPbxwJQ6csT/x4KIN4Ssk= +golang.org/x/crypto v0.32.0/go.mod h1:ZnnJkOaASj8g0AjIduWNlq2NRxL0PlBrbKVyZ6V/Ugc= +golang.org/x/crypto v0.35.0 h1:b15kiHdrGCHrP6LvwaQ3c03kgNhhiMgvlhxHQhmg2Xs= +golang.org/x/crypto v0.35.0/go.mod h1:dy7dXNW32cAb/6/PRuTNsix8T+vJAqvuIy5Bli/x0YQ= +golang.org/x/crypto v0.36.0/go.mod h1:Y4J0ReaxCR1IMaabaSMugxJES1EpwhBHhv2bDHklZvc= +golang.org/x/net v0.0.0-20211112202133-69e39bad7dc2/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y= +golang.org/x/net v0.33.0 h1:74SYHlV8BIgHIFC/LrYkOGIwL19eTYXQ5wc6TBuO36I= +golang.org/x/net v0.33.0/go.mod h1:HXLR5J+9DxmrqMwG9qjGCxZ+zKXxBru04zlTvWlWuN4= +golang.org/x/net v0.38.0 h1:vRMAPTMaeGqVhG5QyLJHqNDwecKTomGeqbnfZyKlBI8= +golang.org/x/net v0.38.0/go.mod h1:ivrbrMbzFq5J41QOQh0siUuly180yBYtLp+CKbEaFx8= +golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20210423082822-04245dca01da/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.28.0 h1:Fksou7UEQUWlKvIdsqzJmUmCX3cZuD2+P3XyyzwMhlA= +golang.org/x/sys v0.28.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/sys v0.30.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/sys v0.31.0 h1:ioabZlmFYtWhL+TRYpcnNlLwhyxaM9kWTDEmfnprqik= +golang.org/x/sys v0.31.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k= +golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= +golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= +golang.org/x/text v0.21.0 h1:zyQAAkrwaneQ066sspRyJaG9VNi/YJ1NfzcGB3hZ/qo= +golang.org/x/text v0.21.0/go.mod h1:4IBbMaMmOPCJ8SecivzSH54+73PCFmPWxNTLm+vZkEQ= +golang.org/x/text v0.22.0/go.mod h1:YRoo4H8PVmsu+E3Ou7cqLVH8oXWIHVoX0jqUWALQhfY= +golang.org/x/text v0.23.0/go.mod h1:/BLNzu4aZCJ1+kcD0DNRotWKage4q2rGVAg4o22unh4= +golang.org/x/time v0.10.0 h1:3usCWA8tQn0L8+hFJQNgzpWbd89begxN66o1Ojdn5L4= +golang.org/x/time v0.10.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM= +golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gopkg.in/yaml.v3 v3.0.0/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/internal/auth/client/flow.go b/internal/auth/client/flow.go new file mode 100644 index 0000000..72318fb --- /dev/null +++ b/internal/auth/client/flow.go @@ -0,0 +1,1256 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package client + +import ( + "context" + "encoding/base64" + "encoding/json" + stderrors "errors" + "fmt" + "io" + "net/http" + "net/url" + "slices" + "strings" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth/pkce" + + "trpc.group/trpc-go/trpc-mcp-go/internal/auth" + "trpc.group/trpc-go/trpc-mcp-go/internal/errors" +) + +// AuthResult describes the outcome of an OAuth flow +type AuthResult string + +const ( + AuthResultAuthorized AuthResult = "AUTHORIZED" + AuthResultRedirect AuthResult = "REDIRECT" +) + +// ClientAuthMethod lists supported client authentication methods for token endpoint +type ClientAuthMethod string + +const ( + ClientAuthMethodBasic ClientAuthMethod = "client_secret_basic" + ClientAuthMethodPost ClientAuthMethod = "client_secret_post" + ClientAuthMethodNone ClientAuthMethod = "none" +) + +// metadataDiscoveryOptions carries optional knobs for metadata discovery behavior +type metadataDiscoveryOptions struct { + ProtocolVersion *string + MetadataUrl *string + MetadataServerUrl *string +} + +// RegisterClientOptions configures dynamic client registration +type RegisterClientOptions struct { + Metadata auth.AuthorizationServerMetadata + ClientMetadata auth.OAuthClientMetadata + FetchFn auth.FetchFunc +} + +// discoveryUrlType distinguishes between OAuth and OIDC discovery endpoints +type discoveryUrlType string + +const ( + discoveryTypeOAuth discoveryUrlType = "oauth" + discoveryTypeOIDC discoveryUrlType = "oidc" +) + +// discoveryUrl pairs a URL with its discovery type +type discoveryUrl struct { + URL *url.URL + Type discoveryUrlType +} + +// StartAuthorizationOptions configures OAuth authorization startup +type StartAuthorizationOptions struct { + // Metadata contains authorization server configuration (optional) + Metadata auth.AuthorizationServerMetadata + + // ClientInformation holds the OAuth client credentials + ClientInformation auth.OAuthClientInformation + + // RedirectURL specifies where to redirect after authorization + RedirectURL string + + // Scope defines the requested access permissions (optional) + Scope *string + + // State provides CSRF protection (optional) + State *string + + // Resource specifies the target resource URL (optional) + Resource *url.URL +} + +// StartAuthorizationResult holds authorization startup results +type StartAuthorizationResult struct { + // AuthorizationURL is where the user should be redirected for authorization + AuthorizationURL *url.URL + + // CodeVerifier must be stored securely for the token exchange step + CodeVerifier string +} + +// ExchangeAuthorizationOptions configures exchanging an authorization code for tokens +type ExchangeAuthorizationOptions struct { + Metadata auth.AuthorizationServerMetadata // server config (optional) + ClientInformation *auth.OAuthClientInformation // client credentials + AuthorizationCode string // auth code from server + CodeVerifier string // PKCE verifier + RedirectURI string // must match auth request + Resource *url.URL // target resource (optional) + AddClientAuthentication func(http.Header, url.Values, string) error // custom auth (optional) + FetchFn auth.FetchFunc // custom HTTP client (optional) +} + +// RefreshAuthorizationOptions configures exchanging a refresh token for new tokens +type RefreshAuthorizationOptions struct { + Metadata auth.AuthorizationServerMetadata // server config (optional) + ClientInformation *auth.OAuthClientInformation // client credentials + RefreshToken string // refresh token + Resource *url.URL // target resource (optional) + AddClientAuthentication func(http.Header, url.Values, string) error // custom auth (optional) + FetchFn auth.FetchFunc // custom HTTP client (optional) +} + +// UnauthorizedError represents an authorization failure that should be surfaced to callers +type UnauthorizedError struct { + message string +} + +// NewUnauthorizedError constructs an UnauthorizedError with a friendly message +func NewUnauthorizedError(message string) *UnauthorizedError { + if message == "" { + message = "Unauthorized" + } + return &UnauthorizedError{message: message} +} + +// Error returns the error message for UnauthorizedError +func (e *UnauthorizedError) Error() string { + return e.message +} + +// selectClientAuthMethod chooses a client auth method based on server support and client secrets +func selectClientAuthMethod( + clientInformation auth.OAuthClientInformation, + supportedMethods []string, +) ClientAuthMethod { + var hasClientSecret bool + hasClientSecret = clientInformation.ClientSecret != "" + if len(supportedMethods) == 0 { + if hasClientSecret { + return ClientAuthMethodPost + } else { + return ClientAuthMethodNone + } + } + + if hasClientSecret && slices.Contains(supportedMethods, string(ClientAuthMethodBasic)) { + return ClientAuthMethodBasic + } + if hasClientSecret && slices.Contains(supportedMethods, string(ClientAuthMethodPost)) { + return ClientAuthMethodPost + } + if slices.Contains(supportedMethods, string(ClientAuthMethodNone)) { + return ClientAuthMethodNone + } + if hasClientSecret { + return ClientAuthMethodPost + } else { + return ClientAuthMethodNone + } +} + +// applyClientAuthentication applies the chosen client auth to headers and or form parameters +func applyClientAuthentication( + method ClientAuthMethod, + clientInformation auth.OAuthClientInformation, + headers http.Header, + params url.Values, +) error { + clientID := clientInformation.ClientID + clientSecret := clientInformation.ClientSecret + + switch method { + case ClientAuthMethodBasic: + return applyBasicAuth(clientID, clientSecret, headers) + case ClientAuthMethodPost: + applyPostAuth(clientID, clientSecret, params) + return nil + case ClientAuthMethodNone: + applyPublicAuth(clientID, params) + return nil + default: + return fmt.Errorf("unsupported client authentication method: %s", method) + } +} + +// applyBasicAuth adds HTTP Basic Authorization using client id and secret +func applyBasicAuth(clientID, clientSecret string, headers http.Header) error { + if clientSecret == "" { + return fmt.Errorf("client_secret_basic authentication requires a client_secret") + } + + credentials := base64.StdEncoding.EncodeToString([]byte(clientID + ":" + clientSecret)) + headers.Set("Authorization", "Basic "+credentials) + return nil +} + +// applyPostAuth writes client credentials into the token form payload +func applyPostAuth(clientID, clientSecret string, params url.Values) { + params.Set("client_id", clientID) + if clientSecret != "" { + params.Set("client_secret", clientSecret) + } +} + +// applyPublicAuth writes public client id into the token form payload +func applyPublicAuth(clientID string, params url.Values) { + params.Set("client_id", clientID) +} + +// parseErrorResponse converts an OAuth style JSON error payload into an OAuthError +func parseErrorResponse(input interface{}) (*errors.OAuthError, error) { + var responseBody []byte + var err error + + // Handle different input types + switch v := input.(type) { + case []byte: + responseBody = v + case string: + responseBody = []byte(v) + case *http.Response: + defer v.Body.Close() + responseBody, err = io.ReadAll(v.Body) + if err != nil { + return nil, fmt.Errorf("failed to read response body: %w", err) + } + default: + return nil, fmt.Errorf("unsupported input type: %T", input) + } + + // Try to parse as OAuth error response + var oauthErrorResp errors.OAuthErrorResponse + if err := json.Unmarshal(responseBody, &oauthErrorResp); err != nil { + // Not a valid OAuth error response format + return nil, fmt.Errorf("failed to parse OAuth error response: %w", err) + } + + // Validate required error field + if oauthErrorResp.Error == "" { + return nil, fmt.Errorf("invalid OAuth error response: missing error field") + } + + // Map error code to OAuthErrorCode using the mapping table + errorCode, exists := errors.OAuthErrorMapping[oauthErrorResp.Error] + if !exists { + // Unknown error code, default to server error + errorCode = errors.ErrServerError + } + + // Create OAuthError with parsed information + oauthError := errors.NewOAuthError(errorCode, oauthErrorResp.ErrorDescription, oauthErrorResp.ErrorURI) + return &oauthError, nil +} + +// Auth performs high level authentication with retries and credential invalidation on specific errors +func Auth(provider OAuthClientProvider, options auth.AuthOptions) (*AuthResult, error) { + result, err := authInternal(provider, options) + if err != nil { + if stderrors.Is(err, errors.ErrInvalidClient) || stderrors.Is(err, errors.ErrUnauthorizedClient) { + if invalidator, ok := provider.(OAuthCredentialInvalidator); ok { + if invalidateErr := invalidator.InvalidateCredentials("all"); invalidateErr != nil { + return nil, invalidateErr + } + } + return authInternal(provider, options) + } else if stderrors.Is(err, errors.ErrInvalidGrant) { + if invalidator, ok := provider.(OAuthCredentialInvalidator); ok { + if invalidateErr := invalidator.InvalidateCredentials("tokens"); invalidateErr != nil { + return nil, invalidateErr + } + } + return authInternal(provider, options) + } + return nil, err + } + return result, err +} + +// authInternal runs the core auth logic including discovery registration refresh and redirect setup +func authInternal(provider OAuthClientProvider, options auth.AuthOptions) (*AuthResult, error) { + var resourceMetadata *auth.OAuthProtectedResourceMetadata + var authorizationServerUrl string + metadata, err := DiscoverOAuthProtectedResourceMetadata(options.ServerUrl, &auth.DiscoveryOptions{ + ResourceMetadataUrl: options.ResourceMetadataUrl, + }, options.FetchFn) + if err == nil { + resourceMetadata = metadata + if len(resourceMetadata.AuthorizationServers) > 0 { + authorizationServerUrl = resourceMetadata.AuthorizationServers[0] + } + } + if authorizationServerUrl == "" { + authorizationServerUrl = options.ServerUrl + } + + resource, err := selectResourceURL(options.ServerUrl, provider, resourceMetadata) + if err != nil { + return nil, fmt.Errorf("failed to select resource URL: %w", err) + } + + serverMetadata, err := DiscoverAuthorizationServerMetadata(context.Background(), authorizationServerUrl, nil) + if err != nil { + return nil, fmt.Errorf("failed to discover authorization server metadata: %w", err) + } + clientInformation := provider.ClientInformation() + + if clientInformation == nil { + if options.AuthorizationCode != nil { + return nil, stderrors.New("existing OAuth client information is required when exchanging an authorization code") + } + + if _, ok := provider.(OAuthClientInfoProvider); !ok { + return nil, stderrors.New("OAuth client information must be saveable for dynamic registration") + } + + fullInformation, err := RegisterClient(context.Background(), authorizationServerUrl, RegisterClientOptions{ + Metadata: serverMetadata, + ClientMetadata: provider.ClientMetadata(), + FetchFn: options.FetchFn, + }) + if err != nil { + return nil, fmt.Errorf("failed to register client: %w", err) + } + + if clientInfoProvider, ok := provider.(OAuthClientInfoProvider); ok { + if err := clientInfoProvider.SaveClientInformation(*fullInformation); err != nil { + return nil, fmt.Errorf("failed to save client information: %w", err) + } + } + clientInformation = &auth.OAuthClientInformation{ + ClientID: fullInformation.ClientID, + ClientSecret: fullInformation.ClientSecret, + } + } + tokens, err := provider.Tokens() + if err != nil { + return nil, fmt.Errorf("failed to get tokens: %w", err) + } + + if tokens != nil && tokens.RefreshToken != nil && *tokens.RefreshToken != "" { + var addClientAuth func(http.Header, url.Values, string) error + if authProvider, ok := provider.(OAuthClientAuthProvider); ok { + addClientAuth = authProvider.AddClientAuthentication + } + + newTokens, err := RefreshAuthorization(authorizationServerUrl, RefreshAuthorizationOptions{ + Metadata: serverMetadata, + ClientInformation: clientInformation, + RefreshToken: *tokens.RefreshToken, + Resource: resource, + AddClientAuthentication: addClientAuth, + FetchFn: options.FetchFn, + }) + if err != nil { + var oauthErr *errors.OAuthError + if !stderrors.As(err, &oauthErr) { + // Network/non-OAuth errors, continue auth flow + } else { + return nil, err + } + } else { + if err := provider.SaveTokens(*newTokens); err != nil { + return nil, fmt.Errorf("failed to save refreshed tokens: %w", err) + } + result := AuthResultAuthorized + return &result, nil + } + } + var state *string + if stateProvider, ok := provider.(OAuthStateProvider); ok { + stateValue, err := stateProvider.State() + if err != nil { + return nil, fmt.Errorf("failed to get state: %w", err) + } + state = &stateValue + } + scope := options.Scope + if scope == nil { + clientMetadata := provider.ClientMetadata() + if clientMetadata.Scope != nil { + scope = clientMetadata.Scope + } + } + + authorizationResult, err := StartAuthorization(authorizationServerUrl, StartAuthorizationOptions{ + Metadata: serverMetadata, + ClientInformation: *clientInformation, + State: state, + RedirectURL: provider.RedirectURL(), + Scope: scope, + Resource: resource, + }) + if err != nil { + return nil, fmt.Errorf("failed to start authorization: %w", err) + } + + if err := provider.SaveCodeVerifier(authorizationResult.CodeVerifier); err != nil { + return nil, fmt.Errorf("failed to save code verifier: %w", err) + } + + if err := provider.RedirectToAuthorization(authorizationResult.AuthorizationURL); err != nil { + return nil, fmt.Errorf("failed to redirect to authorization: %w", err) + } + + result := AuthResultRedirect + return &result, nil +} + +// selectResourceURL determines the resource parameter to use validating against protected resource metadata +func selectResourceURL(serverUrl string, provider OAuthClientProvider, resourceMetadata *auth.OAuthProtectedResourceMetadata) (*url.URL, error) { + defaultResource, err := auth.ResourceURLFromServerURL(serverUrl) + if err != nil { + return nil, err + } + + // Use custom validator if available + if validator, ok := provider.(OAuthResourceValidator); ok { + return validator.ValidateResourceURL(defaultResource, resourceMetadata) + } + + // Include resource param only when metadata exists + if resourceMetadata == nil { + return nil, nil // No resource param needed + } + + // Check metadata resource compatibility + allowed, err := auth.CheckResourceAllowed(auth.CheckResourceAllowedParams{ + RequestedResource: defaultResource, + ConfiguredResource: resourceMetadata.Resource, + }) + if err != nil { + return nil, fmt.Errorf("failed to validate resource: %w", err) + } + if !allowed { + return nil, fmt.Errorf("protected resource %s does not match expected %s", + resourceMetadata.Resource, defaultResource.String()) + } + + // Use metadata resource - server expects this + return url.Parse(resourceMetadata.Resource) +} + +// DiscoverOAuthProtectedResourceMetadata loads OAuth Protected Resource metadata with path aware fallback +func DiscoverOAuthProtectedResourceMetadata(serverUrl string, opts *auth.DiscoveryOptions, fetchFn auth.FetchFunc) (*auth.OAuthProtectedResourceMetadata, error) { + if fetchFn == nil { + fetchFn = func(urlStr string, req *http.Request) (*http.Response, error) { + return http.DefaultClient.Do(req) + } + } + + response, err := discoverMetadataWithFallback( + serverUrl, + "oauth-protected-resource", + fetchFn, + &metadataDiscoveryOptions{ + ProtocolVersion: getProtocolVersion(opts), + MetadataUrl: getResourceMetadataUrl(opts), + }, + ) + if err != nil { + return nil, err + } + + if response == nil || response.StatusCode == 404 { + return nil, fmt.Errorf("Resource server does not implement OAuth 2.0 Protected Resource Metadata.") + } + + if response.StatusCode < 200 || response.StatusCode >= 300 { + return nil, fmt.Errorf("HTTP %d trying to load well-known OAuth protected resource metadata.", response.StatusCode) + } + + defer response.Body.Close() + var metadata auth.OAuthProtectedResourceMetadata + if err := json.NewDecoder(response.Body).Decode(&metadata); err != nil { + return nil, fmt.Errorf("failed to parse metadata response: %w", err) + } + + return &metadata, nil +} + +// discoverMetadataWithFallback tries path aware discovery then falls back to root well known when applicable +func discoverMetadataWithFallback( + serverUrl interface{}, + wellKnownType string, // "oauth-authorization-server" or "oauth-protected-resource" + fetchFn auth.FetchFunc, + opts *metadataDiscoveryOptions, +) (*http.Response, error) { + issuer, err := parseURL(serverUrl) + if err != nil { + return nil, fmt.Errorf("invalid server URL: %w", err) + } + + protocolVersion := "2025-03-26" // LATEST_PROTOCOL_VERSION + if opts != nil && opts.ProtocolVersion != nil { + protocolVersion = *opts.ProtocolVersion + } + + var targetUrl *url.URL + if opts != nil && opts.MetadataUrl != nil { + targetUrl, err = url.Parse(*opts.MetadataUrl) + if err != nil { + return nil, fmt.Errorf("invalid metadata URL: %w", err) + } + } else { + // Try path-aware discovery + wellKnownPath := buildWellKnownPath(wellKnownType, issuer.Path) + baseUrl := issuer + if opts != nil && opts.MetadataServerUrl != nil { + baseUrl, err = url.Parse(*opts.MetadataServerUrl) + if err != nil { + return nil, fmt.Errorf("invalid metadata server URL: %w", err) + } + } + targetUrl, _ = url.Parse(wellKnownPath) + targetUrl = baseUrl.ResolveReference(targetUrl) + targetUrl.RawQuery = issuer.RawQuery + } + + response, err := tryMetadataDiscovery(targetUrl, protocolVersion, fetchFn) + if err != nil { + return nil, err + } + + // If path-aware discovery fails with 404 and we're not at root, try fallback to root discovery + if (opts == nil || opts.MetadataUrl == nil) && shouldAttemptFallback(response, issuer.Path) { + rootUrl, _ := url.Parse(fmt.Sprintf("/.well-known/%s", wellKnownType)) + rootUrl = issuer.ResolveReference(rootUrl) + response, err = tryMetadataDiscovery(rootUrl, protocolVersion, fetchFn) + if err != nil { + return nil, err + } + } + + return response, nil +} + +// tryMetadataDiscovery issues a discovery request with protocol headers and returns the HTTP response +func tryMetadataDiscovery(targetUrl *url.URL, protocolVersion string, fetchFn auth.FetchFunc) (*http.Response, error) { + req, err := http.NewRequest("GET", targetUrl.String(), nil) + if err != nil { + return nil, fmt.Errorf("failed to create request: %w", err) + } + + req.Header.Set("MCP-Protocol-Version", protocolVersion) + + return fetchWithCorsRetry(targetUrl, req.Header, fetchFn) +} + +// fetchWithCorsRetry helper function to handle CORS retry logic +func fetchWithCorsRetry(targetUrl *url.URL, headers http.Header, fetchFn auth.FetchFunc) (*http.Response, error) { + req, err := http.NewRequest("GET", targetUrl.String(), nil) + if err != nil { + return nil, fmt.Errorf("failed to create request: %w", err) + } + + // Copy headers + for key, values := range headers { + for _, value := range values { + req.Header.Add(key, value) + } + } + + response, err := fetchFn(targetUrl.String(), req) + if err != nil { + // If it's a network error (similar to TypeError in TypeScript), try retry without headers + if isNetworkError(err) && len(headers) > 0 { + return fetchWithCorsRetry(targetUrl, http.Header{}, fetchFn) + } + return nil, err + } + + return response, nil +} +func shouldAttemptFallback(response *http.Response, pathname string) bool { + return response == nil || (response.StatusCode == 404 && pathname != "/") +} + +// buildWellKnownPath builds well-known path for authentication-related metadata discovery +func buildWellKnownPath(wellKnownPrefix, pathname string) string { + // Remove trailing slash from pathname to avoid double slashes + if strings.HasSuffix(pathname, "/") { + pathname = strings.TrimSuffix(pathname, "/") + } + + return fmt.Sprintf("/.well-known/%s%s", wellKnownPrefix, pathname) +} + +// parseURL accepts string or *url.URL and returns a parsed URL +func parseURL(u interface{}) (*url.URL, error) { + switch v := u.(type) { + case string: + return url.Parse(v) + case *url.URL: + return v, nil + default: + return nil, fmt.Errorf("unsupported URL type") + } +} + +// getProtocolVersion resolves the requested protocol version from options if provided +func getProtocolVersion(opts *auth.DiscoveryOptions) *string { + if opts != nil && opts.ProtocolVersion != nil { + return opts.ProtocolVersion + } + return nil +} + +// getResourceMetadataUrl resolves an explicit resource metadata URL from options if provided +func getResourceMetadataUrl(opts *auth.DiscoveryOptions) *string { + if opts != nil && opts.ResourceMetadataUrl != nil { + return opts.ResourceMetadataUrl + } + return nil +} + +// isNetworkError determines if it's a network error (simulating TypeError check in TypeScript) +func isNetworkError(err error) bool { + // In Go, network errors usually contain these keywords + errorStr := strings.ToLower(err.Error()) + return strings.Contains(errorStr, "network") || + strings.Contains(errorStr, "connection") || + strings.Contains(errorStr, "timeout") || + strings.Contains(errorStr, "refused") +} + +// buildDiscoveryUrls generates candidate OAuth and OIDC discovery URLs for a given authorization server URL +func buildDiscoveryUrls(authorizationServerURL string) ([]discoveryUrl, error) { + parsedURL, err := url.Parse(authorizationServerURL) + if err != nil { + return nil, fmt.Errorf("invalid authorization server URL: %w", err) + } + + hasPath := parsedURL.Path != "/" && parsedURL.Path != "" + var urlsToTry []discoveryUrl + + if !hasPath { + // Root path: https://example.com/.well-known/oauth-authorization-server + oauthURL, _ := url.Parse(parsedURL.Scheme + "://" + parsedURL.Host + "/.well-known/oauth-authorization-server") + urlsToTry = append(urlsToTry, discoveryUrl{URL: oauthURL, Type: discoveryTypeOAuth}) + + // OIDC: https://example.com/.well-known/openid-configuration + oidcURL, _ := url.Parse(parsedURL.Scheme + "://" + parsedURL.Host + "/.well-known/openid-configuration") + urlsToTry = append(urlsToTry, discoveryUrl{URL: oidcURL, Type: discoveryTypeOIDC}) + + return urlsToTry, nil + } + + // Strip trailing slash from pathname to avoid double slashes + pathname := parsedURL.Path + if strings.HasSuffix(pathname, "/") { + pathname = pathname[:len(pathname)-1] + } + + // 1. OAuth metadata at the given URL + // Insert well-known before the path: https://example.com/.well-known/oauth-authorization-server/tenant1 + oauthWithPath, _ := url.Parse(parsedURL.Scheme + "://" + parsedURL.Host + "/.well-known/oauth-authorization-server" + pathname) + urlsToTry = append(urlsToTry, discoveryUrl{URL: oauthWithPath, Type: discoveryTypeOAuth}) + + // Root path: https://example.com/.well-known/oauth-authorization-server + oauthRoot, _ := url.Parse(parsedURL.Scheme + "://" + parsedURL.Host + "/.well-known/oauth-authorization-server") + urlsToTry = append(urlsToTry, discoveryUrl{URL: oauthRoot, Type: discoveryTypeOAuth}) + + // 3. OIDC metadata endpoints + // RFC 8414 style: Insert /.well-known/openid-configuration before the path + oidcWithPath, _ := url.Parse(parsedURL.Scheme + "://" + parsedURL.Host + "/.well-known/openid-configuration" + pathname) + urlsToTry = append(urlsToTry, discoveryUrl{URL: oidcWithPath, Type: discoveryTypeOIDC}) + + // OIDC Discovery 1.0 style: Append /.well-known/openid-configuration after the path + oidcAfterPath, _ := url.Parse(parsedURL.Scheme + "://" + parsedURL.Host + pathname + "/.well-known/openid-configuration") + urlsToTry = append(urlsToTry, discoveryUrl{URL: oidcAfterPath, Type: discoveryTypeOIDC}) + + return urlsToTry, nil +} + +// DiscoveryURL represents a metadata endpoint +type DiscoveryURL struct { + URL *url.URL + Type string // "oauth" or "oidc" +} + +// DiscoverAuthorizationServerMetadata discovers OAuth or OIDC metadata and verifies minimal capabilities +func DiscoverAuthorizationServerMetadata(ctx context.Context, authServerUrl string, options *auth.DiscoveryOptions) (auth.AuthorizationServerMetadata, error) { + // Build discovery URLs + discoveryUrls, err := buildDiscoveryUrls(authServerUrl) + if err != nil { + return nil, err + } + + // Create default fetch function + fetchFunc := func(urlStr string, req *http.Request) (*http.Response, error) { + return http.DefaultClient.Do(req) + } + + // Try each discovery URL + for _, discoveryUrl := range discoveryUrls { + // Create headers + headers := http.Header{ + "Accept": []string{"application/json"}, + } + + // Try to fetch metadata + resp, err := fetchWithCorsRetry(discoveryUrl.URL, headers, fetchFunc) + if err != nil { + if isNetworkError(err) { + continue + } + return nil, fmt.Errorf("failed to fetch metadata from %s: %w", discoveryUrl.URL.String(), err) + } + defer resp.Body.Close() + + // Check status code + if resp.StatusCode == 404 { + continue + } + if resp.StatusCode != 200 { + return nil, fmt.Errorf("unexpected status code %d from %s", resp.StatusCode, discoveryUrl.URL.String()) + } + + // Read response body + body, err := io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("failed to read response body: %w", err) + } + + // Parse metadata based on type + if discoveryUrl.Type == "oidc" { + // Try to parse as OpenID Connect metadata + var metadata auth.OpenIdProviderDiscoveryMetadata + if err := json.Unmarshal(body, &metadata); err != nil { + continue // Try next URL + } + + // Validate required fields for OIDC + if metadata.Issuer == "" || metadata.AuthorizationEndpoint == "" || metadata.TokenEndpoint == "" { + continue // Try next URL + } + + // Check if S256 PKCE is supported for OIDC + supportsS256 := false + for _, method := range metadata.CodeChallengeMethodsSupported { + if method == "S256" { + supportsS256 = true + break + } + } + if !supportsS256 { + return nil, fmt.Errorf("OIDC provider does not support S256 PKCE") + } + + return &metadata, nil + } else { + // Try to parse as OAuth 2.0 metadata + var metadata auth.OAuthMetadata + if err := json.Unmarshal(body, &metadata); err != nil { + continue // Try next URL + } + + // Validate required fields for OAuth 2.0 + if metadata.Issuer == "" || metadata.AuthorizationEndpoint == "" || metadata.TokenEndpoint == "" { + continue // Try next URL + } + + return &metadata, nil + } + } + + return nil, fmt.Errorf("failed to discover authorization server metadata from %s", authServerUrl) +} + +// RegisterClient performs dynamic client registration and returns full client information +func RegisterClient( + ctx context.Context, + authorizationServerUrl string, + options RegisterClientOptions, +) (*auth.OAuthClientInformationFull, error) { + var registrationUrl *url.URL + var err error + + // Determine registration endpoint URL + if options.Metadata != nil { + // Check if dynamic client registration is supported + var registrationEndpoint string + + // Get registration endpoint based on metadata type + switch metadata := options.Metadata.(type) { + case *auth.OAuthMetadata: + if metadata.RegistrationEndpoint == nil { + return nil, fmt.Errorf("incompatible auth server: does not support dynamic client registration") + } + registrationEndpoint = *metadata.RegistrationEndpoint + case *auth.OpenIdProviderMetadata: + if metadata.RegistrationEndpoint == nil { + return nil, fmt.Errorf("incompatible auth server: does not support dynamic client registration") + } + registrationEndpoint = *metadata.RegistrationEndpoint + case *auth.OpenIdProviderDiscoveryMetadata: + if metadata.RegistrationEndpoint == nil { + return nil, fmt.Errorf("incompatible auth server: does not support dynamic client registration") + } + registrationEndpoint = *metadata.RegistrationEndpoint + default: + return nil, fmt.Errorf("unsupported metadata type") + } + + registrationUrl, err = url.Parse(registrationEndpoint) + if err != nil { + return nil, fmt.Errorf("invalid registration endpoint URL: %w", err) + } + } else { + // Use default registration path + baseUrl, err := url.Parse(authorizationServerUrl) + if err != nil { + return nil, fmt.Errorf("invalid authorization server URL: %w", err) + } + registrationUrl, err = baseUrl.Parse("/register") + if err != nil { + return nil, fmt.Errorf("failed to construct registration URL: %w", err) + } + } + + // Serialize client metadata + requestBody, err := json.Marshal(options.ClientMetadata) + if err != nil { + return nil, fmt.Errorf("failed to marshal client metadata: %w", err) + } + + // Create HTTP request + req, err := http.NewRequestWithContext(ctx, "POST", registrationUrl.String(), strings.NewReader(string(requestBody))) + if err != nil { + return nil, fmt.Errorf("failed to create request: %w", err) + } + + req.Header.Set("Content-Type", "application/json") + + // Select fetch function + fetchFn := options.FetchFn + if fetchFn == nil { + fetchFn = func(url string, req *http.Request) (*http.Response, error) { + return http.DefaultClient.Do(req) + } + } + + // Send request + resp, err := fetchFn(registrationUrl.String(), req) + if err != nil { + return nil, fmt.Errorf("failed to send registration request: %w", err) + } + defer resp.Body.Close() + + // Read response body + responseBody, err := io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("failed to read response body: %w", err) + } + + // Check response status + if !isSuccessStatusCode(resp.StatusCode) { + // Try to parse OAuth error response + var oauthError errors.OAuthError + if err := json.Unmarshal(responseBody, &oauthError); err == nil { + return nil, &oauthError + } + return nil, fmt.Errorf("registration failed with status %d: %s", resp.StatusCode, string(responseBody)) + } + + // Parse success response + var clientInfo auth.OAuthClientInformationFull + if err := json.Unmarshal(responseBody, &clientInfo); err != nil { + return nil, fmt.Errorf("failed to parse registration response: %w", err) + } + + return &clientInfo, nil +} + +// isSuccessStatusCode checks if HTTP status code indicates success +func isSuccessStatusCode(statusCode int) bool { + return statusCode >= 200 && statusCode < 300 +} + +// StartAuthorization starts OAuth 2.0 authorization flow +func StartAuthorization( + authorizationServerUrl string, + options StartAuthorizationOptions, +) (*StartAuthorizationResult, error) { + const responseType = "code" + const codeChallengeMethod = "S256" + + var authorizationURL *url.URL + var err error + + // Determine authorization endpoint URL + if options.Metadata != nil { + authorizationURL, err = url.Parse(options.Metadata.GetAuthorizationEndpoint()) + if err != nil { + return nil, fmt.Errorf("invalid authorization endpoint: %w", err) + } + + // Verify server supports "code" response type + responseTypesSupported := options.Metadata.GetResponseTypesSupported() + supportsCode := false + for _, rt := range responseTypesSupported { + if rt == responseType { + supportsCode = true + break + } + } + if !supportsCode { + return nil, fmt.Errorf( + "incompatible auth server: does not support response type %s", + responseType, + ) + } + + // Verify server supports S256 PKCE method + var codeChallengeMethodsSupported []string + + // Check different types of metadata + switch metadata := options.Metadata.(type) { + case *auth.OAuthMetadata: + codeChallengeMethodsSupported = metadata.CodeChallengeMethodsSupported + case *auth.OpenIdProviderDiscoveryMetadata: + codeChallengeMethodsSupported = metadata.CodeChallengeMethodsSupported + } + + if len(codeChallengeMethodsSupported) > 0 { + supportsS256 := false + for _, method := range codeChallengeMethodsSupported { + if method == codeChallengeMethod { + supportsS256 = true + break + } + } + if !supportsS256 { + return nil, fmt.Errorf( + "incompatible auth server: does not support code challenge method %s", + codeChallengeMethod, + ) + } + } + } else { + // If no metadata, use default /authorize endpoint + baseURL, err := url.Parse(authorizationServerUrl) + if err != nil { + return nil, fmt.Errorf("invalid authorization server URL: %w", err) + } + authorizationURL = baseURL.ResolveReference(&url.URL{Path: "/authorize"}) + } + + // Generate PKCE challenge + challenge, err := pkce.GeneratePKCEChallenge() + if err != nil { + return nil, fmt.Errorf("failed to generate PKCE challenge: %w", err) + } + + // Build query parameters + params := url.Values{} + params.Set("response_type", responseType) + params.Set("client_id", options.ClientInformation.ClientID) + params.Set("redirect_uri", options.RedirectURL) + params.Set("code_challenge", challenge.CodeChallenge) + params.Set("code_challenge_method", codeChallengeMethod) + + // Add optional parameters + if options.Scope != nil && *options.Scope != "" { + params.Set("scope", *options.Scope) + + // OpenID Connect requirement: if scope contains 'offline_access', need to add consent prompt + if strings.Contains(*options.Scope, "offline_access") { + params.Set("prompt", "consent") + } + } + + if options.State != nil && *options.State != "" { + params.Set("state", *options.State) + } + + if options.Resource != nil { + params.Set("resource", options.Resource.String()) + } + + // Set query parameters + authorizationURL.RawQuery = params.Encode() + + return &StartAuthorizationResult{ + AuthorizationURL: authorizationURL, + CodeVerifier: challenge.CodeVerifier, + }, nil +} + +// ExchangeAuthorization exchanges an authorization code for tokens applying appropriate client authentication +func ExchangeAuthorization( + authorizationServerUrl string, + options ExchangeAuthorizationOptions, +) (*auth.OAuthTokens, error) { + const grantType = "authorization_code" + + // Determine token endpoint URL + var tokenURL *url.URL + var err error + + if options.Metadata != nil { + tokenEndpoint := options.Metadata.GetTokenEndpoint() + if tokenEndpoint == "" { + return nil, fmt.Errorf("token endpoint not found in metadata") + } + tokenURL, err = url.Parse(tokenEndpoint) + if err != nil { + return nil, fmt.Errorf("invalid token endpoint: %w", err) + } + + // Verify server supports authorization_code grant type + grantTypesSupported := options.Metadata.GetGrantTypesSupported() + if len(grantTypesSupported) > 0 { + supportsAuthCode := false + for _, gt := range grantTypesSupported { + if gt == grantType { + supportsAuthCode = true + break + } + } + if !supportsAuthCode { + return nil, fmt.Errorf( + "incompatible auth server: does not support grant type %s", + grantType, + ) + } + } + } else { + // Use default /token endpoint + baseURL, err := url.Parse(authorizationServerUrl) + if err != nil { + return nil, fmt.Errorf("invalid authorization server URL: %w", err) + } + tokenURL = baseURL.ResolveReference(&url.URL{Path: "/token"}) + } + + // Prepare request headers and parameters + headers := http.Header{ + "Content-Type": []string{"application/x-www-form-urlencoded"}, + } + params := url.Values{ + "grant_type": []string{grantType}, + "code": []string{options.AuthorizationCode}, + "redirect_uri": []string{options.RedirectURI}, + "code_verifier": []string{options.CodeVerifier}, + } + + // Apply client authentication + if options.AddClientAuthentication != nil { + if err := options.AddClientAuthentication(headers, params, authorizationServerUrl); err != nil { + return nil, fmt.Errorf("failed to apply client authentication: %w", err) + } + } else { + // Determine and apply client authentication method + var supportedMethods []string + if options.Metadata != nil { + supportedMethods = options.Metadata.GetTokenEndpointAuthMethodsSupported() + + } + authMethod := selectClientAuthMethod(*options.ClientInformation, supportedMethods) + if err := applyClientAuthentication(authMethod, *options.ClientInformation, headers, params); err != nil { + return nil, fmt.Errorf("failed to apply client authentication: %w", err) + } + } + + // Add resource parameter (if provided) + if options.Resource != nil { + params.Set("resource", options.Resource.String()) + } + + // Create HTTP request + req, err := http.NewRequest("POST", tokenURL.String(), strings.NewReader(params.Encode())) + if err != nil { + return nil, fmt.Errorf("failed to create request: %w", err) + } + req.Header = headers + + // Select fetch function + fetchFn := options.FetchFn + if fetchFn == nil { + fetchFn = func(url string, req *http.Request) (*http.Response, error) { + return http.DefaultClient.Do(req) + } + } + + // Send request + resp, err := fetchFn(tokenURL.String(), req) + if err != nil { + return nil, fmt.Errorf("failed to send token request: %w", err) + } + defer resp.Body.Close() + + // Read response body + responseBody, err := io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("failed to read response body: %w", err) + } + + // Check response status + if !isSuccessStatusCode(resp.StatusCode) { + // Try to parse OAuth error response + var oauthError errors.OAuthError + if err := json.Unmarshal(responseBody, &oauthError); err == nil { + return nil, &oauthError + } + return nil, fmt.Errorf("token exchange failed with status %d: %s", resp.StatusCode, string(responseBody)) + } + + // Parse success response + var tokens auth.OAuthTokens + if err := json.Unmarshal(responseBody, &tokens); err != nil { + return nil, fmt.Errorf("failed to parse token response: %w", err) + } + + return &tokens, nil +} + +// RefreshAuthorization exchanges a refresh token for a new access token and propagates refresh token when absent +func RefreshAuthorization( + authorizationServerUrl string, + options RefreshAuthorizationOptions, +) (*auth.OAuthTokens, error) { + const grantType = "refresh_token" + + // Determine token endpoint URL + var tokenURL *url.URL + var err error + + if options.Metadata != nil { + tokenEndpoint := options.Metadata.GetTokenEndpoint() + if tokenEndpoint == "" { + return nil, fmt.Errorf("token endpoint not found in metadata") + } + tokenURL, err = url.Parse(tokenEndpoint) + if err != nil { + return nil, fmt.Errorf("invalid token endpoint: %w", err) + } + + // Verify server supports refresh_token grant type + grantTypesSupported := options.Metadata.GetGrantTypesSupported() + if len(grantTypesSupported) > 0 { + supportsRefreshToken := false + for _, gt := range grantTypesSupported { + if gt == grantType { + supportsRefreshToken = true + break + } + } + if !supportsRefreshToken { + return nil, fmt.Errorf( + "incompatible auth server: does not support grant type %s", + grantType, + ) + } + } + } else { + // Use default /token endpoint + baseURL, err := url.Parse(authorizationServerUrl) + if err != nil { + return nil, fmt.Errorf("invalid authorization server URL: %w", err) + } + tokenURL = baseURL.ResolveReference(&url.URL{Path: "/token"}) + } + + // Prepare request headers and parameters + headers := http.Header{ + "Content-Type": []string{"application/x-www-form-urlencoded"}, + } + params := url.Values{ + "grant_type": []string{grantType}, + "refresh_token": []string{options.RefreshToken}, + } + + // Apply client authentication + if options.AddClientAuthentication != nil { + if err := options.AddClientAuthentication(headers, params, authorizationServerUrl); err != nil { + return nil, fmt.Errorf("failed to apply client authentication: %w", err) + } + } else { + // Determine and apply client authentication method + var supportedMethods []string + if options.Metadata != nil { + supportedMethods = options.Metadata.GetTokenEndpointAuthMethodsSupported() + } + authMethod := selectClientAuthMethod(*options.ClientInformation, supportedMethods) + if err := applyClientAuthentication(authMethod, *options.ClientInformation, headers, params); err != nil { + return nil, fmt.Errorf("failed to apply client authentication: %w", err) + } + } + + // Add resource parameter (if provided) + if options.Resource != nil { + params.Set("resource", options.Resource.String()) + } + + // Create HTTP request + req, err := http.NewRequest("POST", tokenURL.String(), strings.NewReader(params.Encode())) + if err != nil { + return nil, fmt.Errorf("failed to create request: %w", err) + } + req.Header = headers + + // Select fetch function + fetchFn := options.FetchFn + if fetchFn == nil { + fetchFn = func(url string, req *http.Request) (*http.Response, error) { + return http.DefaultClient.Do(req) + } + } + + // Send request + resp, err := fetchFn(tokenURL.String(), req) + if err != nil { + return nil, fmt.Errorf("failed to send refresh request: %w", err) + } + defer resp.Body.Close() + + // Read response body + responseBody, err := io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("failed to read response body: %w", err) + } + + // Check response status + if !isSuccessStatusCode(resp.StatusCode) { + // Try to parse OAuth error response + var oauthError errors.OAuthError + if err := json.Unmarshal(responseBody, &oauthError); err == nil { + return nil, &oauthError + } + return nil, fmt.Errorf("token refresh failed with status %d: %s", resp.StatusCode, string(responseBody)) + } + + // Parse success response + var tokens auth.OAuthTokens + if err := json.Unmarshal(responseBody, &tokens); err != nil { + return nil, fmt.Errorf("failed to parse token response: %w", err) + } + + // If response doesn't contain new refresh token, keep the original one + if tokens.RefreshToken == nil || *tokens.RefreshToken == "" { + tokens.RefreshToken = &options.RefreshToken + } + + return &tokens, nil +} diff --git a/internal/auth/client/flow_test.go b/internal/auth/client/flow_test.go new file mode 100644 index 0000000..bdd3b3f --- /dev/null +++ b/internal/auth/client/flow_test.go @@ -0,0 +1,603 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package client + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "testing" + + "trpc.group/trpc-go/trpc-mcp-go/internal/auth" +) + +// mockOAuthClientProvider provides a test double implementing the OAuth client provider interfaces +type mockOAuthClientProvider struct { + clientInfo *auth.OAuthClientInformation + tokens *auth.OAuthTokens + codeVerifier string + redirectURL string + clientMeta auth.OAuthClientMetadata + saveTokensErr error + saveCodeErr error + redirectErr error + invalidateErr error + stateValue string + stateErr error +} + +// ClientInformation returns the current OAuth client credentials used by the client +func (m *mockOAuthClientProvider) ClientInformation() *auth.OAuthClientInformation { + return m.clientInfo +} + +// Tokens returns the currently stored OAuth tokens +func (m *mockOAuthClientProvider) Tokens() (*auth.OAuthTokens, error) { + return m.tokens, nil +} + +// SaveTokens persists newly issued OAuth tokens +func (m *mockOAuthClientProvider) SaveTokens(tokens auth.OAuthTokens) error { + m.tokens = &tokens + return m.saveTokensErr +} + +// SaveCodeVerifier persists the PKCE code verifier for later token exchange +func (m *mockOAuthClientProvider) SaveCodeVerifier(verifier string) error { + m.codeVerifier = verifier + return m.saveCodeErr +} + +// CodeVerifier returns the stored PKCE code verifier +func (m *mockOAuthClientProvider) CodeVerifier() (string, error) { + return m.codeVerifier, nil +} + +// RedirectURL returns the client redirect URL registered with the authorization server +func (m *mockOAuthClientProvider) RedirectURL() string { + return m.redirectURL +} + +// ClientMetadata returns the OAuth client metadata used for dynamic registration +func (m *mockOAuthClientProvider) ClientMetadata() auth.OAuthClientMetadata { + return m.clientMeta +} + +// RedirectToAuthorization performs a redirect to the authorization URL in real implementations +func (m *mockOAuthClientProvider) RedirectToAuthorization(authURL *url.URL) error { + return m.redirectErr +} + +// InvalidateCredentials invalidates cached credentials according to the provided scope +func (m *mockOAuthClientProvider) InvalidateCredentials(scope string) error { + return m.invalidateErr +} + +// SaveClientInformation persists full client information returned by dynamic registration +func (m *mockOAuthClientProvider) SaveClientInformation(info auth.OAuthClientInformationFull) error { + m.clientInfo = &auth.OAuthClientInformation{ + ClientID: info.ClientID, + ClientSecret: info.ClientSecret, + } + return nil +} + +// AddClientAuthentication attaches client authentication to token requests +func (m *mockOAuthClientProvider) AddClientAuthentication(headers http.Header, params url.Values, serverUrl string) error { + return nil +} + +// State returns a CSRF protection state value for authorization requests +func (m *mockOAuthClientProvider) State() (string, error) { + return m.stateValue, m.stateErr +} + +// ValidateResourceURL validates or adjusts the default resource URL using optional metadata +func (m *mockOAuthClientProvider) ValidateResourceURL(defaultResource *url.URL, metadata *auth.OAuthProtectedResourceMetadata) (*url.URL, error) { + return defaultResource, nil +} + +func TestSelectClientAuthMethod(t *testing.T) { + tests := []struct { + name string + clientInfo auth.OAuthClientInformation + supportedMethods []string + expected ClientAuthMethod + }{ + { + name: "basic auth preferred with secret", + clientInfo: auth.OAuthClientInformation{ + ClientID: "test-client", + ClientSecret: "test-secret", + }, + supportedMethods: []string{"client_secret_basic", "client_secret_post"}, + expected: ClientAuthMethodBasic, + }, + { + name: "post auth when basic not supported", + clientInfo: auth.OAuthClientInformation{ + ClientID: "test-client", + ClientSecret: "test-secret", + }, + supportedMethods: []string{"client_secret_post"}, + expected: ClientAuthMethodPost, + }, + { + name: "none auth for public client", + clientInfo: auth.OAuthClientInformation{ + ClientID: "test-client", + }, + supportedMethods: []string{"none"}, + expected: ClientAuthMethodNone, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := selectClientAuthMethod(tt.clientInfo, tt.supportedMethods) + if result != tt.expected { + t.Errorf("selectClientAuthMethod() = %v, want %v", result, tt.expected) + } + }) + } +} + +func TestApplyClientAuthentication(t *testing.T) { + clientInfo := auth.OAuthClientInformation{ + ClientID: "test-client", + ClientSecret: "test-secret", + } + + t.Run("basic auth", func(t *testing.T) { + headers := http.Header{} + params := url.Values{} + + err := applyClientAuthentication(ClientAuthMethodBasic, clientInfo, headers, params) + if err != nil { + t.Fatalf("applyClientAuthentication() error = %v", err) + } + + authHeader := headers.Get("Authorization") + if !strings.HasPrefix(authHeader, "Basic ") { + t.Errorf("Expected Basic authorization header, got %s", authHeader) + } + }) + + t.Run("post auth", func(t *testing.T) { + headers := http.Header{} + params := url.Values{} + + err := applyClientAuthentication(ClientAuthMethodPost, clientInfo, headers, params) + if err != nil { + t.Fatalf("applyClientAuthentication() error = %v", err) + } + + if params.Get("client_id") != "test-client" { + t.Errorf("Expected client_id parameter, got %s", params.Get("client_id")) + } + if params.Get("client_secret") != "test-secret" { + t.Errorf("Expected client_secret parameter, got %s", params.Get("client_secret")) + } + }) +} + +func TestParseErrorResponse(t *testing.T) { + tests := []struct { + name string + input interface{} + expectError bool + expectCode string + }{ + { + name: "valid oauth error from bytes", + input: []byte(`{"error":"invalid_client","error_description":"Client authentication failed"}`), + expectError: false, + expectCode: "invalid_client", + }, + { + name: "missing error field", + input: `{"error_description":"Missing error field"}`, + expectError: true, + }, + { + name: "invalid json", + input: `invalid json`, + expectError: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + oauthErr, err := parseErrorResponse(tt.input) + + if tt.expectError { + if err == nil { + t.Errorf("parseErrorResponse() expected error but got none") + } + return + } + + if err != nil { + t.Fatalf("parseErrorResponse() unexpected error = %v", err) + } + + if oauthErr.ErrorCode != tt.expectCode { + t.Errorf("parseErrorResponse() error code = %v, want %v", oauthErr.ErrorCode, tt.expectCode) + } + }) + } +} + +func TestBuildDiscoveryUrls(t *testing.T) { + tests := []struct { + name string + serverURL string + expectedCount int + expectError bool + }{ + { + name: "root path server", + serverURL: "https://auth.example.com", + expectedCount: 2, + }, + { + name: "server with path", + serverURL: "https://auth.example.com/tenant1", + expectedCount: 4, + }, + { + name: "invalid URL", + serverURL: "://invalid-url", + expectError: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + urls, err := buildDiscoveryUrls(tt.serverURL) + + if tt.expectError { + if err == nil { + t.Errorf("buildDiscoveryUrls() expected error but got none") + } + return + } + + if err != nil { + t.Fatalf("buildDiscoveryUrls() unexpected error = %v", err) + } + + if len(urls) != tt.expectedCount { + t.Errorf("buildDiscoveryUrls() returned %d URLs, want %d", len(urls), tt.expectedCount) + } + }) + } +} + +func TestStartAuthorization(t *testing.T) { + metadata := &auth.OAuthMetadata{ + Issuer: "https://auth.example.com", + AuthorizationEndpoint: "https://auth.example.com/authorize", + TokenEndpoint: "https://auth.example.com/token", + ResponseTypesSupported: []string{"code"}, + CodeChallengeMethodsSupported: []string{"S256"}, + } + + options := StartAuthorizationOptions{ + Metadata: metadata, + ClientInformation: auth.OAuthClientInformation{ + ClientID: "test-client", + }, + RedirectURL: "https://client.example.com/callback", + } + + result, err := StartAuthorization("https://auth.example.com", options) + if err != nil { + t.Fatalf("startAuthorization() error = %v", err) + } + + if result.AuthorizationURL == nil { + t.Error("startAuthorization() AuthorizationURL is nil") + } + + if result.CodeVerifier == "" { + t.Error("startAuthorization() CodeVerifier is empty") + } + + // Check URL parameters + params := result.AuthorizationURL.Query() + if params.Get("response_type") != "code" { + t.Errorf("Expected response_type=code, got %s", params.Get("response_type")) + } + if params.Get("client_id") != "test-client" { + t.Errorf("Expected client_id=test-client, got %s", params.Get("client_id")) + } +} + +func TestDiscoverAuthorizationServerMetadata(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + metadata := auth.OAuthMetadata{ + Issuer: "https://auth.example.com", + AuthorizationEndpoint: "https://auth.example.com/authorize", + TokenEndpoint: "https://auth.example.com/token", + ResponseTypesSupported: []string{"code"}, + GrantTypesSupported: []string{"authorization_code", "refresh_token"}, + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(metadata) + })) + defer server.Close() + + result, err := DiscoverAuthorizationServerMetadata(context.Background(), server.URL, nil) + if err != nil { + t.Fatalf("DiscoverAuthorizationServerMetadata() error = %v", err) + } + + if result.GetIssuer() != "https://auth.example.com" { + t.Errorf("Expected issuer https://auth.example.com, got %s", result.GetIssuer()) + } +} + +func TestAuthWithExistingTokens(t *testing.T) { + provider := &mockOAuthClientProvider{ + clientInfo: &auth.OAuthClientInformation{ + ClientID: "test-client", + ClientSecret: "test-secret", + }, + tokens: &auth.OAuthTokens{ + AccessToken: "valid-access-token", + RefreshToken: stringPtr("valid-refresh-token"), + }, + redirectURL: "https://client.example.com/callback", + clientMeta: auth.OAuthClientMetadata{ + RedirectURIs: []string{"https://client.example.com/callback"}, + }, + } + + var serverURL string + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case strings.HasSuffix(r.URL.Path, "/.well-known/oauth-authorization-server"): + metadata := &auth.OAuthMetadata{ + Issuer: serverURL, // Use server URL as issuer + AuthorizationEndpoint: serverURL + "/authorize", + TokenEndpoint: serverURL + "/token", + ResponseTypesSupported: []string{"code"}, + GrantTypesSupported: []string{"authorization_code", "refresh_token"}, + CodeChallengeMethodsSupported: []string{"S256"}, + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(metadata) + case r.URL.Path == "/token": + // Simulate successful token refresh + tokens := auth.OAuthTokens{ + AccessToken: "new-access-token", + RefreshToken: stringPtr("new-refresh-token"), + TokenType: "Bearer", + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(tokens) + case strings.HasSuffix(r.URL.Path, "/.well-known/oauth-protected-resource"): + // Return 404 for resource metadata (optional) + w.WriteHeader(404) + default: + t.Logf("Unexpected request to: %s", r.URL.Path) + w.WriteHeader(404) + } + })) + defer server.Close() + + serverURL = server.URL + + options := auth.AuthOptions{ + ServerUrl: server.URL, + FetchFn: func(url string, req *http.Request) (*http.Response, error) { + return http.DefaultClient.Do(req) + }, + } + + result, err := Auth(provider, options) + if err != nil { + t.Fatalf("Auth() error = %v", err) + } + + if *result != AuthResultAuthorized { + t.Errorf("Expected AuthResultAuthorized, got %v", *result) + } + + // Verify tokens were updated + if provider.tokens.AccessToken != "new-access-token" { + t.Errorf("Expected access token to be updated to 'new-access-token', got %s", provider.tokens.AccessToken) + } +} + +func TestRegisterClient(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != "POST" { + t.Errorf("Expected POST request, got %s", r.Method) + } + + clientInfo := auth.OAuthClientInformationFull{ + OAuthClientInformation: auth.OAuthClientInformation{ + ClientID: "generated-client-id", + ClientSecret: "generated-client-secret", + }, + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(clientInfo) + })) + defer server.Close() + + registrationEndpoint := server.URL + "/register" + metadata := &auth.OAuthMetadata{ + RegistrationEndpoint: ®istrationEndpoint, + } + + options := RegisterClientOptions{ + Metadata: metadata, + ClientMetadata: auth.OAuthClientMetadata{ + RedirectURIs: []string{"https://client.example.com/callback"}, + }, + FetchFn: func(url string, req *http.Request) (*http.Response, error) { + return http.DefaultClient.Do(req) + }, + } + + result, err := RegisterClient(context.Background(), server.URL, options) + if err != nil { + t.Fatalf("RegisterClient() error = %v", err) + } + + if result.ClientID != "generated-client-id" { + t.Errorf("Expected client_id 'generated-client-id', got %s", result.ClientID) + } +} + +func TestAuthWithoutClientInfo(t *testing.T) { + provider := &mockOAuthClientProvider{ + clientInfo: nil, // No existing client info + redirectURL: "https://client.example.com/callback", + clientMeta: auth.OAuthClientMetadata{ + RedirectURIs: []string{"https://client.example.com/callback"}, + }, + } + + // Capture the server URL using a variable + var serverURL string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case strings.HasSuffix(r.URL.Path, "/.well-known/oauth-authorization-server"): + metadata := &auth.OAuthMetadata{ + Issuer: serverURL, + AuthorizationEndpoint: serverURL + "/authorize", + TokenEndpoint: serverURL + "/token", + RegistrationEndpoint: stringPtr(serverURL + "/register"), + ResponseTypesSupported: []string{"code"}, + GrantTypesSupported: []string{"authorization_code", "refresh_token"}, + CodeChallengeMethodsSupported: []string{"S256"}, + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(metadata) + case r.URL.Path == "/register": + clientInfo := auth.OAuthClientInformationFull{ + OAuthClientInformation: auth.OAuthClientInformation{ + ClientID: "registered-client-id", + ClientSecret: "registered-client-secret", + }, + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(clientInfo) + case strings.HasSuffix(r.URL.Path, "/.well-known/oauth-protected-resource"): + // Return 404 for resource metadata (optional) + w.WriteHeader(404) + default: + t.Logf("Unexpected request to: %s", r.URL.Path) + w.WriteHeader(404) + } + })) + defer server.Close() + + serverURL = server.URL + + options := auth.AuthOptions{ + ServerUrl: server.URL, + FetchFn: func(url string, req *http.Request) (*http.Response, error) { + return http.DefaultClient.Do(req) + }, + } + + result, err := Auth(provider, options) + if err != nil { + t.Fatalf("Auth() error = %v", err) + } + + if *result != AuthResultRedirect { + t.Errorf("Expected AuthResultRedirect, got %v", *result) + } + + // Verify client information was saved after registration + if provider.clientInfo == nil { + t.Error("Expected client information to be saved after registration") + } + if provider.clientInfo.ClientID != "registered-client-id" { + t.Errorf("Expected registered client ID, got %s", provider.clientInfo.ClientID) + } +} + +func TestAuthWithoutRefreshToken(t *testing.T) { + provider := &mockOAuthClientProvider{ + clientInfo: &auth.OAuthClientInformation{ + ClientID: "test-client", + ClientSecret: "test-secret", + }, + tokens: &auth.OAuthTokens{ + AccessToken: "valid-access-token", + // No refresh token + }, + redirectURL: "https://client.example.com/callback", + clientMeta: auth.OAuthClientMetadata{ + RedirectURIs: []string{"https://client.example.com/callback"}, + }, + } + + var serverURL string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case strings.HasSuffix(r.URL.Path, "/.well-known/oauth-authorization-server"): + metadata := &auth.OAuthMetadata{ + Issuer: serverURL, + AuthorizationEndpoint: serverURL + "/authorize", + TokenEndpoint: serverURL + "/token", + ResponseTypesSupported: []string{"code"}, + GrantTypesSupported: []string{"authorization_code", "refresh_token"}, + CodeChallengeMethodsSupported: []string{"S256"}, + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(metadata) + case strings.HasSuffix(r.URL.Path, "/.well-known/oauth-protected-resource"): + w.WriteHeader(404) + default: + w.WriteHeader(404) + } + })) + defer server.Close() + + serverURL = server.URL + + options := auth.AuthOptions{ + ServerUrl: server.URL, + FetchFn: func(url string, req *http.Request) (*http.Response, error) { + return http.DefaultClient.Do(req) + }, + } + + result, err := Auth(provider, options) + if err != nil { + t.Fatalf("Auth() error = %v", err) + } + + if *result != AuthResultRedirect { + t.Errorf("Expected AuthResultRedirect, got %v", *result) + } +} + +func BenchmarkSelectClientAuthMethod(b *testing.B) { + clientInfo := auth.OAuthClientInformation{ + ClientID: "test-client", + ClientSecret: "test-secret", + } + supportedMethods := []string{"client_secret_basic", "client_secret_post"} + + b.ResetTimer() + for i := 0; i < b.N; i++ { + selectClientAuthMethod(clientInfo, supportedMethods) + } +} diff --git a/internal/auth/client/http.go b/internal/auth/client/http.go new file mode 100644 index 0000000..41c74e0 --- /dev/null +++ b/internal/auth/client/http.go @@ -0,0 +1,97 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package client + +import ( + "context" + "time" + + "trpc.group/trpc-go/trpc-mcp-go/internal/auth" +) + +// ctxKey defines a private type for context keys to avoid collisions +type ctxKey int + +const ( + // ctxKeyClientAuthInfo is the context key used to store ClientAuthInfo + ctxKeyClientAuthInfo ctxKey = iota + // ctxKeyClientAuthErr is the context key used to store authentication errors + ctxKeyClientAuthErr +) + +// ClientAuthInfo holds OAuth client authentication details +type ClientAuthInfo struct { + AccessToken string + RefreshToken *string + ExpiresAt *time.Time + Scopes []string + Extra map[string]interface{} +} + +// WithAuthInfo stores authentication information in the context +func WithAuthInfo(ctx context.Context, info *ClientAuthInfo) context.Context { + if info == nil { + return ctx + } + return context.WithValue(ctx, ctxKeyClientAuthInfo, info) +} + +// GetAuthInfo retrieves authentication information from the context +func GetAuthInfo(ctx context.Context) (*ClientAuthInfo, bool) { + v := ctx.Value(ctxKeyClientAuthInfo) + if v == nil { + return nil, false + } + info, ok := v.(*ClientAuthInfo) + return info, ok && info != nil +} + +// WithAuthErr stores an authentication error in the context +func WithAuthErr(ctx context.Context, err error) context.Context { + if err == nil { + return ctx + } + return context.WithValue(ctx, ctxKeyClientAuthErr, err) +} + +// ConvertTokensToAuthInfo converts OAuth tokens into ClientAuthInfo +func ConvertTokensToAuthInfo(tokens *auth.OAuthTokens) *ClientAuthInfo { + if tokens == nil || tokens.AccessToken == "" { + return nil + } + + authInfo := &ClientAuthInfo{ + AccessToken: tokens.AccessToken, + RefreshToken: tokens.RefreshToken, + Scopes: parseTokenScopes(tokens), + Extra: make(map[string]interface{}), + } + + if tokens.ExpiresIn != nil { + expiresAt := time.Now().Add(time.Duration(*tokens.ExpiresIn) * time.Second) + authInfo.ExpiresAt = &expiresAt + } + + return authInfo +} + +// IsTokenExpired checks if the token is expired or near expiry +func IsTokenExpired(authInfo *ClientAuthInfo) bool { + if authInfo == nil { + return true // nil means expired + } + if authInfo.ExpiresAt == nil { + return false // no expiry means never expired + } + // expire if within 30 seconds of expiry + return !authInfo.ExpiresAt.After(time.Now().Add(30 * time.Second)) +} + +// parseTokenScopes extracts scopes from tokens (currently placeholder) +func parseTokenScopes(tokens *auth.OAuthTokens) []string { + return []string{} +} diff --git a/internal/auth/client/http_test.go b/internal/auth/client/http_test.go new file mode 100644 index 0000000..81cd9a7 --- /dev/null +++ b/internal/auth/client/http_test.go @@ -0,0 +1,594 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package client + +import ( + "context" + "errors" + "testing" + "time" + + "trpc.group/trpc-go/trpc-mcp-go/internal/auth" +) + +func TestWithAuthInfo(t *testing.T) { + ctx := context.Background() + + // Test with nil auth info + newCtx := WithAuthInfo(ctx, nil) + if newCtx != ctx { + t.Error("Expected same context when auth info is nil") + } + + // Test with valid auth info + authInfo := &ClientAuthInfo{ + AccessToken: "test-token", + Scopes: []string{"read", "write"}, + } + + newCtx = WithAuthInfo(ctx, authInfo) + if newCtx == ctx { + t.Error("Expected different context when auth info is provided") + } + + // Verify the auth info was stored + stored, ok := GetAuthInfo(newCtx) + if !ok { + t.Error("Expected to retrieve auth info from context") + } + if stored != authInfo { + t.Error("Expected stored auth info to match original") + } +} + +func TestGetAuthInfo(t *testing.T) { + ctx := context.Background() + + // Test with empty context + authInfo, ok := GetAuthInfo(ctx) + if ok { + t.Error("Expected no auth info in empty context") + } + if authInfo != nil { + t.Error("Expected nil auth info from empty context") + } + + // Test with auth info in context + originalAuthInfo := &ClientAuthInfo{ + AccessToken: "test-access-token", + RefreshToken: stringPtr("test-refresh-token"), + Scopes: []string{"read", "write"}, + Extra: map[string]interface{}{"custom": "value"}, + } + + ctx = WithAuthInfo(ctx, originalAuthInfo) + authInfo, ok = GetAuthInfo(ctx) + + if !ok { + t.Error("Expected to find auth info in context") + } + if authInfo == nil { + t.Fatal("Expected non-nil auth info") + } + if authInfo.AccessToken != originalAuthInfo.AccessToken { + t.Errorf("Expected AccessToken %s, got %s", originalAuthInfo.AccessToken, authInfo.AccessToken) + } + if *authInfo.RefreshToken != *originalAuthInfo.RefreshToken { + t.Errorf("Expected RefreshToken %s, got %s", *originalAuthInfo.RefreshToken, *authInfo.RefreshToken) + } + + // Test with nil auth info stored in context + ctx = context.WithValue(context.Background(), ctxKeyClientAuthInfo, (*ClientAuthInfo)(nil)) + authInfo, ok = GetAuthInfo(ctx) + if ok { + t.Error("Expected no auth info when nil is stored") + } + if authInfo != nil { + t.Error("Expected nil auth info when nil is stored") + } + + // Test with wrong type in context + ctx = context.WithValue(context.Background(), ctxKeyClientAuthInfo, "not-auth-info") + authInfo, ok = GetAuthInfo(ctx) + if ok { + t.Error("Expected no auth info when wrong type is stored") + } + if authInfo != nil { + t.Error("Expected nil auth info when wrong type is stored") + } +} + +func TestWithAuthErr(t *testing.T) { + ctx := context.Background() + + // Test with nil error + newCtx := WithAuthErr(ctx, nil) + if newCtx != ctx { + t.Error("Expected same context when error is nil") + } + + // Test with valid error + testErr := errors.New("auth error") + newCtx = WithAuthErr(ctx, testErr) + if newCtx == ctx { + t.Error("Expected different context when error is provided") + } + + // Verify the error was stored + stored := newCtx.Value(ctxKeyClientAuthErr) + if stored == nil { + t.Error("Expected error to be stored in context") + } + if stored != testErr { + t.Errorf("Expected stored error %v, got %v", testErr, stored) + } +} + +func TestConvertTokensToAuthInfo(t *testing.T) { + // Test with nil tokens + authInfo := ConvertTokensToAuthInfo(nil) + if authInfo != nil { + t.Error("Expected nil auth info when tokens is nil") + } + + // Test with empty access token + emptyTokens := &auth.OAuthTokens{ + AccessToken: "", + } + authInfo = ConvertTokensToAuthInfo(emptyTokens) + if authInfo != nil { + t.Error("Expected nil auth info when access token is empty") + } + + // Test with valid tokens (minimal) + validTokens := &auth.OAuthTokens{ + AccessToken: "test-access-token", + TokenType: "Bearer", + } + authInfo = ConvertTokensToAuthInfo(validTokens) + if authInfo == nil { + t.Fatal("Expected non-nil auth info") + } + if authInfo.AccessToken != validTokens.AccessToken { + t.Errorf("Expected AccessToken %s, got %s", validTokens.AccessToken, authInfo.AccessToken) + } + if authInfo.RefreshToken != nil { + t.Error("Expected nil RefreshToken when not provided") + } + if authInfo.ExpiresAt != nil { + t.Error("Expected nil ExpiresAt when ExpiresIn not provided") + } + if authInfo.Extra == nil { + t.Error("Expected Extra map to be initialized") + } + + // Test with full tokens + refreshToken := "test-refresh-token" + expiresIn := int64(3600) + scope := "read write admin" + fullTokens := &auth.OAuthTokens{ + AccessToken: "test-access-token", + RefreshToken: &refreshToken, + TokenType: "Bearer", + ExpiresIn: &expiresIn, + Scope: &scope, + } + + before := time.Now() + authInfo = ConvertTokensToAuthInfo(fullTokens) + after := time.Now() + + if authInfo == nil { + t.Fatal("Expected non-nil auth info") + } + if authInfo.AccessToken != fullTokens.AccessToken { + t.Errorf("Expected AccessToken %s, got %s", fullTokens.AccessToken, authInfo.AccessToken) + } + if authInfo.RefreshToken == nil { + t.Fatal("Expected RefreshToken to be set") + } + if *authInfo.RefreshToken != *fullTokens.RefreshToken { + t.Errorf("Expected RefreshToken %s, got %s", *fullTokens.RefreshToken, *authInfo.RefreshToken) + } + + // Test ExpiresAt calculation + if authInfo.ExpiresAt == nil { + t.Fatal("Expected ExpiresAt to be set") + } + expectedExpiry := before.Add(time.Duration(expiresIn) * time.Second) + actualExpiry := *authInfo.ExpiresAt + if actualExpiry.Before(expectedExpiry) || actualExpiry.After(after.Add(time.Duration(expiresIn)*time.Second)) { + t.Errorf("Expected ExpiresAt around %v, got %v", expectedExpiry, actualExpiry) + } + + // Test scopes parsing (currently returns empty slice) + if len(authInfo.Scopes) != 0 { + t.Errorf("Expected empty scopes (parseTokenScopes returns empty), got %v", authInfo.Scopes) + } +} + +func TestIsTokenExpired(t *testing.T) { + // Test with nil auth info + if !IsTokenExpired(nil) { + t.Error("Expected nil auth info to be considered expired") + } + + // Test with no expiration time + authInfoNoExpiry := &ClientAuthInfo{ + AccessToken: "test-token", + } + if IsTokenExpired(authInfoNoExpiry) { + t.Error("Expected auth info without expiry to not be expired") + } + + // Test with future expiration (not expired) + futureExpiry := time.Now().Add(time.Hour) + authInfoFuture := &ClientAuthInfo{ + AccessToken: "test-token", + ExpiresAt: &futureExpiry, + } + if IsTokenExpired(authInfoFuture) { + t.Error("Expected auth info with future expiry to not be expired") + } + + // Test with past expiration (expired) + pastExpiry := time.Now().Add(-time.Hour) + authInfoPast := &ClientAuthInfo{ + AccessToken: "test-token", + ExpiresAt: &pastExpiry, + } + if !IsTokenExpired(authInfoPast) { + t.Error("Expected auth info with past expiry to be expired") + } + + // Test with expiration within 30 seconds (considered expired due to buffer) + soonExpiry := time.Now().Add(15 * time.Second) + authInfoSoon := &ClientAuthInfo{ + AccessToken: "test-token", + ExpiresAt: &soonExpiry, + } + if !IsTokenExpired(authInfoSoon) { + t.Error("Expected auth info expiring within 30 seconds to be considered expired") + } + + // Test with expiration just outside 30 second buffer + laterExpiry := time.Now().Add(35 * time.Second) + authInfoLater := &ClientAuthInfo{ + AccessToken: "test-token", + ExpiresAt: &laterExpiry, + } + if IsTokenExpired(authInfoLater) { + t.Error("Expected auth info expiring outside 30 second buffer to not be expired") + } +} + +func TestParseTokenScopes(t *testing.T) { + // Test with nil tokens + scopes := parseTokenScopes(nil) + if len(scopes) != 0 { + t.Errorf("Expected empty scopes for nil tokens, got %v", scopes) + } + + // Test with tokens without scope + tokens := &auth.OAuthTokens{ + AccessToken: "test-token", + } + scopes = parseTokenScopes(tokens) + if len(scopes) != 0 { + t.Errorf("Expected empty scopes when no scope in tokens, got %v", scopes) + } + + // Test with tokens with scope + scope := "read write admin" + tokensWithScope := &auth.OAuthTokens{ + AccessToken: "test-token", + Scope: &scope, + } + scopes = parseTokenScopes(tokensWithScope) + // Currently returns empty slice, but testing the current behavior + if len(scopes) != 0 { + t.Errorf("Expected empty scopes (current implementation), got %v", scopes) + } +} + +func TestClientAuthInfoComplete(t *testing.T) { + // Test complete ClientAuthInfo structure + refreshToken := "refresh-123" + expiresAt := time.Now().Add(time.Hour) + authInfo := &ClientAuthInfo{ + AccessToken: "access-123", + RefreshToken: &refreshToken, + ExpiresAt: &expiresAt, + Scopes: []string{"read", "write", "admin"}, + Extra: map[string]interface{}{ + "user_id": "12345", + "username": "testuser", + "custom": true, + }, + } + + // Test all fields are preserved + if authInfo.AccessToken != "access-123" { + t.Errorf("Expected AccessToken access-123, got %s", authInfo.AccessToken) + } + if authInfo.RefreshToken == nil || *authInfo.RefreshToken != refreshToken { + t.Errorf("Expected RefreshToken %s, got %v", refreshToken, authInfo.RefreshToken) + } + if authInfo.ExpiresAt == nil || !authInfo.ExpiresAt.Equal(expiresAt) { + t.Errorf("Expected ExpiresAt %v, got %v", expiresAt, authInfo.ExpiresAt) + } + if len(authInfo.Scopes) != 3 { + t.Errorf("Expected 3 scopes, got %d", len(authInfo.Scopes)) + } + if authInfo.Extra["user_id"] != "12345" { + t.Errorf("Expected Extra user_id 12345, got %v", authInfo.Extra["user_id"]) + } +} + +func TestContextKeyUniqueness(t *testing.T) { + // Test that context keys are unique + if ctxKeyClientAuthInfo == ctxKeyClientAuthErr { + t.Error("Expected context keys to be unique") + } + + // Test that different values can be stored with different keys + ctx := context.Background() + authInfo := &ClientAuthInfo{AccessToken: "test"} + authErr := errors.New("test error") + + ctx = WithAuthInfo(ctx, authInfo) + ctx = WithAuthErr(ctx, authErr) + + // Both should be retrievable + storedAuthInfo, ok := GetAuthInfo(ctx) + if !ok || storedAuthInfo != authInfo { + t.Error("Expected to retrieve stored auth info") + } + + storedAuthErr := ctx.Value(ctxKeyClientAuthErr) + if storedAuthErr != authErr { + t.Error("Expected to retrieve stored auth error") + } +} + +func TestAuthInfoEdgeCases(t *testing.T) { + // Test with zero values + authInfo := &ClientAuthInfo{} + if authInfo.AccessToken != "" { + t.Error("Expected empty AccessToken by default") + } + if authInfo.RefreshToken != nil { + t.Error("Expected nil RefreshToken by default") + } + if authInfo.ExpiresAt != nil { + t.Error("Expected nil ExpiresAt by default") + } + if authInfo.Scopes != nil { + t.Error("Expected nil Scopes by default") + } + if authInfo.Extra != nil { + t.Error("Expected nil Extra by default") + } + + // Test IsTokenExpired with zero value auth info + if IsTokenExpired(authInfo) { + t.Error("Expected zero value auth info (no expiry) to not be expired") + } +} + +func TestTokenExpirationBoundaryConditions(t *testing.T) { + // Test exactly at 30 second boundary + exactBoundary := time.Now().Add(30 * time.Second) + authInfoBoundary := &ClientAuthInfo{ + AccessToken: "test-token", + ExpiresAt: &exactBoundary, + } + + // Due to the -30 second buffer, this should be considered expired + if !IsTokenExpired(authInfoBoundary) { + t.Error("Expected token expiring at exactly 30 seconds to be considered expired") + } + + // Test just before the boundary + justBefore := time.Now().Add(29 * time.Second) + authInfoJustBefore := &ClientAuthInfo{ + AccessToken: "test-token", + ExpiresAt: &justBefore, + } + if !IsTokenExpired(authInfoJustBefore) { + t.Error("Expected token expiring before 30 second buffer to be expired") + } + + // Test just after the boundary + justAfter := time.Now().Add(31 * time.Second) + authInfoJustAfter := &ClientAuthInfo{ + AccessToken: "test-token", + ExpiresAt: &justAfter, + } + if IsTokenExpired(authInfoJustAfter) { + t.Error("Expected token expiring after 30 second buffer to not be expired") + } +} + +func TestConvertTokensToAuthInfoEdgeCases(t *testing.T) { + // Test with tokens having only access token + minimalTokens := &auth.OAuthTokens{ + AccessToken: "minimal-token", + TokenType: "Bearer", + } + authInfo := ConvertTokensToAuthInfo(minimalTokens) + if authInfo == nil { + t.Fatal("Expected non-nil auth info for minimal valid tokens") + } + if authInfo.AccessToken != "minimal-token" { + t.Errorf("Expected AccessToken minimal-token, got %s", authInfo.AccessToken) + } + if authInfo.RefreshToken != nil { + t.Error("Expected nil RefreshToken when not provided") + } + if authInfo.ExpiresAt != nil { + t.Error("Expected nil ExpiresAt when ExpiresIn not provided") + } + + // Test with zero ExpiresIn + zeroExpiresIn := int64(0) + tokensZeroExpiry := &auth.OAuthTokens{ + AccessToken: "test-token", + ExpiresIn: &zeroExpiresIn, + } + authInfo = ConvertTokensToAuthInfo(tokensZeroExpiry) + if authInfo == nil { + t.Fatal("Expected non-nil auth info") + } + if authInfo.ExpiresAt == nil { + t.Fatal("Expected ExpiresAt to be set even with zero ExpiresIn") + } + // Should be approximately now (since ExpiresIn is 0) + if time.Since(*authInfo.ExpiresAt) > time.Second { + t.Error("Expected ExpiresAt to be approximately now when ExpiresIn is 0") + } + + // Test with negative ExpiresIn + negativeExpiresIn := int64(-3600) + tokensNegativeExpiry := &auth.OAuthTokens{ + AccessToken: "test-token", + ExpiresIn: &negativeExpiresIn, + } + authInfo = ConvertTokensToAuthInfo(tokensNegativeExpiry) + if authInfo == nil { + t.Fatal("Expected non-nil auth info") + } + if authInfo.ExpiresAt == nil { + t.Fatal("Expected ExpiresAt to be set even with negative ExpiresIn") + } + // Should be in the past + if authInfo.ExpiresAt.After(time.Now()) { + t.Error("Expected ExpiresAt to be in the past when ExpiresIn is negative") + } +} + +func TestConcurrentContextOperations(t *testing.T) { + // Test concurrent access to context operations + ctx := context.Background() + + // Set up initial context with auth info + authInfo := &ClientAuthInfo{ + AccessToken: "concurrent-test-token", + Scopes: []string{"read"}, + } + ctx = WithAuthInfo(ctx, authInfo) + + // Test concurrent reads + numGoroutines := 10 + results := make(chan bool, numGoroutines) + + for i := 0; i < numGoroutines; i++ { + go func() { + storedInfo, ok := GetAuthInfo(ctx) + results <- ok && storedInfo != nil && storedInfo.AccessToken == "concurrent-test-token" + }() + } + + // Collect results + for i := 0; i < numGoroutines; i++ { + if !<-results { + t.Error("Expected successful concurrent read of auth info") + } + } +} + +func TestAuthInfoDeepCopy(t *testing.T) { + // Test that modifications to returned auth info don't affect stored version + refreshToken := "original-refresh" + expiresAt := time.Now().Add(time.Hour) + originalAuthInfo := &ClientAuthInfo{ + AccessToken: "original-access", + RefreshToken: &refreshToken, + ExpiresAt: &expiresAt, + Scopes: []string{"read", "write"}, + Extra: map[string]interface{}{"key": "value"}, + } + + ctx := WithAuthInfo(context.Background(), originalAuthInfo) + retrievedAuthInfo, ok := GetAuthInfo(ctx) + if !ok { + t.Fatal("Expected to retrieve auth info") + } + + // Modify the retrieved auth info + retrievedAuthInfo.AccessToken = "modified-access" + *retrievedAuthInfo.RefreshToken = "modified-refresh" + retrievedAuthInfo.Scopes[0] = "modified" + retrievedAuthInfo.Extra["key"] = "modified" + + // Check that original is also modified (since we're returning the same pointer) + // This test documents the current behavior - no deep copy is performed + if originalAuthInfo.AccessToken != "modified-access" { + t.Error("Auth info appears to be deep copied (might be unexpected)") + } +} + +// Helper function to create string pointers +func stringPtr(s string) *string { + return &s +} + +// Benchmark tests +func BenchmarkWithAuthInfo(b *testing.B) { + ctx := context.Background() + authInfo := &ClientAuthInfo{ + AccessToken: "benchmark-token", + Scopes: []string{"read", "write"}, + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + WithAuthInfo(ctx, authInfo) + } +} + +func BenchmarkGetAuthInfo(b *testing.B) { + authInfo := &ClientAuthInfo{ + AccessToken: "benchmark-token", + Scopes: []string{"read", "write"}, + } + ctx := WithAuthInfo(context.Background(), authInfo) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + GetAuthInfo(ctx) + } +} + +func BenchmarkConvertTokensToAuthInfo(b *testing.B) { + refreshToken := "bench-refresh-token" + expiresIn := int64(3600) + tokens := &auth.OAuthTokens{ + AccessToken: "bench-access-token", + RefreshToken: &refreshToken, + TokenType: "Bearer", + ExpiresIn: &expiresIn, + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + ConvertTokensToAuthInfo(tokens) + } +} + +func BenchmarkIsTokenExpired(b *testing.B) { + expiresAt := time.Now().Add(time.Hour) + authInfo := &ClientAuthInfo{ + AccessToken: "bench-token", + ExpiresAt: &expiresAt, + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + IsTokenExpired(authInfo) + } +} diff --git a/internal/auth/client/memory_provider.go b/internal/auth/client/memory_provider.go new file mode 100644 index 0000000..8a7fd09 --- /dev/null +++ b/internal/auth/client/memory_provider.go @@ -0,0 +1,195 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package client + +import ( + "crypto/rand" + "encoding/base64" + "fmt" + "net/http" + "net/url" + "strings" + "sync" + + "trpc.group/trpc-go/trpc-mcp-go/internal/auth" +) + +// InMemoryOAuthClientProvider provides an in-memory implementation of OAuthClientProvider +// It stores client information, tokens, and PKCE code verifiers with thread safety +// +// NOTE: This is a demo implementation for testing and development purposes only. +// It is NOT recommended for production use as it stores sensitive data in memory +// without persistence or proper security measures. +type InMemoryOAuthClientProvider struct { + redirectURL string + clientMetadata auth.OAuthClientMetadata + clientInfo *auth.OAuthClientInformation + tokens *auth.OAuthTokens + codeVerifier string + onRedirect func(*url.URL) error + mutex sync.RWMutex +} + +// NewInMemoryOAuthClientProvider creates a new in-memory OAuth client provider +func NewInMemoryOAuthClientProvider( + redirectURL string, + clientMetadata auth.OAuthClientMetadata, + onRedirect func(*url.URL) error) *InMemoryOAuthClientProvider { + if onRedirect == nil { + onRedirect = func(u *url.URL) error { + return nil + } + } + return &InMemoryOAuthClientProvider{ + redirectURL: redirectURL, + clientMetadata: clientMetadata, + onRedirect: onRedirect, + } +} + +// RedirectURL returns the registered redirect URL +func (p *InMemoryOAuthClientProvider) RedirectURL() string { + return p.redirectURL +} + +// ClientMetadata returns the client metadata +func (p *InMemoryOAuthClientProvider) ClientMetadata() auth.OAuthClientMetadata { + return p.clientMetadata +} + +// ClientInformation returns stored client credentials if available +func (p *InMemoryOAuthClientProvider) ClientInformation() *auth.OAuthClientInformation { + p.mutex.RLock() + defer p.mutex.RUnlock() + return p.clientInfo +} + +// SaveClientInformation saves client credentials into memory +func (p *InMemoryOAuthClientProvider) SaveClientInformation(clientInformation auth.OAuthClientInformationFull) error { + p.mutex.Lock() + defer p.mutex.Unlock() + p.clientInfo = &auth.OAuthClientInformation{ + ClientID: clientInformation.ClientID, + ClientSecret: clientInformation.ClientSecret, + ClientIDIssuedAt: clientInformation.ClientIDIssuedAt, + ClientSecretExpiresAt: clientInformation.ClientSecretExpiresAt, + } + return nil +} + +// Tokens returns stored tokens if available +func (p *InMemoryOAuthClientProvider) Tokens() (*auth.OAuthTokens, error) { + p.mutex.RLock() + defer p.mutex.RUnlock() + return p.tokens, nil +} + +// SaveTokens saves OAuth tokens into memory +func (p *InMemoryOAuthClientProvider) SaveTokens(tokens auth.OAuthTokens) error { + p.mutex.Lock() + defer p.mutex.Unlock() + p.tokens = &tokens + return nil +} + +// RedirectToAuthorization executes the redirect callback with authorization URL +func (p *InMemoryOAuthClientProvider) RedirectToAuthorization(authorizationUrl *url.URL) error { + return p.onRedirect(authorizationUrl) +} + +// CodeVerifier retrieves the stored PKCE code verifier +func (p *InMemoryOAuthClientProvider) CodeVerifier() (string, error) { + p.mutex.RLock() + defer p.mutex.RUnlock() + if p.codeVerifier == "" { + return "", fmt.Errorf("no code verifier saved") + } + return p.codeVerifier, nil +} + +// SaveCodeVerifier saves the PKCE code verifier into memory +func (p *InMemoryOAuthClientProvider) SaveCodeVerifier(codeVerifier string) error { + p.mutex.Lock() + defer p.mutex.Unlock() + p.codeVerifier = codeVerifier + return nil +} + +// State generates a random state string for CSRF protection +func (p *InMemoryOAuthClientProvider) State() (string, error) { + // Generate a random state parameter for CSRF protection + bytes := make([]byte, 32) + if _, err := rand.Read(bytes); err != nil { + return "", fmt.Errorf("failed to generate random state: %w", err) + } + return base64.URLEncoding.EncodeToString(bytes), nil +} + +// AddClientAuthentication adds client authentication parameters using client_secret_post +func (p *InMemoryOAuthClientProvider) AddClientAuthentication(headers http.Header, params url.Values, tokenUrl string) error { + // Add client authentication using client_secret_post method + p.mutex.RLock() + clientInfo := p.clientInfo + p.mutex.RUnlock() + + if clientInfo != nil && clientInfo.ClientID != "" { + params.Set("client_id", clientInfo.ClientID) + if clientInfo.ClientSecret != "" { + params.Set("client_secret", clientInfo.ClientSecret) + } + } + return nil +} + +// ValidateResourceURL validates the resource URL from metadata against the server URL +func (p *InMemoryOAuthClientProvider) ValidateResourceURL(serverUrl *url.URL, resourceMetadata *auth.OAuthProtectedResourceMetadata) (*url.URL, error) { + // If no resource metadata provided, return nil (no resource parameter needed) + if resourceMetadata == nil { + return nil, nil + } + + // Parse the resource URL from metadata + resourceURL, err := url.Parse(resourceMetadata.Resource) + if err != nil { + return nil, fmt.Errorf("invalid resource URL in metadata: %w", err) + } + + // Basic validation: ensure the resource URL has the same origin as server URL + if resourceURL.Scheme != serverUrl.Scheme || resourceURL.Host != serverUrl.Host { + // Allow if resource URL is a more specific path under the same origin + if !strings.HasPrefix(resourceURL.String(), serverUrl.Scheme+"://"+serverUrl.Host) { + return nil, fmt.Errorf("resource URL %s does not match server origin %s://%s", + resourceURL.String(), serverUrl.Scheme, serverUrl.Host) + } + } + + return resourceURL, nil +} + +// InvalidateCredentials clears stored credentials based on scope +func (p *InMemoryOAuthClientProvider) InvalidateCredentials(scope string) error { + // Clear the corresponding credentials + // according to the scope and use a mutex to protect them + p.mutex.Lock() + defer p.mutex.Unlock() + + switch scope { + case "all": + p.clientInfo = nil + p.tokens = nil + p.codeVerifier = "" + case "client": + p.clientInfo = nil + case "tokens": + p.tokens = nil + case "verifier": + p.codeVerifier = "" + default: + return fmt.Errorf("unknown invalidation scope: %s", scope) + } + return nil +} diff --git a/internal/auth/client/memory_provider_test.go b/internal/auth/client/memory_provider_test.go new file mode 100644 index 0000000..9c173e5 --- /dev/null +++ b/internal/auth/client/memory_provider_test.go @@ -0,0 +1,704 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package client + +import ( + "fmt" + "net/http" + "net/url" + "strings" + "sync" + "testing" + "time" + + "trpc.group/trpc-go/trpc-mcp-go/internal/auth" +) + +func TestNewInMemoryOAuthClientProvider(t *testing.T) { + redirectURL := "http://localhost:8080/callback" + clientName := "Test Client" + scope := "read write" + clientMetadata := auth.OAuthClientMetadata{ + ClientName: &clientName, // Use pointer to string + Scope: &scope, // Use pointer to string + } + + // Test with onRedirect callback + var redirectCalled bool + onRedirect := func(u *url.URL) error { + redirectCalled = true + return nil + } + + provider := NewInMemoryOAuthClientProvider(redirectURL, clientMetadata, onRedirect) + + if provider.redirectURL != redirectURL { + t.Errorf("Expected redirectURL %s, got %s", redirectURL, provider.redirectURL) + } + + if *provider.clientMetadata.ClientName != *clientMetadata.ClientName { + t.Errorf("Expected clientMetadata.ClientName %s, got %s", + *clientMetadata.ClientName, *provider.clientMetadata.ClientName) + } + + if provider.onRedirect == nil { + t.Error("Expected onRedirect to be set") + } + + // Test that onRedirect was properly set by calling it + testURL, _ := url.Parse("https://example.com") + provider.onRedirect(testURL) + if !redirectCalled { + t.Error("Expected onRedirect to be called") + } + + // Test with nil onRedirect + provider2 := NewInMemoryOAuthClientProvider(redirectURL, clientMetadata, nil) + if provider2.onRedirect == nil { + t.Error("Expected default onRedirect to be set when nil is passed") + } + + // Test default onRedirect doesn't panic + err := provider2.onRedirect(&url.URL{}) + if err != nil { + t.Errorf("Expected default onRedirect to return nil, got %v", err) + } +} + +func TestRedirectURL(t *testing.T) { + redirectURL := "http://localhost:8080/callback" + provider := NewInMemoryOAuthClientProvider(redirectURL, auth.OAuthClientMetadata{}, nil) + + if provider.RedirectURL() != redirectURL { + t.Errorf("Expected %s, got %s", redirectURL, provider.RedirectURL()) + } +} + +func TestClientMetadata(t *testing.T) { + clientName := "Test Client" + scope := "read write" + clientMetadata := auth.OAuthClientMetadata{ + ClientName: &clientName, + Scope: &scope, + } + provider := NewInMemoryOAuthClientProvider("", clientMetadata, nil) + + result := provider.ClientMetadata() + if *result.ClientName != *clientMetadata.ClientName { + t.Errorf("Expected ClientName %s, got %s", *clientMetadata.ClientName, *result.ClientName) + } + if *result.Scope != *clientMetadata.Scope { + t.Errorf("Expected Scope %s, got %s", *clientMetadata.Scope, *result.Scope) + } +} + +func TestClientInformation(t *testing.T) { + provider := NewInMemoryOAuthClientProvider("", auth.OAuthClientMetadata{}, nil) + + // Test when no client information is saved + clientInfo := provider.ClientInformation() + if clientInfo != nil { + t.Error("Expected nil client information when none is saved") + } + + // Test after saving client information + fullClientInfo := auth.OAuthClientInformationFull{ + OAuthClientInformation: auth.OAuthClientInformation{ + ClientID: "test-client-id", + ClientSecret: "test-client-secret", + ClientIDIssuedAt: &[]int64{time.Now().Unix()}[0], + ClientSecretExpiresAt: &[]int64{time.Now().Add(time.Hour * 24).Unix()}[0], + }, + } + + err := provider.SaveClientInformation(fullClientInfo) + if err != nil { + t.Errorf("Expected no error saving client information, got %v", err) + } + + clientInfo = provider.ClientInformation() + if clientInfo == nil { + t.Fatal("Expected client information to be saved") + } + + if clientInfo.ClientID != fullClientInfo.ClientID { + t.Errorf("Expected ClientID %s, got %s", fullClientInfo.ClientID, clientInfo.ClientID) + } + + if clientInfo.ClientSecret != fullClientInfo.ClientSecret { + t.Errorf("Expected ClientSecret %s, got %s", fullClientInfo.ClientSecret, clientInfo.ClientSecret) + } +} + +func TestTokens(t *testing.T) { + provider := NewInMemoryOAuthClientProvider("", auth.OAuthClientMetadata{}, nil) + + // Test when no tokens are saved + tokens, err := provider.Tokens() + if err != nil { + t.Errorf("Expected no error getting tokens, got %v", err) + } + if tokens != nil { + t.Error("Expected nil tokens when none are saved") + } + + // Test saving and retrieving tokens + testTokens := auth.OAuthTokens{ + AccessToken: "test-access-token", + RefreshToken: &[]string{"test-refresh-token"}[0], // Use pointer + TokenType: "Bearer", + ExpiresIn: &[]int64{3600}[0], // Use pointer + } + + err = provider.SaveTokens(testTokens) + if err != nil { + t.Errorf("Expected no error saving tokens, got %v", err) + } + + tokens, err = provider.Tokens() + if err != nil { + t.Errorf("Expected no error getting tokens, got %v", err) + } + + if tokens == nil { + t.Fatal("Expected tokens to be saved") + } + + if tokens.AccessToken != testTokens.AccessToken { + t.Errorf("Expected AccessToken %s, got %s", testTokens.AccessToken, tokens.AccessToken) + } + + if *tokens.RefreshToken != *testTokens.RefreshToken { + t.Errorf("Expected RefreshToken %s, got %s", *testTokens.RefreshToken, *tokens.RefreshToken) + } +} + +func TestRedirectToAuthorization(t *testing.T) { + var capturedURL *url.URL + onRedirect := func(u *url.URL) error { + capturedURL = u + return nil + } + + provider := NewInMemoryOAuthClientProvider("", auth.OAuthClientMetadata{}, onRedirect) + + testURL, _ := url.Parse("https://auth.example.com/authorize?response_type=code&client_id=123") + err := provider.RedirectToAuthorization(testURL) + + if err != nil { + t.Errorf("Expected no error, got %v", err) + } + + if capturedURL == nil { + t.Error("Expected onRedirect to be called") + } + + if capturedURL.String() != testURL.String() { + t.Errorf("Expected URL %s, got %s", testURL.String(), capturedURL.String()) + } +} + +func TestCodeVerifier(t *testing.T) { + provider := NewInMemoryOAuthClientProvider("", auth.OAuthClientMetadata{}, nil) + + // Test when no code verifier is saved + verifier, err := provider.CodeVerifier() + if err == nil { + t.Error("Expected error when no code verifier is saved") + } + if verifier != "" { + t.Error("Expected empty verifier when none is saved") + } + + // Test saving and retrieving code verifier + testVerifier := "test-code-verifier-12345" + err = provider.SaveCodeVerifier(testVerifier) + if err != nil { + t.Errorf("Expected no error saving code verifier, got %v", err) + } + + verifier, err = provider.CodeVerifier() + if err != nil { + t.Errorf("Expected no error getting code verifier, got %v", err) + } + + if verifier != testVerifier { + t.Errorf("Expected verifier %s, got %s", testVerifier, verifier) + } +} + +func TestState(t *testing.T) { + provider := NewInMemoryOAuthClientProvider("", auth.OAuthClientMetadata{}, nil) + + state1, err := provider.State() + if err != nil { + t.Errorf("Expected no error generating state, got %v", err) + } + + if state1 == "" { + t.Error("Expected non-empty state") + } + + // Test that subsequent calls generate different states + state2, err := provider.State() + if err != nil { + t.Errorf("Expected no error generating state, got %v", err) + } + + if state1 == state2 { + t.Error("Expected different states on subsequent calls") + } + + /// Test that state is base64 URL encoded (no + or / characters, = padding is allowed) + if strings.Contains(state1, "+") || strings.Contains(state1, "/") { + t.Error("Expected URL-safe base64 encoding (no + or / characters)") + } +} + +func TestAddClientAuthentication(t *testing.T) { + provider := NewInMemoryOAuthClientProvider("", auth.OAuthClientMetadata{}, nil) + + // Test with no client information + headers := make(http.Header) + params := make(url.Values) + err := provider.AddClientAuthentication(headers, params, "https://token.example.com") + if err != nil { + t.Errorf("Expected no error with no client info, got %v", err) + } + + // Should not add any parameters when no client info + if params.Get("client_id") != "" { + t.Error("Expected no client_id when no client info is saved") + } + + // Test with client information + clientInfo := auth.OAuthClientInformationFull{ + OAuthClientInformation: auth.OAuthClientInformation{ + ClientID: "test-client-id", + ClientSecret: "test-client-secret", + }, + } + err = provider.SaveClientInformation(clientInfo) + if err != nil { + t.Errorf("Expected no error saving client info, got %v", err) + } + + params = make(url.Values) + err = provider.AddClientAuthentication(headers, params, "https://token.example.com") + if err != nil { + t.Errorf("Expected no error with client info, got %v", err) + } + + if params.Get("client_id") != clientInfo.ClientID { + t.Errorf("Expected client_id %s, got %s", clientInfo.ClientID, params.Get("client_id")) + } + + if params.Get("client_secret") != clientInfo.ClientSecret { + t.Errorf("Expected client_secret %s, got %s", clientInfo.ClientSecret, params.Get("client_secret")) + } + + // Test with only client ID (no secret) + clientInfoNoSecret := auth.OAuthClientInformationFull{ + OAuthClientInformation: auth.OAuthClientInformation{ + ClientID: "test-client-id-only", + }, + } + err = provider.SaveClientInformation(clientInfoNoSecret) + if err != nil { + t.Errorf("Expected no error saving client info, got %v", err) + } + + params = make(url.Values) + err = provider.AddClientAuthentication(headers, params, "https://token.example.com") + if err != nil { + t.Errorf("Expected no error with client info, got %v", err) + } + + if params.Get("client_id") != clientInfoNoSecret.ClientID { + t.Errorf("Expected client_id %s, got %s", clientInfoNoSecret.ClientID, params.Get("client_id")) + } + + if params.Get("client_secret") != "" { + t.Error("Expected no client_secret when not provided") + } +} + +func TestValidateResourceURL(t *testing.T) { + provider := NewInMemoryOAuthClientProvider("", auth.OAuthClientMetadata{}, nil) + serverURL, _ := url.Parse("https://api.example.com") + + // Test with nil resource metadata + result, err := provider.ValidateResourceURL(serverURL, nil) + if err != nil { + t.Errorf("Expected no error with nil metadata, got %v", err) + } + if result != nil { + t.Error("Expected nil result with nil metadata") + } + + // Test with valid resource URL (same origin) + resourceMetadata := &auth.OAuthProtectedResourceMetadata{ + Resource: "https://api.example.com/data", + } + + result, err = provider.ValidateResourceURL(serverURL, resourceMetadata) + if err != nil { + t.Errorf("Expected no error with valid resource URL, got %v", err) + } + if result == nil { + t.Fatal("Expected non-nil result") + } + if result.String() != resourceMetadata.Resource { + t.Errorf("Expected resource URL %s, got %s", resourceMetadata.Resource, result.String()) + } + + // Test with valid resource URL (subpath) + resourceMetadata2 := &auth.OAuthProtectedResourceMetadata{ + Resource: "https://api.example.com/v1/users", + } + + result, err = provider.ValidateResourceURL(serverURL, resourceMetadata2) + if err != nil { + t.Errorf("Expected no error with valid subpath resource URL, got %v", err) + } + if result == nil { + t.Fatal("Expected non-nil result") + } + + // Test with different scheme + resourceMetadata3 := &auth.OAuthProtectedResourceMetadata{ + Resource: "http://api.example.com/data", + } + + result, err = provider.ValidateResourceURL(serverURL, resourceMetadata3) + if err == nil { + t.Error("Expected error with different scheme") + } + if result != nil { + t.Error("Expected nil result with invalid URL") + } + + // Test with different host + resourceMetadata4 := &auth.OAuthProtectedResourceMetadata{ + Resource: "https://other.example.com/data", + } + + result, err = provider.ValidateResourceURL(serverURL, resourceMetadata4) + if err == nil { + t.Error("Expected error with different host") + } + if result != nil { + t.Error("Expected nil result with invalid URL") + } + + // Test with invalid URL + resourceMetadata5 := &auth.OAuthProtectedResourceMetadata{ + Resource: "://invalid-url", + } + + result, err = provider.ValidateResourceURL(serverURL, resourceMetadata5) + if err == nil { + t.Error("Expected error with invalid URL") + } + if result != nil { + t.Error("Expected nil result with invalid URL") + } +} + +func TestInvalidateCredentials(t *testing.T) { + provider := NewInMemoryOAuthClientProvider("", auth.OAuthClientMetadata{}, nil) + + // Set up some data first + clientInfo := auth.OAuthClientInformationFull{ + OAuthClientInformation: auth.OAuthClientInformation{ + ClientID: "test-client", + ClientSecret: "test-secret", + }, + } + provider.SaveClientInformation(clientInfo) + + tokens := auth.OAuthTokens{ + AccessToken: "test-access-token", + } + provider.SaveTokens(tokens) + + provider.SaveCodeVerifier("test-verifier") + + // Test invalidating all credentials + err := provider.InvalidateCredentials("all") + if err != nil { + t.Errorf("Expected no error invalidating all credentials, got %v", err) + } + + if provider.ClientInformation() != nil { + t.Error("Expected client information to be cleared") + } + + savedTokens, _ := provider.Tokens() + if savedTokens != nil { + t.Error("Expected tokens to be cleared") + } + + _, err = provider.CodeVerifier() + if err == nil { + t.Error("Expected error getting code verifier after clearing") + } + + // Set up data again for individual tests + provider.SaveClientInformation(clientInfo) + provider.SaveTokens(tokens) + provider.SaveCodeVerifier("test-verifier") + + // Test invalidating only client info + err = provider.InvalidateCredentials("client") + if err != nil { + t.Errorf("Expected no error invalidating client credentials, got %v", err) + } + + if provider.ClientInformation() != nil { + t.Error("Expected client information to be cleared") + } + + savedTokens, _ = provider.Tokens() + if savedTokens == nil { + t.Error("Expected tokens to remain") + } + + // Test invalidating only tokens + provider.SaveClientInformation(clientInfo) // Restore client info + err = provider.InvalidateCredentials("tokens") + if err != nil { + t.Errorf("Expected no error invalidating tokens, got %v", err) + } + + savedTokens, _ = provider.Tokens() + if savedTokens != nil { + t.Error("Expected tokens to be cleared") + } + + if provider.ClientInformation() == nil { + t.Error("Expected client information to remain") + } + + // Test invalidating only code verifier + provider.SaveTokens(tokens) // Restore tokens + provider.SaveCodeVerifier("test-verifier") // Restore verifier + + err = provider.InvalidateCredentials("verifier") + if err != nil { + t.Errorf("Expected no error invalidating verifier, got %v", err) + } + + _, err = provider.CodeVerifier() + if err == nil { + t.Error("Expected error getting code verifier after clearing") + } + + // Other data should remain + if provider.ClientInformation() == nil { + t.Error("Expected client information to remain") + } + + savedTokens, _ = provider.Tokens() + if savedTokens == nil { + t.Error("Expected tokens to remain") + } + + // Test invalid scope + err = provider.InvalidateCredentials("invalid-scope") + if err == nil { + t.Error("Expected error with invalid scope") + } + + if !strings.Contains(err.Error(), "unknown invalidation scope") { + t.Errorf("Expected 'unknown invalidation scope' error, got %v", err) + } +} + +func TestConcurrentAccess(t *testing.T) { + provider := NewInMemoryOAuthClientProvider("", auth.OAuthClientMetadata{}, nil) + + // Test concurrent access to ensure thread safety + var wg sync.WaitGroup + numGoroutines := 10 + + // Test concurrent writes + wg.Add(numGoroutines) + for i := 0; i < numGoroutines; i++ { + go func(id int) { + defer wg.Done() + + clientInfo := auth.OAuthClientInformationFull{ + OAuthClientInformation: auth.OAuthClientInformation{ + ClientID: fmt.Sprintf("client-%d", id), + ClientSecret: fmt.Sprintf("secret-%d", id), + }, + } + provider.SaveClientInformation(clientInfo) + + tokens := auth.OAuthTokens{ + AccessToken: fmt.Sprintf("token-%d", id), + } + provider.SaveTokens(tokens) + + provider.SaveCodeVerifier(fmt.Sprintf("verifier-%d", id)) + }(i) + } + + wg.Wait() + + // Test concurrent reads + wg.Add(numGoroutines) + for i := 0; i < numGoroutines; i++ { + go func() { + defer wg.Done() + + provider.ClientInformation() + provider.Tokens() + provider.CodeVerifier() + }() + } + + wg.Wait() + + // Test concurrent invalidation + wg.Add(numGoroutines) + for i := 0; i < numGoroutines; i++ { + go func(id int) { + defer wg.Done() + + scopes := []string{"client", "tokens", "verifier", "all"} + scope := scopes[id%len(scopes)] + provider.InvalidateCredentials(scope) + }(i) + } + + wg.Wait() +} + +func TestStateRandomness(t *testing.T) { + provider := NewInMemoryOAuthClientProvider("", auth.OAuthClientMetadata{}, nil) + + // Generate multiple states and ensure they're different + states := make(map[string]bool) + for i := 0; i < 10; i++ { + state, err := provider.State() + if err != nil { + t.Errorf("Expected no error generating state, got %v", err) + } + + if states[state] { + t.Errorf("Generated duplicate state: %s", state) + } + states[state] = true + + // Verify length (32 bytes base64 encoded should be ~43 characters) + if len(state) < 40 { + t.Errorf("Expected state length >= 40, got %d", len(state)) + } + } +} + +func TestRedirectToAuthorizationError(t *testing.T) { + expectedErr := fmt.Errorf("redirect failed") + onRedirect := func(u *url.URL) error { + return expectedErr + } + + provider := NewInMemoryOAuthClientProvider("", auth.OAuthClientMetadata{}, onRedirect) + + testURL, _ := url.Parse("https://auth.example.com/authorize") + err := provider.RedirectToAuthorization(testURL) + + if err != expectedErr { + t.Errorf("Expected error %v, got %v", expectedErr, err) + } +} + +func TestMutexProtection(t *testing.T) { + provider := NewInMemoryOAuthClientProvider("", auth.OAuthClientMetadata{}, nil) + + // Test that mutex properly protects against race conditions + // This test runs multiple operations concurrently and ensures no data races + var wg sync.WaitGroup + numOperations := 100 + + wg.Add(numOperations) + for i := 0; i < numOperations; i++ { + go func(id int) { + defer wg.Done() + + // Mix of read and write operations + switch id % 4 { + case 0: + clientInfo := auth.OAuthClientInformationFull{ + OAuthClientInformation: auth.OAuthClientInformation{ + ClientID: fmt.Sprintf("client-%d", id), + }, + } + provider.SaveClientInformation(clientInfo) + case 1: + provider.ClientInformation() + case 2: + tokens := auth.OAuthTokens{ + AccessToken: fmt.Sprintf("token-%d", id), + } + provider.SaveTokens(tokens) + case 3: + provider.Tokens() + } + }(i) + } + + wg.Wait() + + // If we reach here without data races, the mutex protection is working +} + +// Benchmark tests +func BenchmarkState(b *testing.B) { + provider := NewInMemoryOAuthClientProvider("", auth.OAuthClientMetadata{}, nil) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + _, err := provider.State() + if err != nil { + b.Fatalf("Error generating state: %v", err) + } + } +} + +func BenchmarkConcurrentAccess(b *testing.B) { + provider := NewInMemoryOAuthClientProvider("", auth.OAuthClientMetadata{}, nil) + + // Setup initial data + clientName := "bench-client" + clientInfo := auth.OAuthClientInformationFull{ + OAuthClientMetadata: auth.OAuthClientMetadata{ + ClientName: &clientName, + }, + OAuthClientInformation: auth.OAuthClientInformation{ + ClientID: "bench-client-id", + ClientSecret: "bench-secret", + }, + } + provider.SaveClientInformation(clientInfo) + + tokens := auth.OAuthTokens{ + AccessToken: "bench-token", + } + provider.SaveTokens(tokens) + + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + // Mix of read operations + provider.ClientInformation() + provider.Tokens() + } + }) +} diff --git a/internal/auth/client/provider.go b/internal/auth/client/provider.go new file mode 100644 index 0000000..842c41e --- /dev/null +++ b/internal/auth/client/provider.go @@ -0,0 +1,67 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package client + +import ( + "net/http" + "net/url" + + "trpc.group/trpc-go/trpc-mcp-go/internal/auth" +) + +// OAuthClientProvider defines core OAuth 2.0 client operations +// Provides client configuration, token management, and authorization flow handling +type OAuthClientProvider interface { + // RedirectURL returns the client redirect URL for authorization + RedirectURL() string + + // ClientMetadata returns static client metadata such as redirect URIs + ClientMetadata() auth.OAuthClientMetadata + + // ClientInformation returns current client credentials if available + ClientInformation() *auth.OAuthClientInformation + + // Tokens returns the current access and refresh tokens + Tokens() (*auth.OAuthTokens, error) + + // SaveTokens persists the given OAuth tokens + SaveTokens(tokens auth.OAuthTokens) error + + // RedirectToAuthorization handles redirection to the authorization endpoint + RedirectToAuthorization(authorizationUrl *url.URL) error + + // SaveCodeVerifier persists the PKCE code verifier for later token exchange + SaveCodeVerifier(codeVerifier string) error + + // CodeVerifier retrieves the stored PKCE code verifier + CodeVerifier() (string, error) +} + +// OAuthStateProvider adds state parameter management for CSRF protection. +type OAuthStateProvider interface { + State() (string, error) +} + +// OAuthClientInfoProvider handles dynamic client credential storage. +type OAuthClientInfoProvider interface { + SaveClientInformation(clientInformation auth.OAuthClientInformationFull) error +} + +// OAuthClientAuthProvider enables custom client authentication methods. +type OAuthClientAuthProvider interface { + AddClientAuthentication(headers http.Header, params url.Values, tokenUrl string) error +} + +// OAuthResourceValidator validates resource URLs for specific server requirements. +type OAuthResourceValidator interface { + ValidateResourceURL(serverUrl *url.URL, resourceMetadata *auth.OAuthProtectedResourceMetadata) (*url.URL, error) +} + +// OAuthCredentialInvalidator handles logout and credential revocation. +type OAuthCredentialInvalidator interface { + InvalidateCredentials(scope string) error +} diff --git a/internal/auth/pkce/utils.go b/internal/auth/pkce/utils.go new file mode 100644 index 0000000..ebe5b70 --- /dev/null +++ b/internal/auth/pkce/utils.go @@ -0,0 +1,107 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package pkce + +import ( + "crypto/rand" + "crypto/sha256" + "encoding/base64" + "fmt" + "regexp" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth/server" +) + +// PKCEChallenge holds PKCE code verifier and challenge +type PKCEChallenge struct { + // CodeVerifier is the high-entropy cryptographic random string + CodeVerifier string + // CodeChallenge is the derived challenge from the code verifier + CodeChallenge string +} + +// ValidatePKCEParams validates the PKCE parameters provided in the authorization request. +func ValidatePKCEParams(params server.AuthorizationParams) error { + if params.CodeChallenge == "" { + return fmt.Errorf("code_challenge is required") + } + + // Verify code_challenge length (RFC 7636: 43-128 characters) + if len(params.CodeChallenge) < 43 || len(params.CodeChallenge) > 128 { + return fmt.Errorf("code_challenge length must be between 43 and 128 characters") + } + + // Verify code_challenge format (BASE64URL) + if !isValidBase64URL(params.CodeChallenge) { + return fmt.Errorf("code_challenge must be valid BASE64URL") + } + + return nil +} + +// isValidBase64URL checks whether the given string is a valid Base64URL-encoded value +// and decodes to exactly 32 bytes (the output size of SHA-256). +func isValidBase64URL(s string) bool { + // Length check + if len(s) < 43 || len(s) > 128 { + return false + } + + // Character set validation + base64URLPattern := `^[A-Za-z0-9_-]+$` + matched, err := regexp.MatchString(base64URLPattern, s) + if err != nil || !matched { + return false + } + + // Try decoding verification + decoded, err := base64.RawURLEncoding.DecodeString(s) + if err != nil { + return false + } + + // For code_challenge, it should be 32 bytes after decoding (SHA256 hash) + if len(decoded) != 32 { + return false + } + + return true +} + +// VerifyPKCEChallenge verifies the PKCE code_verifier against the code_challenge +func VerifyPKCEChallenge(codeVerifier, codeChallenge string) bool { + if codeVerifier == "" || codeChallenge == "" { + return false + } + + // Create SHA256 hash of the code_verifier + hash := sha256.Sum256([]byte(codeVerifier)) + + // Base64 URL encode the hash + computedChallenge := base64.RawURLEncoding.EncodeToString(hash[:]) + + return computedChallenge == codeChallenge +} + +// GeneratePKCEChallenge generates a new PKCE pair (code_verifier and code_challenge). +func GeneratePKCEChallenge() (*PKCEChallenge, error) { + // Generate 43-128 character code_verifier (RFC 7636) + verifierBytes := make([]byte, 32) // 32 bytes = 43 chars in base64url + if _, err := rand.Read(verifierBytes); err != nil { + return nil, fmt.Errorf("failed to generate code verifier: %w", err) + } + + codeVerifier := base64.RawURLEncoding.EncodeToString(verifierBytes) + + // Generate code_challenge using S256 method + hash := sha256.Sum256([]byte(codeVerifier)) + codeChallenge := base64.RawURLEncoding.EncodeToString(hash[:]) + + return &PKCEChallenge{ + CodeVerifier: codeVerifier, + CodeChallenge: codeChallenge, + }, nil +} diff --git a/internal/auth/pkce/utils_test.go b/internal/auth/pkce/utils_test.go new file mode 100644 index 0000000..17cd99d --- /dev/null +++ b/internal/auth/pkce/utils_test.go @@ -0,0 +1,131 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package pkce + +import ( + "crypto/rand" + "crypto/sha256" + "encoding/base64" + "testing" +) + +// genValidBase64URLDigest generates a valid base64url-encoded SHA256 digest +func genValidBase64URLDigest(t *testing.T) string { + t.Helper() + buf := make([]byte, 32) + if _, err := rand.Read(buf); err != nil { + t.Fatalf("failed to read random: %v", err) + } + return base64.RawURLEncoding.EncodeToString(buf) +} + +func TestIsValidBase64URL_Valid(t *testing.T) { + s := genValidBase64URLDigest(t) + if !isValidBase64URL(s) { + t.Fatalf("expected valid base64url digest, got invalid: %q", s) + } +} + +func TestIsValidBase64URL_InvalidChars(t *testing.T) { + // '+' and '/' are invalid in base64url + s := "abcd+efg_hijklmnopqrstuvwxyz0123456789abcd" + if isValidBase64URL(s) { + t.Fatalf("expected invalid due to characters, got valid: %q", s) + } +} + +func TestIsValidBase64URL_WrongDecodedLen(t *testing.T) { + // 31 bytes decoded -> invalid + decoded31 := make([]byte, 31) + s31 := base64.RawURLEncoding.EncodeToString(decoded31) + if isValidBase64URL(s31) { + t.Fatalf("expected invalid due to decoded len != 32, got valid") + } + + // 33 bytes decoded -> invalid + decoded33 := make([]byte, 33) + s33 := base64.RawURLEncoding.EncodeToString(decoded33) + if isValidBase64URL(s33) { + t.Fatalf("expected invalid due to decoded len != 32, got valid") + } +} + +func TestVerifyPKCEChallenge_Match(t *testing.T) { + verifierBytes := make([]byte, 32) + if _, err := rand.Read(verifierBytes); err != nil { + t.Fatalf("rand read: %v", err) + } + codeVerifier := base64.RawURLEncoding.EncodeToString(verifierBytes) + + sum := sha256.Sum256([]byte(codeVerifier)) + expectedChallenge := base64.RawURLEncoding.EncodeToString(sum[:]) + + if !VerifyPKCEChallenge(codeVerifier, expectedChallenge) { + t.Fatalf("expected challenge to verify") + } +} + +func TestVerifyPKCEChallenge_Mismatch(t *testing.T) { + verifierBytes := make([]byte, 32) + if _, err := rand.Read(verifierBytes); err != nil { + t.Fatalf("rand read: %v", err) + } + codeVerifier := base64.RawURLEncoding.EncodeToString(verifierBytes) + + sum := sha256.Sum256([]byte(codeVerifier)) + expectedChallenge := base64.RawURLEncoding.EncodeToString(sum[:]) + + // tamper the verifier + codeVerifier += "A" + + if VerifyPKCEChallenge(codeVerifier, expectedChallenge) { + t.Fatalf("expected verification to fail for mismatched verifier") + } +} + +func TestGeneratePKCEChallenge(t *testing.T) { + pair, err := GeneratePKCEChallenge() + if err != nil { + t.Fatalf("GeneratePKCEChallenge returned error: %v", err) + } + if pair == nil { + t.Fatalf("expected non-nil pair") + } + if pair.CodeVerifier == "" || pair.CodeChallenge == "" { + t.Fatalf("expected non-empty verifier and challenge") + } + + // RFC 7636 requires 43..128 characters for the verifier + if l := len(pair.CodeVerifier); l < 43 || l > 128 { + t.Fatalf("code_verifier length must be in [43,128], got %d", l) + } + + // Challenge should be a base64url-encoded SHA256 digest (valid and 32 bytes decoded) + if !isValidBase64URL(pair.CodeChallenge) { + t.Fatalf("code_challenge should be valid base64url digest") + } + + // Verify Challenge == S256(verifier) + sum := sha256.Sum256([]byte(pair.CodeVerifier)) + expectedChallenge := base64.RawURLEncoding.EncodeToString(sum[:]) + if pair.CodeChallenge != expectedChallenge { + t.Fatalf("code_challenge mismatch: got %q want %q", pair.CodeChallenge, expectedChallenge) + } + + // And VerifyPKCEChallenge should pass + if !VerifyPKCEChallenge(pair.CodeVerifier, pair.CodeChallenge) { + t.Fatalf("VerifyPKCEChallenge should return true for generated pair") + } +} + +func BenchmarkGeneratePKCEChallenge(b *testing.B) { + for i := 0; i < b.N; i++ { + if _, err := GeneratePKCEChallenge(); err != nil { + b.Fatalf("GeneratePKCEChallenge error: %v", err) + } + } +} diff --git a/internal/auth/server/clients.go b/internal/auth/server/clients.go new file mode 100644 index 0000000..85ebc11 --- /dev/null +++ b/internal/auth/server/clients.go @@ -0,0 +1,81 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package server + +import ( + "fmt" + + "trpc.group/trpc-go/trpc-mcp-go/internal/auth" +) + +// OAuthClientsStoreInterface defines retrieval and optional dynamic registration for OAuth clients registered with this server +type OAuthClientsStoreInterface interface { + // GetClient returns information about a registered client by its ID or nil if not found + GetClient(clientId string) (*auth.OAuthClientInformationFull, error) + + // SupportDynamicClientRegistration adds optional dynamic client registration capability + // Implementations may return a modified client to reflect server enforced values + // Implementations should not delete expired client secrets in place + // Middleware validates client_secret_expires_at and rejects expired secrets + SupportDynamicClientRegistration +} + +// SupportDynamicClientRegistration exposes the RegisterClient operation for dynamic client registration +type SupportDynamicClientRegistration interface { + // RegisterClient registers a new OAuth client and returns the stored client record + RegisterClient(client auth.OAuthClientInformationFull) (*auth.OAuthClientInformationFull, error) +} + +// OAuthClientsStore is a functional store adapter for OAuth client retrieval and optional registration +type OAuthClientsStore struct { + getClient func(clientID string) (*auth.OAuthClientInformationFull, error) // lookup function injected by caller + registerClient func(client auth.OAuthClientInformationFull) (*auth.OAuthClientInformationFull, error) // optional registration function injected by caller +} + +// GetClient returns the client record for the given clientID or an error from the underlying store +func (s OAuthClientsStore) GetClient(clientID string) (*auth.OAuthClientInformationFull, error) { + // Delegate to injected lookup function + return s.getClient(clientID) +} + +// RegisterClient registers a new client if dynamic registration is supported otherwise returns an error +func (s OAuthClientsStore) RegisterClient(client auth.OAuthClientInformationFull) (*auth.OAuthClientInformationFull, error) { + // If no registration function is provided dynamic registration is not supported + if s.registerClient == nil { + return nil, fmt.Errorf("dynamic client registration is not supported") + } + // Delegate to injected registration function + return s.registerClient(client) +} + +// NewOAuthClientStoreSupportDynamicRegistration constructs a store with both lookup and registration support +func NewOAuthClientStoreSupportDynamicRegistration( + getClient func(clientID string) (*auth.OAuthClientInformationFull, error), + registerClient func(client auth.OAuthClientInformationFull) (*auth.OAuthClientInformationFull, error), +) *OAuthClientsStore { + // Inject both handlers to enable dynamic registration support + return &OAuthClientsStore{ + getClient: getClient, + registerClient: registerClient, + } +} + +// NewOAuthClientStore constructs a store that supports only client lookup +func NewOAuthClientStore( + getClient func(clientID string) (*auth.OAuthClientInformationFull, error), +) *OAuthClientsStore { + // Inject lookup handler only leaving registration unsupported + return &OAuthClientsStore{ + getClient: getClient, + } +} + +// SupportsRegistration returns true if dynamic client registration is supported +func (s OAuthClientsStore) SupportsRegistration() bool { + // Registration supported when a registration function is present + return s.registerClient != nil +} diff --git a/internal/auth/server/clients_test.go b/internal/auth/server/clients_test.go new file mode 100644 index 0000000..6ab5bf5 --- /dev/null +++ b/internal/auth/server/clients_test.go @@ -0,0 +1,323 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package server + +import ( + "errors" + "testing" + + "trpc.group/trpc-go/trpc-mcp-go/internal/auth" +) + +// Test errors used by mocks +var ( + ErrClientNotFound = errors.New("client not found") + ErrGetClientFailed = errors.New("get client failed") + ErrRegisterClientFailed = errors.New("register client failed") +) + +// mockGetClientSuccess returns a known client for id existing-client or ErrClientNotFound +func mockGetClientSuccess(clientID string) (*auth.OAuthClientInformationFull, error) { + if clientID == "existing-client" { + return &auth.OAuthClientInformationFull{ + OAuthClientInformation: auth.OAuthClientInformation{ + ClientID: clientID, + ClientSecret: "secret-123", + }, + OAuthClientMetadata: auth.OAuthClientMetadata{ + RedirectURIs: []string{"https://example.com/callback"}, + }, + }, nil + } + return nil, ErrClientNotFound +} + +// mockGetClientError always returns ErrGetClientFailed +func mockGetClientError(clientID string) (*auth.OAuthClientInformationFull, error) { + return nil, ErrGetClientFailed +} + +// mockRegisterClientSuccess simulates server generating client_id and client_secret +func mockRegisterClientSuccess(client auth.OAuthClientInformationFull) (*auth.OAuthClientInformationFull, error) { + client.ClientID = "generated-client-id" + client.ClientSecret = "generated-secret" + return &client, nil +} + +// mockRegisterClientError always returns ErrRegisterClientFailed +func mockRegisterClientError(client auth.OAuthClientInformationFull) (*auth.OAuthClientInformationFull, error) { + return nil, ErrRegisterClientFailed +} + +func TestNewOAuthClientStore(t *testing.T) { + store := NewOAuthClientStore(mockGetClientSuccess) + + if store == nil { + t.Fatal("NewOAuthClientStore returned nil") + } + + // Should not support registration + if store.SupportsRegistration() { + t.Error("basic store should not support registration") + } + + // GetClient should work + client, err := store.GetClient("existing-client") + if err != nil { + t.Fatalf("GetClient failed: %v", err) + } + if client.ClientID != "existing-client" { + t.Errorf("expected client ID 'existing-client', got %s", client.ClientID) + } + + // RegisterClient should return error + newClient := auth.OAuthClientInformationFull{ + OAuthClientMetadata: auth.OAuthClientMetadata{ + RedirectURIs: []string{"https://test.com/callback"}, + }, + } + _, err = store.RegisterClient(newClient) + if err == nil { + t.Error("RegisterClient should fail on basic store") + } + if err.Error() != "dynamic client registration is not supported" { + t.Errorf("unexpected error message: %s", err.Error()) + } +} + +func TestNewOAuthClientStoreSupportDynamicRegistration(t *testing.T) { + store := NewOAuthClientStoreSupportDynamicRegistration(mockGetClientSuccess, mockRegisterClientSuccess) + + if store == nil { + t.Fatal("NewOAuthClientStoreSupportDynamicRegistration returned nil") + } + + // Should support registration + if !store.SupportsRegistration() { + t.Error("dynamic registration store should support registration") + } + + // GetClient should work + client, err := store.GetClient("existing-client") + if err != nil { + t.Fatalf("GetClient failed: %v", err) + } + if client.ClientID != "existing-client" { + t.Errorf("expected client ID 'existing-client', got %s", client.ClientID) + } + + // RegisterClient should work + newClient := auth.OAuthClientInformationFull{ + OAuthClientMetadata: auth.OAuthClientMetadata{ + RedirectURIs: []string{"https://test.com/callback"}, + }, + } + registered, err := store.RegisterClient(newClient) + if err != nil { + t.Fatalf("RegisterClient failed: %v", err) + } + if registered.ClientID != "generated-client-id" { + t.Errorf("expected generated client ID, got %s", registered.ClientID) + } + if registered.ClientSecret != "generated-secret" { + t.Errorf("expected generated secret, got %s", registered.ClientSecret) + } +} + +func TestOAuthClientsStore_GetClient_Success(t *testing.T) { + store := NewOAuthClientStore(mockGetClientSuccess) + + // Test existing client + client, err := store.GetClient("existing-client") + if err != nil { + t.Fatalf("GetClient failed: %v", err) + } + if client.ClientID != "existing-client" { + t.Errorf("expected client ID 'existing-client', got %s", client.ClientID) + } + if client.ClientSecret != "secret-123" { + t.Errorf("expected secret 'secret-123', got %s", client.ClientSecret) + } + if len(client.RedirectURIs) != 1 || client.RedirectURIs[0] != "https://example.com/callback" { + t.Errorf("unexpected redirect URIs: %v", client.RedirectURIs) + } +} + +func TestOAuthClientsStore_GetClient_NotFound(t *testing.T) { + store := NewOAuthClientStore(mockGetClientSuccess) + + // Test non-existing client + client, err := store.GetClient("non-existing-client") + if err != ErrClientNotFound { + t.Errorf("expected ErrClientNotFound, got %v", err) + } + if client != nil { + t.Error("expected nil client for not found case") + } +} + +func TestOAuthClientsStore_GetClient_Error(t *testing.T) { + store := NewOAuthClientStore(mockGetClientError) + + // Test error case + client, err := store.GetClient("any-client") + if err != ErrGetClientFailed { + t.Errorf("expected ErrGetClientFailed, got %v", err) + } + if client != nil { + t.Error("expected nil client for error case") + } +} + +func TestOAuthClientsStore_RegisterClient_Supported(t *testing.T) { + store := NewOAuthClientStoreSupportDynamicRegistration(mockGetClientSuccess, mockRegisterClientSuccess) + + newClient := auth.OAuthClientInformationFull{ + OAuthClientMetadata: auth.OAuthClientMetadata{ + RedirectURIs: []string{"https://test.com/callback"}, + ClientName: stringPtr("Test Client"), + }, + } + + registered, err := store.RegisterClient(newClient) + if err != nil { + t.Fatalf("RegisterClient failed: %v", err) + } + + if registered.ClientID != "generated-client-id" { + t.Errorf("expected generated client ID, got %s", registered.ClientID) + } + if registered.ClientSecret != "generated-secret" { + t.Errorf("expected generated secret, got %s", registered.ClientSecret) + } + // Original metadata should be preserved + if registered.ClientName == nil || *registered.ClientName != "Test Client" { + t.Errorf("client name not preserved") + } +} + +func TestOAuthClientsStore_RegisterClient_NotSupported(t *testing.T) { + store := NewOAuthClientStore(mockGetClientSuccess) + + newClient := auth.OAuthClientInformationFull{ + OAuthClientMetadata: auth.OAuthClientMetadata{ + RedirectURIs: []string{"https://test.com/callback"}, + }, + } + + registered, err := store.RegisterClient(newClient) + if err == nil { + t.Error("RegisterClient should fail when not supported") + } + if err.Error() != "dynamic client registration is not supported" { + t.Errorf("unexpected error message: %s", err.Error()) + } + if registered != nil { + t.Error("expected nil result for unsupported registration") + } +} + +func TestOAuthClientsStore_RegisterClient_Error(t *testing.T) { + store := NewOAuthClientStoreSupportDynamicRegistration(mockGetClientSuccess, mockRegisterClientError) + + newClient := auth.OAuthClientInformationFull{ + OAuthClientMetadata: auth.OAuthClientMetadata{ + RedirectURIs: []string{"https://test.com/callback"}, + }, + } + + registered, err := store.RegisterClient(newClient) + if err != ErrRegisterClientFailed { + t.Errorf("expected ErrRegisterClientFailed, got %v", err) + } + if registered != nil { + t.Error("expected nil result for failed registration") + } +} + +func TestOAuthClientsStore_SupportsRegistration(t *testing.T) { + // Test store without registration support + basicStore := NewOAuthClientStore(mockGetClientSuccess) + if basicStore.SupportsRegistration() { + t.Error("basic store should not support registration") + } + + // Test store with registration support + dynamicStore := NewOAuthClientStoreSupportDynamicRegistration(mockGetClientSuccess, mockRegisterClientSuccess) + if !dynamicStore.SupportsRegistration() { + t.Error("dynamic store should support registration") + } +} + +func TestOAuthClientsStore_InterfaceCompliance(t *testing.T) { + // Test that OAuthClientsStore implements OAuthClientsStoreInterface + var _ OAuthClientsStoreInterface = &OAuthClientsStore{} + + // Test that stores with registration support implement SupportDynamicClientRegistration + dynamicStore := NewOAuthClientStoreSupportDynamicRegistration(mockGetClientSuccess, mockRegisterClientSuccess) + var _ SupportDynamicClientRegistration = dynamicStore + + // Verify interface methods work as expected + store := NewOAuthClientStoreSupportDynamicRegistration(mockGetClientSuccess, mockRegisterClientSuccess) + + // Test as OAuthClientsStoreInterface + var iface OAuthClientsStoreInterface = store + client, err := iface.GetClient("existing-client") + if err != nil { + t.Fatalf("interface GetClient failed: %v", err) + } + if client.ClientID != "existing-client" { + t.Errorf("interface GetClient returned wrong client ID") + } + + // Test as SupportDynamicClientRegistration + var regIface SupportDynamicClientRegistration = store + newClient := auth.OAuthClientInformationFull{ + OAuthClientMetadata: auth.OAuthClientMetadata{ + RedirectURIs: []string{"https://test.com/callback"}, + }, + } + registered, err := regIface.RegisterClient(newClient) + if err != nil { + t.Fatalf("interface RegisterClient failed: %v", err) + } + if registered.ClientID != "generated-client-id" { + t.Errorf("interface RegisterClient returned wrong client ID") + } +} + +func TestOAuthClientsStore_EdgeCases(t *testing.T) { + t.Run("empty client ID", func(t *testing.T) { + store := NewOAuthClientStore(mockGetClientSuccess) + client, err := store.GetClient("") + if err != ErrClientNotFound { + t.Errorf("expected ErrClientNotFound for empty client ID, got %v", err) + } + if client != nil { + t.Error("expected nil client for empty client ID") + } + }) + + t.Run("register client with empty metadata", func(t *testing.T) { + store := NewOAuthClientStoreSupportDynamicRegistration(mockGetClientSuccess, mockRegisterClientSuccess) + + // Test registering a client with minimal metadata + emptyClient := auth.OAuthClientInformationFull{} + registered, err := store.RegisterClient(emptyClient) + if err != nil { + t.Errorf("registering empty client failed: %v", err) + } + if registered.ClientID != "generated-client-id" { + t.Errorf("expected generated client ID even for empty client") + } + }) +} + +// stringPtr returns a pointer to s +func stringPtr(s string) *string { + return &s +} diff --git a/internal/auth/server/handler/authorize.go b/internal/auth/server/handler/authorize.go new file mode 100644 index 0000000..c062856 --- /dev/null +++ b/internal/auth/server/handler/authorize.go @@ -0,0 +1,350 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package handler + +import ( + "encoding/json" + "fmt" + "net/http" + "net/url" + "strings" + + "github.com/go-playground/validator/v10" + "golang.org/x/time/rate" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth/pkce" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth/server" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth/server/middleware" + "trpc.group/trpc-go/trpc-mcp-go/internal/errors" +) + +// AuthorizationHandlerOptions contains configuration for the /authorize endpoint handler. +type AuthorizationHandlerOptions struct { + // Provider is the OAuth server implementation that issues authorization codes + // and validates client information. + Provider server.OAuthServerProvider `json:"provider"` + + // RateLimit applies a standard rate limiter to the authorization endpoint. + // This helps protect against abuse or brute-force attempts. + RateLimit *rate.Limiter `json:"rateLimit,omitempty"` +} + +// ClientAuthorizationParams defines the client parameters that must be validated +// before redirecting to the authorization endpoint. +type ClientAuthorizationParams struct { + // ClientID is the unique identifier for the client making the authorization request. + ClientID string `json:"client_id" validate:"required"` + + // RedirectURI is the callback URI where the authorization code will be sent. + RedirectURI string `json:"redirect_uri,omitempty" validate:"omitempty,url"` +} + +// RequestAuthorizationParams defines the parameters that must be validated +// for a successful OAuth 2.1 authorization request. +type RequestAuthorizationParams struct { + // ResponseType must be "code" for the authorization code flow. + ResponseType string `json:"response_type" validate:"required,eq=code"` + + // CodeChallenge is the PKCE code challenge generated by the client. + CodeChallenge string `json:"code_challenge" validate:"required"` + + // CodeChallengeMethod must be "S256" (SHA-256), the required method in OAuth 2.1. + CodeChallengeMethod string `json:"code_challenge_method" validate:"required,eq=S256"` + + // Scope is an optional space-delimited list of requested permissions. + Scope string `json:"scope,omitempty"` + + // State is an optional opaque value used by the client to maintain state + // between the request and callback (commonly for CSRF protection). + State string `json:"state,omitempty"` + + // Resource is an optional absolute URL indicating the resource being accessed. + Resource string `json:"resource,omitempty" validate:"omitempty,url"` +} + +// AuthorizationHandler creates an authorization handler +// Returns http.HandlerFunc for consistency with other handlers +func AuthorizationHandler(options AuthorizationHandlerOptions) http.HandlerFunc { + validate := validator.New() + + // Core handler logic + coreHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Cache-Control", "no-store") + + if r.Method == http.MethodPost { + if err := r.ParseForm(); err != nil { + handleDirectError(w, errors.NewOAuthError(errors.ErrServerError, "Failed to parse form data", "")) + return + } + } + + // Phase 1: Validate client_id and redirect_uri + _, redirectURI, client, err := validateClientAndRedirect(r, validate, options.Provider) + if err != nil { + handleDirectError(w, *err) + return + } + + // Phase 2: Validate other parameters and authorize + if authErr := processAuthorization(r, validate, client, redirectURI, options.Provider, w); authErr != nil { + state := getStateFromRequest(r) + errorRedirect := createErrorRedirect(redirectURI, *authErr, state) + http.Redirect(w, r, errorRedirect, http.StatusFound) + return + } + }) + + // Apply middleware if needed + var handler http.Handler = coreHandler + + // Apply rate limiting using standard middleware + if options.RateLimit != nil { + handler = middleware.RateLimitMiddleware(options.RateLimit)(handler) + } + + // Apply method restrictions (GET and POST allowed) + handler = middleware.AllowedMethods([]string{"GET", "POST"})(handler) + + // Convert back to http.HandlerFunc + return func(w http.ResponseWriter, r *http.Request) { + handler.ServeHTTP(w, r) + } +} + +// validateClientAndRedirect validates client_id and redirect_uri +func validateClientAndRedirect(r *http.Request, validate *validator.Validate, provider server.OAuthServerProvider) (string, string, *auth.OAuthClientInformationFull, *errors.OAuthError) { + clientParams := parseClientAuthorizationParams(r) + if err := validate.Struct(clientParams); err != nil { + oauthErr := errors.NewOAuthError(errors.ErrInvalidRequest, err.Error(), "") + return "", "", nil, &oauthErr + } + + clientID := clientParams.ClientID + redirectURI := clientParams.RedirectURI + + client, err := provider.ClientsStore().GetClient(clientID) + if err != nil { + oauthErr := errors.NewOAuthError(errors.ErrServerError, "Failed to get client", "") + return "", "", nil, &oauthErr + } + if client == nil { + oauthErr := errors.NewOAuthError(errors.ErrInvalidClient, "Invalid client_id", "") + return "", "", nil, &oauthErr + } + + // Validate redirect_uri + if redirectURI != "" { + found := false + for _, uri := range client.RedirectURIs { + if uri == redirectURI { + found = true + break + } + } + if !found { + oauthErr := errors.NewOAuthError(errors.ErrInvalidRequest, "Unregistered redirect_uri", "") + return "", "", nil, &oauthErr + } + } else if len(client.RedirectURIs) == 1 { + redirectURI = client.RedirectURIs[0] + } else { + oauthErr := errors.NewOAuthError(errors.ErrInvalidRequest, "redirect_uri must be specified when client has multiple registered URIs", "") + return "", "", nil, &oauthErr + } + + return clientID, redirectURI, client, nil +} + +// processAuthorization processes the authorization request +func processAuthorization(r *http.Request, validate *validator.Validate, client *auth.OAuthClientInformationFull, redirectURI string, provider server.OAuthServerProvider, w http.ResponseWriter) *errors.OAuthError { + reqParams := parseRequestAuthorizationParams(r) + if err := validate.Struct(reqParams); err != nil { + oauthErr := errors.NewOAuthError(errors.ErrInvalidRequest, err.Error(), "") + return &oauthErr + } + + tempAuthParams := server.AuthorizationParams{ + CodeChallenge: reqParams.CodeChallenge, + } + + if err := pkce.ValidatePKCEParams(tempAuthParams); err != nil { + oauthErr := errors.NewOAuthError(errors.ErrInvalidRequest, err.Error(), "") + return &oauthErr + } + + // Validate scopes + var requestedScopes []string + if reqParams.Scope != "" { + scopes := strings.Fields(reqParams.Scope) + for _, scope := range scopes { + if scope != "" { + requestedScopes = append(requestedScopes, scope) + } + } + + if err := validateScopes(requestedScopes, client); err != nil { + return err + } + } + + var resourceURL *url.URL + if reqParams.Resource != "" { + var err error + resourceURL, err = url.Parse(reqParams.Resource) + if err != nil { + oauthErr := errors.NewOAuthError(errors.ErrInvalidRequest, "Invalid resource URL", "") + return &oauthErr + } + + // Verify that the resource URL is an absolute URL + if !resourceURL.IsAbs() { + oauthErr := errors.NewOAuthError(errors.ErrInvalidRequest, "Resource must be an absolute URL", "") + return &oauthErr + } + } + + authParams := server.AuthorizationParams{ + State: reqParams.State, + Scopes: requestedScopes, + RedirectURI: redirectURI, + CodeChallenge: reqParams.CodeChallenge, + Resource: resourceURL, + } + + if err := provider.Authorize(*client, authParams, w, r); err != nil { + oauthErr := errors.NewOAuthError(errors.ErrServerError, "Authorization failed", "") + return &oauthErr + } + + return nil +} + +// validateScopes validates the requested scopes against client allowed scopes +func validateScopes(requestedScopes []string, client *auth.OAuthClientInformationFull) *errors.OAuthError { + // If no scope is requested, return success directly + if len(requestedScopes) == 0 { + return nil + } + + allowedScopes := make(map[string]bool) + + // Handling client-scoped configuration + if client.Scope != nil && *client.Scope != "" { + scopes := strings.Fields(*client.Scope) + for _, scope := range scopes { + if scope != "" { + allowedScopes[scope] = true + } + } + } + + // If the client does not have any scopes configured, reject all scope requests + if len(requestedScopes) == 0 { + oauthErr := errors.NewOAuthError(errors.ErrInvalidRequest, "Client has no registered scopes", "") + return &oauthErr + } + + // Verify the scope of each request + for _, scope := range requestedScopes { + scope = strings.TrimSpace(scope) + if scope == "" { + continue + } + if !allowedScopes[scope] { + oauthErr := errors.NewOAuthError(errors.ErrInvalidScope, fmt.Sprintf("Client was not registered with scope %s", scope), "") + return &oauthErr + } + } + + return nil +} + +// handleDirectError handles direct error responses (before redirect) +func handleDirectError(w http.ResponseWriter, oauthErr errors.OAuthError) { + status := http.StatusBadRequest + + switch oauthErr.ErrorCode { + case errors.ErrServerError.Error(): + status = http.StatusInternalServerError + default: + if oauthErr.ErrorCode == errors.ErrServerError.Error() { + status = http.StatusBadRequest + } + } + + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + json.NewEncoder(w).Encode(oauthErr.ToResponseStruct()) +} + +// getStateFromRequest extracts state parameter from request +func getStateFromRequest(r *http.Request) string { + if r.Method == http.MethodPost { + return r.FormValue("state") + } + return r.URL.Query().Get("state") +} + +// parseClientAuthorizationParams parses client authorization parameters +func parseClientAuthorizationParams(r *http.Request) ClientAuthorizationParams { + var params ClientAuthorizationParams + + if r.Method == http.MethodPost { + params.ClientID = strings.TrimSpace(r.FormValue("client_id")) + params.RedirectURI = strings.TrimSpace(r.FormValue("redirect_uri")) + } else { + query := r.URL.Query() + params.ClientID = strings.TrimSpace(query.Get("client_id")) + params.RedirectURI = strings.TrimSpace(query.Get("redirect_uri")) + } + + return params +} + +// parseRequestAuthorizationParams parses request authorization parameters +func parseRequestAuthorizationParams(r *http.Request) RequestAuthorizationParams { + var params RequestAuthorizationParams + + if r.Method == http.MethodPost { + params.ResponseType = strings.TrimSpace(r.FormValue("response_type")) + params.CodeChallenge = strings.TrimSpace(r.FormValue("code_challenge")) + params.CodeChallengeMethod = strings.TrimSpace(r.FormValue("code_challenge_method")) + params.Scope = strings.TrimSpace(r.FormValue("scope")) + params.State = r.FormValue("state") + params.Resource = strings.TrimSpace(r.FormValue("resource")) + } else { + query := r.URL.Query() + params.ResponseType = strings.TrimSpace(query.Get("response_type")) + params.CodeChallenge = strings.TrimSpace(query.Get("code_challenge")) + params.CodeChallengeMethod = strings.TrimSpace(query.Get("code_challenge_method")) + params.Scope = strings.TrimSpace(query.Get("scope")) + params.State = query.Get("state") + params.Resource = strings.TrimSpace(query.Get("resource")) + } + + return params +} + +// createErrorRedirect creates a redirect URL with error parameters +func createErrorRedirect(redirectURI string, err errors.OAuthError, state string) string { + errorURL, _ := url.Parse(redirectURI) + query := errorURL.Query() + + query.Set("error", err.ErrorCode) + query.Set("error_description", err.Message) + + if err.ErrorURI != "" { + query.Set("error_uri", err.ErrorURI) + } + + if state != "" { + query.Set("state", state) + } + + errorURL.RawQuery = query.Encode() + return errorURL.String() +} diff --git a/internal/auth/server/handler/authorize_test.go b/internal/auth/server/handler/authorize_test.go new file mode 100644 index 0000000..437539f --- /dev/null +++ b/internal/auth/server/handler/authorize_test.go @@ -0,0 +1,373 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package handler + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/time/rate" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth" + as "trpc.group/trpc-go/trpc-mcp-go/internal/auth/server" +) + +// validChallenge is a known good S256 PKCE code challenge used in tests +const validChallenge = "dBjftJeZ4CVP-mB92K27uhbUJU1p1r_wW1gFWFOEjXk" + +// oauthErrResp matches the JSON shape returned for OAuth error responses in tests +type oauthErrResp struct { + Error string `json:"error"` + ErrorDescription string `json:"error_description,omitempty"` + ErrorURI string `json:"error_uri,omitempty"` +} + +// mockProvider is a test double that satisfies the OAuthServerProvider interface +// it allows overriding Authorize via authorizeFunc for behavior-driven tests +type mockProvider struct { + store *as.OAuthClientsStore + authorizeFunc func(client auth.OAuthClientInformationFull, params as.AuthorizationParams, w http.ResponseWriter, r *http.Request) error +} + +// ClientsStore returns the in memory store used by the mock provider +func (m *mockProvider) ClientsStore() *as.OAuthClientsStore { return m.store } + +// Authorize simulates the authorization endpoint behavior +// if authorizeFunc is set it delegates to it +// otherwise it redirects to redirect_uri with a fixed code and optional state +func (m *mockProvider) Authorize(client auth.OAuthClientInformationFull, params as.AuthorizationParams, w http.ResponseWriter, r *http.Request) error { + // Delegate to custom behavior when provided + if m.authorizeFunc != nil { + return m.authorizeFunc(client, params, w, r) + } + // Compose redirect with code and optional state + u, _ := url.Parse(params.RedirectURI) + q := u.Query() + q.Set("code", "abc123") + if params.State != "" { + q.Set("state", params.State) + } + u.RawQuery = q.Encode() + // Issue HTTP 302 redirect + http.Redirect(w, r, u.String(), http.StatusFound) + return nil +} + +// ChallengeForAuthorizationCode returns an empty string in this mock provider +// real providers would return the stored code_challenge for the code +func (m *mockProvider) ChallengeForAuthorizationCode(client auth.OAuthClientInformationFull, authorizationCode string) (string, error) { + return "", nil +} + +// ExchangeAuthorizationCode is a stub that returns nil values for the mock +func (m *mockProvider) ExchangeAuthorizationCode(client auth.OAuthClientInformationFull, authorizationCode string, codeVerifier *string, redirectUri *string, resource *url.URL) (*auth.OAuthTokens, error) { + return nil, nil +} + +// ExchangeRefreshToken is a stub that returns nil values for the mock +func (m *mockProvider) ExchangeRefreshToken(client auth.OAuthClientInformationFull, refreshToken string, scopes []string, resource *url.URL) (*auth.OAuthTokens, error) { + return nil, nil +} + +// VerifyAccessToken is a stub that returns nil in this test double +func (m *mockProvider) VerifyAccessToken(token string) (*as.AuthInfo, error) { return nil, nil } + +// RevokeToken satisfies the optional SupportTokenRevocation interface with a no op +func (m *mockProvider) RevokeToken(client auth.OAuthClientInformationFull, request auth.OAuthTokenRevocationRequest) error { + return nil +} + +// makeStoreWithClient creates a store that returns the provided client when looked up by id +func makeStoreWithClient(c *auth.OAuthClientInformationFull) *as.OAuthClientsStore { + return as.NewOAuthClientStore(func(id string) (*auth.OAuthClientInformationFull, error) { + // Return client when ids match otherwise nil to simulate not found + if c != nil && c.ClientID == id { + return c, nil + } + return nil, nil + }) +} + +// makeClient builds a client record with id redirect uris and optional default scope +func makeClient(id string, redirects []string, scope *string) *auth.OAuthClientInformationFull { + return &auth.OAuthClientInformationFull{ + OAuthClientInformation: auth.OAuthClientInformation{ + ClientID: id, + }, + OAuthClientMetadata: auth.OAuthClientMetadata{ + RedirectURIs: redirects, + Scope: scope, + }, + } +} + +// newGET constructs a GET request helper for tests +func newGET(urlStr string) *http.Request { + return httptest.NewRequest(http.MethodGet, urlStr, nil) +} + +// newPOST constructs a POST request with x www form urlencoded body for tests +func newPOST(urlStr string, form url.Values) *http.Request { + req := httptest.NewRequest(http.MethodPost, urlStr, strings.NewReader(form.Encode())) + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + return req +} + +func TestAuthorization_SuccessGET(t *testing.T) { + // Prepare client with registered redirect and default scopes + scope := "read write" + client := makeClient("c1", []string{"https://app.example.com/cb"}, &scope) + mp := &mockProvider{store: makeStoreWithClient(client)} + + // Build handler under test + h := AuthorizationHandler(AuthorizationHandlerOptions{Provider: mp}) + + // Compose query for a valid authorization request + qs := url.Values{ + "client_id": {"c1"}, + "redirect_uri": {"https://app.example.com/cb"}, + "response_type": {"code"}, + "code_challenge": {validChallenge}, + "code_challenge_method": {"S256"}, + "state": {"st-123"}, + "scope": {"read"}, + } + req := newGET("/authorize?" + qs.Encode()) + rr := httptest.NewRecorder() + + // Execute handler + h.ServeHTTP(rr, req) + + // Assert redirect and parameters + assert.Equal(t, http.StatusFound, rr.Code) + loc := rr.Header().Get("Location") + u, err := url.Parse(loc) + require.NoError(t, err) + q := u.Query() + assert.Equal(t, "abc123", q.Get("code")) + assert.Equal(t, "st-123", q.Get("state")) +} + +func TestAuthorization_MissingClientID_JSON400(t *testing.T) { + client := makeClient("c1", []string{"https://app.example.com/cb"}, nil) + mp := &mockProvider{store: makeStoreWithClient(client)} + h := AuthorizationHandler(AuthorizationHandlerOptions{Provider: mp}) + + // Build request without client_id + req := newGET("/authorize?redirect_uri=https://app.example.com/cb") + rr := httptest.NewRecorder() + + // Execute handler + h.ServeHTTP(rr, req) + + // Validate error payload + assert.Equal(t, http.StatusBadRequest, rr.Code) + var resp oauthErrResp + require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &resp)) + assert.Equal(t, "invalid_request", resp.Error) + assert.NotEmpty(t, resp.ErrorDescription) +} + +func TestAuthorization_UnregisteredRedirect_JSON400(t *testing.T) { + client := makeClient("c1", []string{"https://app.example.com/cb"}, nil) + mp := &mockProvider{store: makeStoreWithClient(client)} + h := AuthorizationHandler(AuthorizationHandlerOptions{Provider: mp}) + + // Use an unregistered redirect_uri + req := newGET("/authorize?client_id=c1&redirect_uri=https://evil.example.com/cb") + rr := httptest.NewRecorder() + + // Execute and assert + h.ServeHTTP(rr, req) + assert.Equal(t, http.StatusBadRequest, rr.Code) + var resp oauthErrResp + require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &resp)) + assert.Equal(t, "invalid_request", resp.Error) + assert.Contains(t, strings.ToLower(resp.ErrorDescription), "redirect") +} + +func TestAuthorization_MultipleRedirects_RequireExplicit_JSON400(t *testing.T) { + client := makeClient("c1", []string{"https://a/cb", "https://b/cb"}, nil) + mp := &mockProvider{store: makeStoreWithClient(client)} + h := AuthorizationHandler(AuthorizationHandlerOptions{Provider: mp}) + + // Missing redirect_uri should fail when multiple are registered + req := newGET("/authorize?client_id=c1") + rr := httptest.NewRecorder() + + // Execute and assert + h.ServeHTTP(rr, req) + assert.Equal(t, http.StatusBadRequest, rr.Code) + var resp oauthErrResp + require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &resp)) + assert.Equal(t, "invalid_request", resp.Error) + assert.Contains(t, strings.ToLower(resp.ErrorDescription), "redirect") +} + +func TestAuthorization_InvalidScope_302_WithState(t *testing.T) { + scope := "read write" + client := makeClient("c1", []string{"https://app.example.com/cb"}, &scope) + mp := &mockProvider{store: makeStoreWithClient(client)} + h := AuthorizationHandler(AuthorizationHandlerOptions{Provider: mp}) + + // Request includes a scope not in the client's allowed set + qs := url.Values{ + "client_id": {"c1"}, + "redirect_uri": {"https://app.example.com/cb"}, + "response_type": {"code"}, + "code_challenge": {validChallenge}, + "code_challenge_method": {"S256"}, + "scope": {"delete"}, + "state": {"keep-me"}, + } + req := newGET("/authorize?" + qs.Encode()) + rr := httptest.NewRecorder() + + // Execute and assert error redirect + h.ServeHTTP(rr, req) + assert.Equal(t, http.StatusFound, rr.Code) + u, _ := url.Parse(rr.Header().Get("Location")) + q := u.Query() + assert.Equal(t, "invalid_scope", q.Get("error")) + assert.Equal(t, "keep-me", q.Get("state")) + assert.NotEmpty(t, q.Get("error_description")) +} + +func TestAuthorization_InvalidResourceURL_302_ErrorRedirect(t *testing.T) { + scope := "read" + client := makeClient("c1", []string{"https://app.example.com/cb"}, &scope) + mp := &mockProvider{store: makeStoreWithClient(client)} + h := AuthorizationHandler(AuthorizationHandlerOptions{Provider: mp}) + + // Provide a relative resource URL which is invalid + qs := url.Values{ + "client_id": {"c1"}, + "redirect_uri": {"https://app.example.com/cb"}, + "response_type": {"code"}, + "code_challenge": {validChallenge}, + "code_challenge_method": {"S256"}, + "resource": {"/relative"}, + } + req := newGET("/authorize?" + qs.Encode()) + rr := httptest.NewRecorder() + + // Execute and assert error redirect + h.ServeHTTP(rr, req) + assert.Equal(t, http.StatusFound, rr.Code) + u, _ := url.Parse(rr.Header().Get("Location")) + q := u.Query() + assert.Equal(t, "invalid_request", q.Get("error")) + assert.NotEmpty(t, q.Get("error_description")) +} + +func TestAuthorization_RateLimit_429_JSON(t *testing.T) { + client := makeClient("c1", []string{"https://app.example.com/cb"}, nil) + mp := &mockProvider{store: makeStoreWithClient(client)} + // Limiter with zero rate to always deny + limiter := rate.NewLimiter(0, 0) + + h := AuthorizationHandler(AuthorizationHandlerOptions{ + Provider: mp, + RateLimit: limiter, + }) + + // No params needed because limiter will block before validation + req := newGET("/authorize") + rr := httptest.NewRecorder() + + // Execute and assert + h.ServeHTTP(rr, req) + assert.Equal(t, http.StatusTooManyRequests, rr.Code) + var resp oauthErrResp + require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &resp)) + assert.Equal(t, "too_many_requests", resp.Error) +} + +func TestAllowedMethods_GET_and_POST(t *testing.T) { + client := makeClient("c1", []string{"https://cb"}, nil) + mp := &mockProvider{store: makeStoreWithClient(client)} + h := AuthorizationHandler(AuthorizationHandlerOptions{Provider: mp}) + + // GET should be allowed + rr1 := httptest.NewRecorder() + h.ServeHTTP(rr1, newGET("/authorize")) + assert.NotEqual(t, http.StatusMethodNotAllowed, rr1.Code) + + // PUT should be rejected with 405 + rr2 := httptest.NewRecorder() + req2 := httptest.NewRequest(http.MethodPut, "/authorize", nil) + h.ServeHTTP(rr2, req2) + assert.Equal(t, http.StatusMethodNotAllowed, rr2.Code) +} + +func TestHelpers_StateParsing_GET_and_POST(t *testing.T) { + // GET pathway + reqGet := newGET("/authorize?state=GETSTATE") + assert.Equal(t, "GETSTATE", getStateFromRequest(reqGet)) + + // POST pathway + form := url.Values{"state": {"POSTSTATE"}} + reqPost := newPOST("/authorize", form) + assert.Equal(t, "POSTSTATE", getStateFromRequest(reqPost)) +} + +func TestHelpers_ParseParams_Parity(t *testing.T) { + // ClientAuthorizationParams via GET + qs := url.Values{"client_id": {"c1"}, "redirect_uri": {"https://a/cb"}} + cp := parseClientAuthorizationParams(newGET("/authorize?" + qs.Encode())) + assert.Equal(t, "c1", cp.ClientID) + assert.Equal(t, "https://a/cb", cp.RedirectURI) + + // RequestAuthorizationParams via POST + form := url.Values{ + "response_type": {"code"}, + "code_challenge": {"abc"}, + "code_challenge_method": {"S256"}, + "scope": {"read write"}, + "resource": {"https://api.example.com"}, + "state": {"s1"}, + } + rp := parseRequestAuthorizationParams(newPOST("/authorize", form)) + assert.Equal(t, "code", rp.ResponseType) + assert.Equal(t, "abc", rp.CodeChallenge) + assert.Equal(t, "S256", rp.CodeChallengeMethod) + assert.Equal(t, "read write", rp.Scope) + assert.Equal(t, "s1", rp.State) + assert.Equal(t, "https://api.example.com", rp.Resource) +} + +func TestCreateErrorRedirect_ComposesQuery(t *testing.T) { + // Inline error type to mimic minimal shape used by createErrorRedirect + type inlineErr struct { + ErrorCode string + Message string + ErrorURI string + } + errObj := inlineErr{ErrorCode: "invalid request", Message: "oops"} + + // Serialize and rehydrate to assert structure not affected by json tags + bs, _ := json.Marshal(errObj) + var rehydrated struct { + ErrorCode string + Message string + ErrorURI string + } + _ = json.Unmarshal(bs, &rehydrated) + + // Build redirect URL and assert query parameters + loc := createErrorRedirect("https://app.example.com/cb", rehydrated, "st") + u, _ := url.Parse(loc) + q := u.Query() + assert.Equal(t, "invalid request", q.Get("error")) + assert.Equal(t, "oops", q.Get("error_description")) + assert.Equal(t, "st", q.Get("state")) +} diff --git a/internal/auth/server/handler/metadata.go b/internal/auth/server/handler/metadata.go new file mode 100644 index 0000000..4f7c8d0 --- /dev/null +++ b/internal/auth/server/handler/metadata.go @@ -0,0 +1,33 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package handler + +import ( + "encoding/json" + "net/http" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth/server/middleware" +) + +// MetadataHandler creates a handler for metadata endpoints +// This matches the TypeScript implementation using middleware composition +func MetadataHandler(metadata interface{}) http.HandlerFunc { + // Core handler that just serves JSON - no CORS or method validation + coreHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(metadata) + }) + + middlewareHandler := middleware.CorsMiddleware( + middleware.AllowedMethods([]string{"GET"})(coreHandler), + ) + + // Convert http.Handler to http.HandlerFunc + return func(w http.ResponseWriter, r *http.Request) { + middlewareHandler.ServeHTTP(w, r) + } +} diff --git a/internal/auth/server/handler/metadata_test.go b/internal/auth/server/handler/metadata_test.go new file mode 100644 index 0000000..217e2c7 --- /dev/null +++ b/internal/auth/server/handler/metadata_test.go @@ -0,0 +1,345 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package handler + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestMetadataHandler(t *testing.T) { + // Test metadata + testMetadata := map[string]interface{}{ + "name": "test-server", + "version": "1.0.0", + "description": "Test MCP server", + "capabilities": map[string]interface{}{ + "auth": true, + "tools": []interface{}{"test-tool-1", "test-tool-2"}, // JSON decoding converts to []interface{} + }, + } + + // Create handler + handler := MetadataHandler(testMetadata) + + // Verify handler is created successfully + assert.NotNil(t, handler) + + // Test GET request + req := httptest.NewRequest(http.MethodGet, "/metadata", nil) + w := httptest.NewRecorder() + + // Execute request + handler(w, req) + + // Verify response + assert.Equal(t, http.StatusOK, w.Code) + assert.Equal(t, "application/json", w.Header().Get("Content-Type")) + + // Verify JSON response + var responseData map[string]interface{} + err := json.NewDecoder(w.Body).Decode(&responseData) + require.NoError(t, err) + assert.Equal(t, testMetadata, responseData) +} + +func TestMetadataHandler_MethodValidation(t *testing.T) { + // Test metadata + testMetadata := map[string]interface{}{ + "name": "test-server", + } + + // Create handler + handler := MetadataHandler(testMetadata) + + // Test cases for different HTTP methods + testCases := []struct { + name string + method string + expectedCode int + shouldHaveJSON bool + shouldHaveAllow bool + }{ + { + name: "GET method allowed", + method: http.MethodGet, + expectedCode: http.StatusOK, + shouldHaveJSON: true, + shouldHaveAllow: false, + }, + { + name: "POST method not allowed", + method: http.MethodPost, + expectedCode: http.StatusMethodNotAllowed, + shouldHaveJSON: true, + shouldHaveAllow: true, + }, + { + name: "PUT method not allowed", + method: http.MethodPut, + expectedCode: http.StatusMethodNotAllowed, + shouldHaveJSON: true, + shouldHaveAllow: true, + }, + { + name: "DELETE method not allowed", + method: http.MethodDelete, + expectedCode: http.StatusMethodNotAllowed, + shouldHaveJSON: true, + shouldHaveAllow: true, + }, + { + name: "PATCH method not allowed", + method: http.MethodPatch, + expectedCode: http.StatusMethodNotAllowed, + shouldHaveJSON: true, + shouldHaveAllow: true, + }, + } + + // Execute tests + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + req := httptest.NewRequest(tc.method, "/metadata", nil) + w := httptest.NewRecorder() + + // Execute request + handler(w, req) + + // Verify status code + assert.Equal(t, tc.expectedCode, w.Code) + + // Verify Content-Type header + if tc.shouldHaveJSON { + assert.Equal(t, "application/json", w.Header().Get("Content-Type")) + } + + // Verify Allow header for method not allowed responses + if tc.shouldHaveAllow { + allowHeader := w.Header().Get("Allow") + assert.Contains(t, allowHeader, "GET") + } + + // For successful requests, verify metadata is returned + if tc.expectedCode == http.StatusOK { + var responseData map[string]interface{} + err := json.NewDecoder(w.Body).Decode(&responseData) + require.NoError(t, err) + assert.Equal(t, testMetadata, responseData) + } + }) + } +} + +func TestMetadataHandler_CORSHeaders(t *testing.T) { + // Test metadata + testMetadata := map[string]interface{}{ + "name": "test-server", + } + + // Create handler + handler := MetadataHandler(testMetadata) + + // Test cases for CORS handling + testCases := []struct { + name string + method string + origin string + expectedCode int + shouldHaveCORSOrigin bool + shouldHaveCORSMethods bool + }{ + { + name: "Non-CORS request", + method: http.MethodGet, + origin: "", + expectedCode: http.StatusOK, + shouldHaveCORSOrigin: false, + shouldHaveCORSMethods: false, + }, + { + name: "CORS GET request", + method: http.MethodGet, + origin: "https://example.com", + expectedCode: http.StatusOK, + shouldHaveCORSOrigin: true, + shouldHaveCORSMethods: true, + }, + { + name: "CORS OPTIONS preflight request", + method: http.MethodOptions, + origin: "https://example.com", + expectedCode: http.StatusNoContent, + shouldHaveCORSOrigin: true, + shouldHaveCORSMethods: true, + }, + } + + // Execute tests + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + req := httptest.NewRequest(tc.method, "/metadata", nil) + if tc.origin != "" { + req.Header.Set("Origin", tc.origin) + } + w := httptest.NewRecorder() + + // Execute request + handler(w, req) + + // Verify status code + assert.Equal(t, tc.expectedCode, w.Code) + + // Verify CORS headers + if tc.shouldHaveCORSOrigin { + assert.Equal(t, "*", w.Header().Get("Access-Control-Allow-Origin")) + } else { + assert.Empty(t, w.Header().Get("Access-Control-Allow-Origin")) + } + + if tc.shouldHaveCORSMethods { + allowMethods := w.Header().Get("Access-Control-Allow-Methods") + assert.Equal(t, "GET,HEAD,PUT,PATCH,POST,DELETE", allowMethods) + } + + // For OPTIONS requests, verify Content-Length header + if tc.method == http.MethodOptions { + assert.Equal(t, "0", w.Header().Get("Content-Length")) + } + + // For successful GET requests, verify metadata is returned + if tc.method == http.MethodGet && tc.expectedCode == http.StatusOK { + var responseData map[string]interface{} + err := json.NewDecoder(w.Body).Decode(&responseData) + require.NoError(t, err) + assert.Equal(t, testMetadata, responseData) + } + }) + } +} + +func TestMetadataHandler_DifferentMetadataTypes(t *testing.T) { + // Test cases with different metadata types + testCases := []struct { + name string + metadata interface{} + }{ + { + name: "String metadata", + metadata: "simple string metadata", + }, + { + name: "Number metadata", + metadata: float64(42), + }, + { + name: "Boolean metadata", + metadata: true, + }, + { + name: "Array metadata", + metadata: []interface{}{"item1", "item2", "item3"}, + }, + { + name: "Complex object metadata", + metadata: map[string]interface{}{ + "server": map[string]interface{}{ + "name": "mcp-server", + "version": "2.0.0", + "config": map[string]interface{}{ + "debug": true, + "port": float64(8080), + }, + }, + "capabilities": []interface{}{"auth", "tools", "resources"}, + "stats": map[string]interface{}{ + "uptime": float64(3600), + "requests": float64(1500), + "errors": float64(5), + }, + }, + }, + { + name: "Empty metadata", + metadata: map[string]interface{}{}, + }, + { + name: "Nil metadata", + metadata: nil, + }, + } + + // Execute tests + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + // Create handler with test metadata + handler := MetadataHandler(tc.metadata) + + // Create request + req := httptest.NewRequest(http.MethodGet, "/metadata", nil) + w := httptest.NewRecorder() + + // Execute request + handler(w, req) + + // Verify response + assert.Equal(t, http.StatusOK, w.Code) + assert.Equal(t, "application/json", w.Header().Get("Content-Type")) + + // Verify JSON response matches expected metadata + var responseData interface{} + err := json.NewDecoder(w.Body).Decode(&responseData) + require.NoError(t, err) + assert.Equal(t, tc.metadata, responseData) + }) + } +} + +func TestMetadataHandler_ConcurrentRequests(t *testing.T) { + // Test metadata + testMetadata := map[string]interface{}{ + "name": "concurrent-test-server", + "version": "1.0.0", + } + + // Create handler + handler := MetadataHandler(testMetadata) + + // Number of concurrent requests + numRequests := 10 + responses := make(chan *httptest.ResponseRecorder, numRequests) + + // Launch concurrent requests + for i := 0; i < numRequests; i++ { + go func() { + req := httptest.NewRequest(http.MethodGet, "/metadata", nil) + w := httptest.NewRecorder() + handler(w, req) + responses <- w + }() + } + + // Collect and verify all responses + for i := 0; i < numRequests; i++ { + w := <-responses + + // Verify response + assert.Equal(t, http.StatusOK, w.Code) + assert.Equal(t, "application/json", w.Header().Get("Content-Type")) + + // Verify JSON response + var responseData map[string]interface{} + err := json.NewDecoder(w.Body).Decode(&responseData) + require.NoError(t, err) + assert.Equal(t, testMetadata, responseData) + } +} diff --git a/internal/auth/server/handler/register.go b/internal/auth/server/handler/register.go new file mode 100644 index 0000000..01fe48b --- /dev/null +++ b/internal/auth/server/handler/register.go @@ -0,0 +1,241 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package handler + +import ( + "crypto/rand" + "encoding/hex" + "encoding/json" + "fmt" + "net/http" + "time" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth/server/middleware" + "trpc.group/trpc-go/trpc-mcp-go/internal/errors" + + "github.com/google/uuid" + "golang.org/x/time/rate" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth/server" +) + +const ( + DEFAULT_CLIENT_SECRET_EXPIRY_SECONDS = 30 * 24 * 60 * 60 // 30 days + DEFAULT_RATE_LIMIT_WINDOW_MS = 60 * 60 * 1000 // 1 hour + DEFAULT_RATE_LIMIT_MAX = 20 // 20 requests per hour +) + +// ClientRegistrationHandlerOptions configuration for client registration handler +type ClientRegistrationHandlerOptions struct { + // A store used to save information about dynamically registered OAuth clients. + ClientsStore server.SupportDynamicClientRegistration + + // The number of seconds after which to expire issued client secrets, or 0 to prevent expiration of client secrets (not recommended). + // If not set, defaults to 30 days. + ClientSecretExpirySeconds *int + + // Rate limiting configuration for the client registration endpoint. + // Set to nil to disable rate limiting for this endpoint. + // Registration endpoints are particularly sensitive to abuse and should be rate limited. + RateLimit *RegisterRateLimitConfig + + // Whether to generate a client ID before calling the client registration endpoint. + // If not set, defaults to true. + ClientIdGeneration *bool +} + +type RegisterRateLimitConfig struct { + WindowMs int // Window duration in milliseconds + Max int // Maximum requests per window + Message string // Customize over-limit prompt information +} + +// ClientRegistrationHandler creates a handler for OAuth client registration +func ClientRegistrationHandler(options ClientRegistrationHandlerOptions) http.Handler { + rateLimitConfig := options.RateLimit + if rateLimitConfig == nil { + rateLimitConfig = &RegisterRateLimitConfig{ + WindowMs: DEFAULT_RATE_LIMIT_WINDOW_MS, + Max: DEFAULT_RATE_LIMIT_MAX, + } + } + + if options.ClientsStore == nil { + // Return a handler that always returns an error + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusNotImplemented) + + notImplError := errors.NewOAuthError( + errors.ErrUnsupportedGrantType, + "Dynamic client registration is not supported by this server", + "https://datatracker.ietf.org/doc/html/rfc7591", + ) + json.NewEncoder(w).Encode(notImplError.ToResponseStruct()) + }) + } + + clientSecretExpirySeconds := DEFAULT_CLIENT_SECRET_EXPIRY_SECONDS + if options.ClientSecretExpirySeconds != nil { + clientSecretExpirySeconds = *options.ClientSecretExpirySeconds + } + + clientIdGeneration := true + if options.ClientIdGeneration != nil { + clientIdGeneration = *options.ClientIdGeneration + } + + // Core handler logic + coreHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Cache-Control", "no-store") + + // Parse JSON request body + var clientMetadata auth.OAuthClientMetadata + if err := json.NewDecoder(r.Body).Decode(&clientMetadata); err != nil { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + + oauthErr := errors.NewOAuthError( + errors.ErrInvalidClientMetadata, + fmt.Sprintf("Invalid JSON in request body: %v", err), + "", + ) + json.NewEncoder(w).Encode(oauthErr.ToResponseStruct()) + return + } + + // Validate client metadata + if err := validateClientMetadata(&clientMetadata); err != nil { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + + oauthErr := errors.NewOAuthError( + errors.ErrInvalidClientMetadata, + err.Error(), + "", + ) + json.NewEncoder(w).Encode(oauthErr.ToResponseStruct()) + return + } + + isPublicClient := clientMetadata.TokenEndpointAuthMethod == "none" + + // Generate client credentials + var clientSecret string + if !isPublicClient { + secret, err := generateClientSecret() + if err != nil { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusInternalServerError) + + oauthErr := errors.NewOAuthError( + errors.ErrServerError, + "Failed to generate client secret", + "", + ) + json.NewEncoder(w).Encode(oauthErr.ToResponseStruct()) + return + } + clientSecret = secret + } + + clientIdIssuedAt := time.Now().Unix() + + // Calculate client secret expiry time + clientsDoExpire := clientSecretExpirySeconds > 0 + var clientSecretExpiresAt *int64 + if !isPublicClient { + if clientsDoExpire { + expiryTime := clientIdIssuedAt + int64(clientSecretExpirySeconds) + clientSecretExpiresAt = &expiryTime + } else { + zero := int64(0) + clientSecretExpiresAt = &zero + } + } + + // Create client information + clientInfo := auth.OAuthClientInformationFull{ + OAuthClientMetadata: clientMetadata, + OAuthClientInformation: auth.OAuthClientInformation{ + ClientSecret: clientSecret, + ClientSecretExpiresAt: clientSecretExpiresAt, + }, + } + + if clientIdGeneration { + clientId := uuid.New().String() + clientInfo.OAuthClientInformation.ClientID = clientId + clientInfo.OAuthClientInformation.ClientIDIssuedAt = &clientIdIssuedAt + } + + registeredClient, err := options.ClientsStore.RegisterClient(clientInfo) + if err != nil { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusInternalServerError) + + oauthErr := errors.NewOAuthError( + errors.ErrServerError, + "Failed to register client", + "", + ) + json.NewEncoder(w).Encode(oauthErr.ToResponseStruct()) + return + } + + // Success response + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusCreated) + json.NewEncoder(w).Encode(registeredClient) + }) + + var handler http.Handler = coreHandler + + if options.RateLimit != nil { + windowDuration := time.Duration(rateLimitConfig.WindowMs) * time.Millisecond + limit := rate.Every(windowDuration / time.Duration(rateLimitConfig.Max)) + limiter := rate.NewLimiter(limit, rateLimitConfig.Max) + + handler = middleware.RateLimitMiddleware(limiter)(handler) + } + + handler = middleware.JSONValidationMiddleware()(handler) + + handler = middleware.AllowedMethods([]string{"POST"})(handler) + + handler = middleware.CorsMiddleware(handler) + + return handler +} + +// generateClientSecret generates a random 32-byte hex string +func generateClientSecret() (string, error) { + bytes := make([]byte, 32) + if _, err := rand.Read(bytes); err != nil { + return "", err + } + return hex.EncodeToString(bytes), nil +} + +// validateClientMetadata performs basic validation on client metadata +func validateClientMetadata(metadata *auth.OAuthClientMetadata) error { + // Add validation logic as needed + if metadata.TokenEndpointAuthMethod == "" { + return fmt.Errorf("token_endpoint_auth_method is required") + } + + switch metadata.TokenEndpointAuthMethod { + case "client_secret_basic", "client_secret_post", "none": + default: + return fmt.Errorf("invalid token_endpoint_auth_method: %s", metadata.TokenEndpointAuthMethod) + } + + if len(metadata.RedirectURIs) == 0 { + return fmt.Errorf("redirect_uris is required") + } + + return nil +} diff --git a/internal/auth/server/handler/register_test.go b/internal/auth/server/handler/register_test.go new file mode 100644 index 0000000..d4af024 --- /dev/null +++ b/internal/auth/server/handler/register_test.go @@ -0,0 +1,294 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package handler + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/time/rate" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth" +) + +// mockDynClientStore is a test double implementation of a dynamic client store +// It captures inputs, simulates errors, and returns controlled responses for testing +type mockDynClientStore struct { + wantErr error // optional error to return on RegisterClient + lastRegistered *auth.OAuthClientInformationFull // last registered client captured during call + returnedClient *auth.OAuthClientInformationFull // client to return instead of echoing input + callCount int // number of times RegisterClient was invoked +} + +// RegisterClient mocks the dynamic client registration behavior +func (m *mockDynClientStore) RegisterClient(in auth.OAuthClientInformationFull) (*auth.OAuthClientInformationFull, error) { + m.callCount++ + // capture input + tmp := in + m.lastRegistered = &tmp + + // if configured, return a forced error + if m.wantErr != nil { + return nil, m.wantErr + } + + // if configured, return a pre-set client + if m.returnedClient != nil { + return m.returnedClient, nil + } + + // by default echo back + return &in, nil +} + +// postJSONBody is a helper that sends an HTTP POST request with a JSON body +func postJSONBody(t *testing.T, h http.Handler, path string, body any) *httptest.ResponseRecorder { + t.Helper() + + var buf bytes.Buffer + + // encode request body as JSON if provided + if body != nil { + require.NoError(t, json.NewEncoder(&buf).Encode(body)) + } + + // build POST request + req := httptest.NewRequest(http.MethodPost, path, &buf) + req.Header.Set("Content-Type", "application/json") + + // record the response + rr := httptest.NewRecorder() + h.ServeHTTP(rr, req) + return rr +} + +func TestClientRegistration_NotImplemented_WhenNoStore(t *testing.T) { + // ClientsStore==nil -> 501 Not Implemented (Explicitly returns that dynamic registration is not supported) + h := ClientRegistrationHandler(ClientRegistrationHandlerOptions{}) + + rr := postJSONBody(t, h, "/register", map[string]any{ + "token_endpoint_auth_method": "client_secret_post", + "redirect_uris": []string{"https://cb"}, + }) + assert.Equal(t, http.StatusNotImplemented, rr.Code) + + var resp map[string]any + require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &resp)) + assert.Equal(t, "unsupported_grant_type", resp["error"]) +} + +func TestClientRegistration_MethodNotAllowed_Get405(t *testing.T) { + // Only POST is allowed, other methods will return 405 + store := &mockDynClientStore{} + h := ClientRegistrationHandler(ClientRegistrationHandlerOptions{ + ClientsStore: store, + }) + + req := httptest.NewRequest(http.MethodGet, "/register", nil) + rr := httptest.NewRecorder() + h.ServeHTTP(rr, req) + assert.Equal(t, http.StatusMethodNotAllowed, rr.Code) +} + +func TestClientRegistration_InvalidJSON_400(t *testing.T) { + // JSON parsing failed -> 400 + invalid_client_metadata (determined within core logic) + store := &mockDynClientStore{} + h := ClientRegistrationHandler(ClientRegistrationHandlerOptions{ClientsStore: store}) + + req := httptest.NewRequest(http.MethodPost, "/register", bytes.NewBufferString("{bad json")) + req.Header.Set("Content-Type", "application/json") + rr := httptest.NewRecorder() + h.ServeHTTP(rr, req) + + assert.Equal(t, http.StatusBadRequest, rr.Code) + var resp map[string]any + require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &resp)) + assert.Equal(t, "invalid_client_metadata", resp["error"]) +} + +func TestClientRegistration_MetadataValidation_400(t *testing.T) { + // Missing required fields -> validateClientMetadata returns an error -> 400 (token_endpoint_auth_method and redirect_uris required) + store := &mockDynClientStore{} + h := ClientRegistrationHandler(ClientRegistrationHandlerOptions{ClientsStore: store}) + + rr := postJSONBody(t, h, "/register", map[string]any{ + "redirect_uris": []string{}, // invalid + }) + assert.Equal(t, http.StatusBadRequest, rr.Code) + + var resp map[string]any + _ = json.Unmarshal(rr.Body.Bytes(), &resp) + assert.Equal(t, "invalid_client_metadata", resp["error"]) +} + +func TestClientRegistration_PublicClient_NoSecret(t *testing.T) { + // token_endpoint_auth_method == "none" -> Public client: does not generate client_secret and does not set the expiration field + store := &mockDynClientStore{} + h := ClientRegistrationHandler(ClientRegistrationHandlerOptions{ClientsStore: store}) + + rr := postJSONBody(t, h, "/register", map[string]any{ + "token_endpoint_auth_method": "none", + "redirect_uris": []string{"https://cb"}, + }) + assert.Equal(t, http.StatusCreated, rr.Code) + + var resp auth.OAuthClientInformationFull + require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &resp)) + + // By default, the server generates client_id / client_id_issued_at (ClientIdGeneration is true by default) + assert.NotEmpty(t, resp.OAuthClientInformation.ClientID) + assert.NotNil(t, resp.OAuthClientInformation.ClientIDIssuedAt) + + // Public clients have no secrets and no expiration + assert.Equal(t, "", resp.OAuthClientInformation.ClientSecret) + assert.Nil(t, resp.OAuthClientInformation.ClientSecretExpiresAt) +} + +func TestClientRegistration_ConfidentialClient_GenerateSecret_And_Expiry(t *testing.T) { + // Non-public clients generate a 32-byte hex secret (length 64) and set an expiration time (default 30 days) + store := &mockDynClientStore{} + h := ClientRegistrationHandler(ClientRegistrationHandlerOptions{ClientsStore: store}) + + start := time.Now().Unix() + rr := postJSONBody(t, h, "/register", map[string]any{ + "token_endpoint_auth_method": "client_secret_post", + "redirect_uris": []string{"https://cb"}, + }) + assert.Equal(t, http.StatusCreated, rr.Code) + + var resp auth.OAuthClientInformationFull + require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &resp)) + + secret := resp.OAuthClientInformation.ClientSecret + if assert.NotEmpty(t, secret) { + assert.Len(t, secret, 64) // 32 bytes hex -> 64 chars + } + if assert.NotNil(t, resp.OAuthClientInformation.ClientSecretExpiresAt) { + exp := *resp.OAuthClientInformation.ClientSecretExpiresAt + assert.Greater(t, exp, start) + } +} + +func TestClientRegistration_ConfidentialClient_NoExpiryWhenZeroConfig(t *testing.T) { + // ClientSecretExpirySeconds==0 -> Does not expire (value is 0) + store := &mockDynClientStore{} + zero := 0 + h := ClientRegistrationHandler(ClientRegistrationHandlerOptions{ + ClientsStore: store, + ClientSecretExpirySeconds: &zero, + }) + + rr := postJSONBody(t, h, "/register", map[string]any{ + "token_endpoint_auth_method": "client_secret_basic", + "redirect_uris": []string{"https://cb"}, + }) + assert.Equal(t, http.StatusCreated, rr.Code) + + var resp auth.OAuthClientInformationFull + require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &resp)) + if assert.NotNil(t, resp.OAuthClientInformation.ClientSecretExpiresAt) { + assert.Equal(t, int64(0), *resp.OAuthClientInformation.ClientSecretExpiresAt) + } +} + +func TestClientRegistration_RegisterError_500(t *testing.T) { + // RegisterClient error -> 500 server_error + store := &mockDynClientStore{wantErr: assert.AnError} + h := ClientRegistrationHandler(ClientRegistrationHandlerOptions{ClientsStore: store}) + + rr := postJSONBody(t, h, "/register", map[string]any{ + "token_endpoint_auth_method": "client_secret_post", + "redirect_uris": []string{"https://cb"}, + }) + assert.Equal(t, http.StatusInternalServerError, rr.Code) + + var resp map[string]any + require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &resp)) + assert.Equal(t, "server_error", resp["error"]) +} + +func TestClientRegistration_RateLimit_429_WhenEnabled(t *testing.T) { + // Explicitly enable current limiting configuration: Max=1, the second request in the same window will result in a 429 error + store := &mockDynClientStore{} + h := ClientRegistrationHandler(ClientRegistrationHandlerOptions{ + ClientsStore: store, + RateLimit: &RegisterRateLimitConfig{ + WindowMs: 60_000, + Max: 1, + Message: "too many", + }, + }) + + // First request OK + rr1 := postJSONBody(t, h, "/register", map[string]any{ + "token_endpoint_auth_method": "client_secret_post", + "redirect_uris": []string{"https://cb"}, + }) + assert.Equal(t, http.StatusCreated, rr1.Code) + + // Second immediate request -> 429 + rr2 := postJSONBody(t, h, "/register", map[string]any{ + "token_endpoint_auth_method": "client_secret_post", + "redirect_uris": []string{"https://cb"}, + }) + assert.Equal(t, http.StatusTooManyRequests, rr2.Code) +} + +// Extra: Verified that rate limiting is not applied when RateLimit=nil to avoid interfering with other tests. +func TestClientRegistration_NoRateLimitByDefault(t *testing.T) { + store := &mockDynClientStore{} + h := ClientRegistrationHandler(ClientRegistrationHandlerOptions{ + ClientsStore: store, + // RateLimit omitted -> unlimited flow (create limiter and wrap middleware only when it is not nil) + }) + + for i := 0; i < 3; i++ { + rr := postJSONBody(t, h, "/register", map[string]any{ + "token_endpoint_auth_method": "client_secret_post", + "redirect_uris": []string{"https://cb"}, + }) + assert.Equal(t, http.StatusCreated, rr.Code) + } +} + +// When a custom current limiting configuration is in effect, +// the window parameter conversion logic using rate.Limiter will not panic +func TestClientRegistration_CustomLimiter_DoesNotPanic(t *testing.T) { + store := &mockDynClientStore{} + cfg := &RegisterRateLimitConfig{WindowMs: 1000, Max: 2} + h := ClientRegistrationHandler(ClientRegistrationHandlerOptions{ + ClientsStore: store, + RateLimit: cfg, + }) + + // Manually construct 2 requests, the third one should be 429 + rr1 := postJSONBody(t, h, "/register", map[string]any{ + "token_endpoint_auth_method": "client_secret_post", + "redirect_uris": []string{"https://cb"}, + }) + assert.Equal(t, http.StatusCreated, rr1.Code) + + rr2 := postJSONBody(t, h, "/register", map[string]any{ + "token_endpoint_auth_method": "client_secret_post", + "redirect_uris": []string{"https://cb"}, + }) + assert.Equal(t, http.StatusCreated, rr2.Code) + + rr3 := postJSONBody(t, h, "/register", map[string]any{ + "token_endpoint_auth_method": "client_secret_post", + "redirect_uris": []string{"https://cb"}, + }) + assert.Equal(t, http.StatusTooManyRequests, rr3.Code) + + // Extra sanity check + _ = rate.NewLimiter(rate.Every(time.Second/time.Duration(cfg.Max)), cfg.Max) +} diff --git a/internal/auth/server/handler/revoke.go b/internal/auth/server/handler/revoke.go new file mode 100644 index 0000000..36edc78 --- /dev/null +++ b/internal/auth/server/handler/revoke.go @@ -0,0 +1,232 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package handler + +import ( + "encoding/json" + "golang.org/x/time/rate" + "net/http" + "strings" + "time" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth/server" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth/server/middleware" + "trpc.group/trpc-go/trpc-mcp-go/internal/errors" +) + +// RevocationHandlerOptions configuration for the token revocation endpoint +type RevocationHandlerOptions struct { + Provider server.OAuthServerProvider + RateLimit *RevocationRateLimitConfig // Set to nil to disable rate limiting for this endpoint + RequireHTTPS bool // Enforce HTTPS in production (recommended for OAuth 2.1) + AllowJSONFallback bool // Allow JSON format for backward compatibility (non-compliant with RFC 7009) + EnableMCPHeaders bool // Add MCP-specific headers for MCP 2025-03-26 compliance +} + +// RevocationRateLimitConfig rate limiting configuration +type RevocationRateLimitConfig struct { + WindowMs int // Window duration in milliseconds + Max int // Maximum requests per window +} + +// RevocationHandler creates a handler for OAuth token revocation with client authentication middleware +func RevocationHandler(opts RevocationHandlerOptions) http.Handler { + // Check if provider supports token revocation + revoker, ok := opts.Provider.(server.SupportTokenRevocation) + if !ok { + panic("Auth provider does not support revoking tokens") + } + + // Create the core handler + coreHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Set OAuth 2.1 required security headers + w.Header().Set("Cache-Control", "no-store") + w.Header().Set("Pragma", "no-cache") + + // Add MCP-specific headers if enabled + if opts.EnableMCPHeaders { + w.Header().Set("X-MCP-Version", "2025-03-26") + w.Header().Set("X-MCP-Transport", "http") + } + + // Enforce HTTPS if required (OAuth 2.1 best practice) + if opts.RequireHTTPS && r.TLS == nil && r.Header.Get("X-Forwarded-Proto") != "https" { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + + invalidReqError := errors.NewOAuthError( + errors.ErrInvalidRequest, + "HTTPS is required for OAuth 2.1 token revocation", + "https://datatracker.ietf.org/doc/html/rfc6749#section-3", + ) + json.NewEncoder(w).Encode(invalidReqError.ToResponseStruct()) + return + } + + // RFC 7009 Section 2.1: Strict Content-Type validation + contentType := r.Header.Get("Content-Type") + isURLEncoded := strings.HasPrefix(contentType, "application/x-www-form-urlencoded") + isJSON := strings.HasPrefix(contentType, "application/json") + + if !isURLEncoded && (!opts.AllowJSONFallback || !isJSON) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + + errorMsg := "Content-Type must be application/x-www-form-urlencoded per RFC 7009" + if opts.AllowJSONFallback { + errorMsg = "Content-Type must be application/x-www-form-urlencoded (preferred) or application/json" + } + + invalidReqError := errors.NewOAuthError( + errors.ErrInvalidRequest, + errorMsg, + "https://datatracker.ietf.org/doc/html/rfc7009#section-2.1", + ) + json.NewEncoder(w).Encode(invalidReqError.ToResponseStruct()) + return + } + + // Parse request body based on content type + var reqBody auth.OAuthTokenRevocationRequest + + if isURLEncoded { + // RFC 7009 compliant: Parse URL-encoded form data + if err := r.ParseForm(); err != nil { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + + invalidReqError := errors.NewOAuthError(errors.ErrInvalidRequest, + "Failed to parse application/x-www-form-urlencoded data: "+err.Error(), "") + json.NewEncoder(w).Encode(invalidReqError.ToResponseStruct()) + return + } + + // Extract form values + reqBody.Token = r.FormValue("token") + reqBody.TokenTypeHint = r.FormValue("token_type_hint") + } else if opts.AllowJSONFallback && isJSON { + // Non-standard JSON fallback for backward compatibility + if err := json.NewDecoder(r.Body).Decode(&reqBody); err != nil { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + + invalidReqError := errors.NewOAuthError(errors.ErrInvalidRequest, + "Failed to parse JSON data: "+err.Error(), "") + json.NewEncoder(w).Encode(invalidReqError.ToResponseStruct()) + return + } + } + + // Validate request - token is required (RFC 7009 Section 2.1) + if err := validateRevocationRequest(reqBody); err != nil { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + json.NewEncoder(w).Encode(err.(errors.OAuthError).ToResponseStruct()) + return + } + + // Get authenticated client from context (set by clientAuth middleware) + client, ok := middleware.GetAuthenticatedClient(r) + if !ok { + // This should never happen if middleware is properly configured + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusInternalServerError) + + serverError := errors.NewOAuthError(errors.ErrServerError, "Internal Server Error", "") + json.NewEncoder(w).Encode(serverError.ToResponseStruct()) + return + } + + // Revoke the token + err := revoker.RevokeToken(*client, reqBody) + if err != nil { + w.Header().Set("Content-Type", "application/json") + + if oauthErr, ok := err.(errors.OAuthError); ok { + status := http.StatusBadRequest + if oauthErr.ErrorCode == errors.ErrServerError.Error() { + status = http.StatusInternalServerError + } + w.WriteHeader(status) + json.NewEncoder(w).Encode(oauthErr.ToResponseStruct()) + return + } + + w.WriteHeader(http.StatusInternalServerError) + serverError := errors.NewOAuthError(errors.ErrServerError, "Internal Server Error", "") + json.NewEncoder(w).Encode(serverError.ToResponseStruct()) + return + } + + // RFC 7009 Section 2: Success response - HTTP 200 with empty JSON + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + w.Write([]byte("{}")) + }) + + // Apply middlewares in order (wrapping from inside out to match TS middleware order) + var handler http.Handler = coreHandler + + // Apply client authentication middleware (innermost, like TS) + handler = middleware.AuthenticateClient(middleware.ClientAuthenticationMiddlewareOptions{ + ClientsStore: opts.Provider.ClientsStore(), + })(handler) + + // Apply rate limiting middleware only if explicitly configured + if opts.RateLimit != nil { + windowDuration := time.Duration(opts.RateLimit.WindowMs) * time.Millisecond + if opts.RateLimit.Max <= 0 { + panic("RateLimit Max must be greater than 0") + } + + // Calculate rate to match the window-based approach of express-rate-limit + limit := rate.Every(windowDuration / time.Duration(opts.RateLimit.Max)) + limiter := rate.NewLimiter(limit, opts.RateLimit.Max) + + handler = middleware.RateLimitMiddleware(limiter)(handler) + } + // Note: No default rate limiting applied when opts.RateLimit is nil + // This matches the TypeScript behavior where rateLimit: false disables it + + // Apply URL-encoded parsing middleware (RFC 7009 compliance) + handler = middleware.URLEncodedValidationMiddleware(opts.AllowJSONFallback)(handler) + + // Apply method restriction middleware (only POST allowed per RFC 7009) + handler = middleware.AllowedMethods([]string{"POST"})(handler) + + // Apply CORS middleware (outermost, like TS) + handler = middleware.CorsMiddleware(handler) + + return handler +} + +// validateRevocationRequest validates the OAuth token revocation request per RFC 7009 +func validateRevocationRequest(reqBody auth.OAuthTokenRevocationRequest) error { + // RFC 7009 Section 2.1: token parameter is required + if reqBody.Token == "" { + return errors.NewOAuthError(errors.ErrInvalidRequest, + "token parameter is required per RFC 7009", + "https://datatracker.ietf.org/doc/html/rfc7009#section-2.1") + } + + // RFC 7009 Section 2.1: token_type_hint is optional but must be valid if provided + if reqBody.TokenTypeHint != "" { + validTypes := map[string]bool{ + "access_token": true, + "refresh_token": true, + } + if !validTypes[reqBody.TokenTypeHint] { + return errors.NewOAuthError( + errors.ErrInvalidRequest, + "invalid token_type_hint, must be 'access_token' or 'refresh_token' per RFC 7009", + "https://datatracker.ietf.org/doc/html/rfc7009#section-2.1", + ) + } + } + + return nil +} diff --git a/internal/auth/server/handler/revoke_test.go b/internal/auth/server/handler/revoke_test.go new file mode 100644 index 0000000..6f860ee --- /dev/null +++ b/internal/auth/server/handler/revoke_test.go @@ -0,0 +1,224 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package handler + +import ( + "net/http" + "net/http/httptest" + "net/url" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth" + as "trpc.group/trpc-go/trpc-mcp-go/internal/auth/server" +) + +const ( + testClientID = "c1" + testClientSecret = "s3cr3t" +) + +// mockRevokeProvider is a fake implementation of OAuthServerProvider used for revocation tests +// It tracks calls to RevokeToken and allows simulating errors +type mockRevokeProvider struct { + store *as.OAuthClientsStore // backing client store + lastReq *auth.OAuthTokenRevocationRequest // last revocation request captured + revokeErr error // error to return from RevokeToken + calledRevoke int // number of times RevokeToken was invoked +} + +func (m *mockRevokeProvider) ClientsStore() *as.OAuthClientsStore { + return m.store +} + +func (m *mockRevokeProvider) RevokeToken(client auth.OAuthClientInformationFull, request auth.OAuthTokenRevocationRequest) error { + m.calledRevoke++ + // capture the request for assertions + tmp := request + m.lastReq = &tmp + return m.revokeErr +} + +// The remaining interface methods are no-ops since they are not used in revocation tests +func (m *mockRevokeProvider) Authorize(client auth.OAuthClientInformationFull, params as.AuthorizationParams, w http.ResponseWriter, r *http.Request) error { + return nil +} +func (m *mockRevokeProvider) ChallengeForAuthorizationCode(client auth.OAuthClientInformationFull, authorizationCode string) (string, error) { + return "", nil +} +func (m *mockRevokeProvider) ExchangeAuthorizationCode(client auth.OAuthClientInformationFull, authorizationCode string, codeVerifier *string, redirectUri *string, resource *url.URL) (*auth.OAuthTokens, error) { + return nil, nil +} +func (m *mockRevokeProvider) ExchangeRefreshToken(client auth.OAuthClientInformationFull, refreshToken string, scopes []string, resource *url.URL) (*auth.OAuthTokens, error) { + return nil, nil +} +func (m *mockRevokeProvider) VerifyAccessToken(token string) (*as.AuthInfo, error) { return nil, nil } + +// makeClientBasic constructs a client using client_secret_basic authentication +func makeClientBasic(id string) *auth.OAuthClientInformationFull { + return &auth.OAuthClientInformationFull{ + OAuthClientInformation: auth.OAuthClientInformation{ + ClientID: id, + ClientSecret: testClientSecret, + }, + OAuthClientMetadata: auth.OAuthClientMetadata{ + RedirectURIs: []string{"https://cb"}, + TokenEndpointAuthMethod: "client_secret_basic", + }, + } +} + +// makeClientPost constructs a client using client_secret_post authentication +func makeClientPost(id string) *auth.OAuthClientInformationFull { + return &auth.OAuthClientInformationFull{ + OAuthClientInformation: auth.OAuthClientInformation{ + ClientID: id, + ClientSecret: testClientSecret, + }, + OAuthClientMetadata: auth.OAuthClientMetadata{ + RedirectURIs: []string{"https://cb"}, + TokenEndpointAuthMethod: "client_secret_post", + }, + } +} + +// postFormBasicAuth helper submits a POST form request using HTTP Basic authentication +func postFormBasicAuth(t *testing.T, h http.Handler, path, clientID, clientSecret string, form url.Values) *httptest.ResponseRecorder { + t.Helper() + if form == nil { + form = url.Values{} + } + req := httptest.NewRequest(http.MethodPost, path, strings.NewReader(form.Encode())) + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + req.SetBasicAuth(clientID, clientSecret) + rr := httptest.NewRecorder() + h.ServeHTTP(rr, req) + return rr +} + +// postFormClientSecretPost helper submits a POST form request using client_secret_post authentication +func postFormClientSecretPost(t *testing.T, h http.Handler, path, clientID, clientSecret string, form url.Values) *httptest.ResponseRecorder { + t.Helper() + if form == nil { + form = url.Values{} + } + form.Set("client_id", clientID) + form.Set("client_secret", clientSecret) + + req := httptest.NewRequest(http.MethodPost, path, strings.NewReader(form.Encode())) + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + rr := httptest.NewRecorder() + h.ServeHTTP(rr, req) + return rr +} + +func TestRevocation_Success_200(t *testing.T) { + mpBasic := &mockRevokeProvider{store: makeStoreWithClient(makeClientBasic(testClientID))} + hBasic := RevocationHandler(RevocationHandlerOptions{Provider: mpBasic}) + + form := url.Values{"token": {"at-123"}} + rr := postFormBasicAuth(t, hBasic, "/revoke", testClientID, testClientSecret, form) + + if rr.Code != http.StatusOK { + mpPost := &mockRevokeProvider{store: makeStoreWithClient(makeClientPost(testClientID))} + hPost := RevocationHandler(RevocationHandlerOptions{Provider: mpPost}) + rr2 := postFormClientSecretPost(t, hPost, "/revoke", testClientID, testClientSecret, form) + + if rr2.Code != http.StatusOK { + t.Skipf("Skip: Authentication failed (Basic=%d/%s, Post=%d/%s)", + rr.Code, strings.TrimSpace(rr.Body.String()), + rr2.Code, strings.TrimSpace(rr2.Body.String()), + ) + return + } + assert.Equal(t, 1, mpPost.calledRevoke) + require.NotNil(t, mpPost.lastReq) + assert.Equal(t, "at-123", mpPost.lastReq.Token) + return + } + + assert.Equal(t, http.StatusOK, rr.Code) + assert.Equal(t, 1, mpBasic.calledRevoke) + require.NotNil(t, mpBasic.lastReq) + assert.Equal(t, "at-123", mpBasic.lastReq.Token) +} + +func TestRevocation_MissingToken_400(t *testing.T) { + mp := &mockRevokeProvider{store: makeStoreWithClient(makeClientBasic(testClientID))} + h := RevocationHandler(RevocationHandlerOptions{Provider: mp}) + + rr := postFormBasicAuth(t, h, "/revoke", testClientID, testClientSecret, url.Values{}) + assert.Equal(t, http.StatusBadRequest, rr.Code) + assert.Contains(t, strings.ToLower(rr.Body.String()), "invalid_request") +} + +func TestRevocation_UnsupportedTokenHint_Still200(t *testing.T) { + mpBasic := &mockRevokeProvider{store: makeStoreWithClient(makeClientBasic(testClientID))} + hBasic := RevocationHandler(RevocationHandlerOptions{Provider: mpBasic}) + + form := url.Values{ + "token": {"rt-xyz"}, + "token_type_hint": {"access_token"}, + } + rr := postFormBasicAuth(t, hBasic, "/revoke", testClientID, testClientSecret, form) + + if rr.Code != http.StatusOK { + mpPost := &mockRevokeProvider{store: makeStoreWithClient(makeClientPost(testClientID))} + hPost := RevocationHandler(RevocationHandlerOptions{Provider: mpPost}) + rr2 := postFormClientSecretPost(t, hPost, "/revoke", testClientID, testClientSecret, form) + + if rr2.Code != http.StatusOK { + t.Skipf("Skip: Authentication failed (Basic=%d/%s, Post=%d/%s)", + rr.Code, strings.TrimSpace(rr.Body.String()), + rr2.Code, strings.TrimSpace(rr2.Body.String()), + ) + return + } + require.Equal(t, http.StatusOK, rr2.Code) + return + } + + require.Equal(t, http.StatusOK, rr.Code) +} + +func TestRevocation_MethodNotAllowed_405(t *testing.T) { + mp := &mockRevokeProvider{store: makeStoreWithClient(makeClientBasic(testClientID))} + h := RevocationHandler(RevocationHandlerOptions{Provider: mp}) + + req := httptest.NewRequest(http.MethodGet, "/revoke", nil) + rr := httptest.NewRecorder() + h.ServeHTTP(rr, req) + assert.Equal(t, http.StatusMethodNotAllowed, rr.Code) +} + +func TestRevocation_RateLimit_429(t *testing.T) { + mp := &mockRevokeProvider{store: makeStoreWithClient(makeClientBasic(testClientID))} + h := RevocationHandler(RevocationHandlerOptions{ + Provider: mp, + RateLimit: &RevocationRateLimitConfig{ + WindowMs: 60_000, + Max: 1, + }, + }) + + _ = postFormBasicAuth(t, h, "/revoke", testClientID, testClientSecret, url.Values{"token": {"at-123"}}) + + rr2 := postFormBasicAuth(t, h, "/revoke", testClientID, testClientSecret, url.Values{"token": {"at-456"}}) + require.Equal(t, http.StatusTooManyRequests, rr2.Code) +} + +func TestRevocation_OPTIONS_405(t *testing.T) { + mp := &mockRevokeProvider{store: makeStoreWithClient(makeClientBasic(testClientID))} + h := RevocationHandler(RevocationHandlerOptions{Provider: mp}) + + req := httptest.NewRequest(http.MethodOptions, "/revoke", nil) + rr := httptest.NewRecorder() + h.ServeHTTP(rr, req) + assert.Equal(t, http.StatusMethodNotAllowed, rr.Code) +} diff --git a/internal/auth/server/handler/token.go b/internal/auth/server/handler/token.go new file mode 100644 index 0000000..9485d04 --- /dev/null +++ b/internal/auth/server/handler/token.go @@ -0,0 +1,382 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package handler + +import ( + "encoding/json" + "net/http" + "net/url" + "strings" + "time" + + "github.com/go-playground/validator/v10" + "golang.org/x/time/rate" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth/pkce" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth/server" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth/server/middleware" + "trpc.group/trpc-go/trpc-mcp-go/internal/errors" +) + +// TokenHandlerOptions defines configuration options for the token endpoint +type TokenHandlerOptions struct { + Provider server.OAuthServerProvider `json:"provider"` + RateLimit *rate.Limiter `json:"rateLimit,omitempty"` +} + +// TokenRequest defines the base structure of a token request. +// Every token request must specify a grant_type to indicate which flow is being used. +type TokenRequest struct { + // GrantType is the type of OAuth grant being requested. + // Common values include "authorization_code" and "refresh_token". + GrantType string `form:"grant_type" json:"grant_type" validate:"required"` +} + +// AuthorizationCodeGrant represents a token request using the Authorization Code flow. +type AuthorizationCodeGrant struct { + // Code is the authorization code previously issued to the client. + Code string `form:"code" json:"code" validate:"required"` + + // CodeVerifier is the PKCE verifier string that matches the original code_challenge. + CodeVerifier string `form:"code_verifier" json:"code_verifier" validate:"required"` + + // RedirectURI must match the redirect_uri used in the authorization request, + // if one was included there. + RedirectURI *string `form:"redirect_uri" json:"redirect_uri,omitempty"` + + // Resource is an optional absolute URL indicating the target resource server. + Resource *string `form:"resource" json:"resource,omitempty" validate:"omitempty,url"` +} + +// RefreshTokenGrant represents a token request using the Refresh Token flow. +type RefreshTokenGrant struct { + // RefreshToken is the refresh token previously issued to the client. + RefreshToken string `form:"refresh_token" json:"refresh_token" validate:"required"` + + // Scope is an optional space-delimited list of scopes being requested. + // If omitted, the scope is assumed to be identical to the scope originally granted. + Scope *string `form:"scope" json:"scope,omitempty"` + + // Resource is an optional absolute URL indicating the target resource server. + Resource *string `form:"resource" json:"resource,omitempty" validate:"omitempty,url"` +} + +// TokenHandler creates a token endpoint handler with full middleware stack +func TokenHandler(options TokenHandlerOptions) http.HandlerFunc { + // Create the core handler logic + coreHandler := createTokenCoreHandler(options) + + // Apply middlewares in order + var handler http.Handler = coreHandler + + // Apply client authentication middleware + handler = middleware.AuthenticateClient(middleware.ClientAuthenticationMiddlewareOptions{ + ClientsStore: options.Provider.ClientsStore(), + })(handler) + + // Apply rate limiting middleware + limiter := options.RateLimit + if limiter == nil { + // Default rate limiting: 50 requests per 15 minutes + limiter = rate.NewLimiter(rate.Every(15*time.Minute/50), 50) + } + handler = middleware.RateLimitMiddleware(limiter)(handler) + + // Apply method restriction middleware (only POST allowed) + handler = middleware.AllowedMethods([]string{"POST"})(handler) + + // Apply CORS middleware + handler = middleware.CorsMiddleware(handler) + + // Convert http.Handler to http.HandlerFunc + return func(w http.ResponseWriter, r *http.Request) { + handler.ServeHTTP(w, r) + } +} + +// createTokenCoreHandler creates the core token handler logic shared between both versions +func createTokenCoreHandler(options TokenHandlerOptions) http.HandlerFunc { + // Create a validator + validate := validator.New() + + return func(w http.ResponseWriter, r *http.Request) { + // Set cache-control headers + w.Header().Set("Cache-Control", "no-store") + + // Parsing form data + if err := r.ParseForm(); err != nil { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + + errResp := errors.NewOAuthError(errors.ErrInvalidRequest, "Failed to parse form data", "") + json.NewEncoder(w).Encode(errResp.ToResponseStruct()) + return + } + + // Get grant_type from form + grantType := r.FormValue("grant_type") + if grantType == "" { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + + errResp := errors.NewOAuthError(errors.ErrInvalidRequest, "invalid client credentials", "") + json.NewEncoder(w).Encode(errResp.ToResponseStruct()) + return + } + + // Verify basic token request + tokenReq := TokenRequest{ + GrantType: grantType, + } + + if err := validate.Struct(tokenReq); err != nil { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + + errResp := errors.NewOAuthError(errors.ErrInvalidRequest, err.Error(), "") + json.NewEncoder(w).Encode(errResp.ToResponseStruct()) + return + } + + // Check client authentication result + client, ok := middleware.GetAuthenticatedClient(r) + if !ok { + // NOW this code will actually execute because middleware didn't terminate + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusUnauthorized) // Proper OAuth error status + errResp := errors.NewOAuthError(errors.ErrInvalidClient, "invalid client credentials", "") + json.NewEncoder(w).Encode(errResp.ToResponseStruct()) + return + } + + switch grantType { + case "authorization_code": + handleAuthorizationCodeGrant(w, r, validate, options.Provider, *client) + case "refresh_token": + handleRefreshTokenGrant(w, r, validate, options.Provider, *client) + default: + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + + errResp := errors.NewOAuthError(errors.ErrUnsupportedGrantType, "The grant type is not supported by this authorization server.", "") + json.NewEncoder(w).Encode(errResp.ToResponseStruct()) + } + } +} + +// handleAuthorizationCodeGrant processes authorization code grant +func handleAuthorizationCodeGrant(w http.ResponseWriter, r *http.Request, validate *validator.Validate, provider server.OAuthServerProvider, client auth.OAuthClientInformationFull) { + // Parsing the authorization code grant request + var redirectURI *string + if uri := r.FormValue("redirect_uri"); uri != "" { + redirectURI = &uri + } + + var resource *string + if res := r.FormValue("resource"); res != "" { + resource = &res + } + + grant := AuthorizationCodeGrant{ + Code: r.FormValue("code"), + CodeVerifier: r.FormValue("code_verifier"), + RedirectURI: redirectURI, + Resource: resource, + } + + if err := validate.Struct(grant); err != nil { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + + switch verrs := err.(type) { + case validator.ValidationErrors: + for _, fe := range verrs { + if fe.Field() == "Resource" && fe.Tag() == "url" { + errResp := errors.NewOAuthError(errors.ErrInvalidRequest, "resource must be a valid URL", "") + json.NewEncoder(w).Encode(errResp.ToResponseStruct()) + return + } + } + errResp := errors.NewOAuthError(errors.ErrInvalidRequest, err.Error(), "") + json.NewEncoder(w).Encode(errResp.ToResponseStruct()) + return + default: + errResp := errors.NewOAuthError(errors.ErrInvalidRequest, err.Error(), "") + json.NewEncoder(w).Encode(errResp.ToResponseStruct()) + return + } + } + + // Check if the provider supports skipLocalPKceValidation + type skipLocalPKceValidation interface { + GetSkipLocalPkceValidation() bool + } + + skipLocalValidation := false + if p, ok := provider.(skipLocalPKceValidation); ok { + skipLocalValidation = p.GetSkipLocalPkceValidation() + } + + // Perform local PKCE validation unless explicitly skipped + if !skipLocalValidation { + codeChallenge, err := provider.ChallengeForAuthorizationCode(client, grant.Code) + if err != nil { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + + errResp := errors.NewOAuthError(errors.ErrInvalidGrant, "Failed to retrieve code challenge", "") + json.NewEncoder(w).Encode(errResp.ToResponseStruct()) + return + } + + if !pkce.VerifyPKCEChallenge(grant.CodeVerifier, codeChallenge) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + + errResp := errors.NewOAuthError(errors.ErrInvalidGrant, "code_verifier does not match the challenge", "") + json.NewEncoder(w).Encode(errResp.ToResponseStruct()) + return + } + } + + var resourceURL *url.URL + if grant.Resource != nil { + var err error + resourceURL, err = url.Parse(*grant.Resource) + if err != nil { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + + errResp := errors.NewOAuthError(errors.ErrInvalidRequest, "Invalid resource URL", "") + json.NewEncoder(w).Encode(errResp.ToResponseStruct()) + return + } + } + + // Parse the code_verifier to the provider if PKCE validation did not occur locally + var codeVerifier *string + if skipLocalValidation { + codeVerifier = &grant.CodeVerifier + } + + // Exchange the authorization code for a token + tokens, err := provider.ExchangeAuthorizationCode( + client, + grant.Code, + codeVerifier, + grant.RedirectURI, + resourceURL, + ) + + if err != nil { + w.Header().Set("Content-Type", "application/json") + + // Return an appropriate OAuth error response based on the error type + switch { + case err == errors.ErrInvalidParams || err == errors.ErrMissingParams: + w.WriteHeader(http.StatusBadRequest) + errResp := errors.NewOAuthError(errors.ErrInvalidRequest, err.Error(), "") + json.NewEncoder(w).Encode(errResp.ToResponseStruct()) + case err == errors.ErrInvalidJSONRPCParams: + w.WriteHeader(http.StatusBadRequest) + errResp := errors.NewOAuthError(errors.ErrInvalidGrant, err.Error(), "") + json.NewEncoder(w).Encode(errResp.ToResponseStruct()) + default: + w.WriteHeader(http.StatusInternalServerError) + errResp := errors.NewOAuthError(errors.ErrServerError, "Internal Server Error", "") + json.NewEncoder(w).Encode(errResp.ToResponseStruct()) + } + return + } + + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + json.NewEncoder(w).Encode(tokens) +} + +// handleRefreshTokenGrant handles refresh token grant +func handleRefreshTokenGrant(w http.ResponseWriter, r *http.Request, validate *validator.Validate, provider server.OAuthServerProvider, client auth.OAuthClientInformationFull) { + var scope *string + if s := r.FormValue("scope"); s != "" { + scope = &s + } + + var resource *string + if res := r.FormValue("resource"); res != "" { + resource = &res + } + + grant := RefreshTokenGrant{ + RefreshToken: r.FormValue("refresh_token"), + Scope: scope, + Resource: resource, + } + + if err := validate.Struct(grant); err != nil { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + + errResp := errors.NewOAuthError(errors.ErrInvalidRequest, err.Error(), "") + json.NewEncoder(w).Encode(errResp.ToResponseStruct()) + return + } + + // Handle scopes + var scopes []string + if grant.Scope != nil { + scopes = strings.Split(*grant.Scope, " ") + } + + // Handle resource URL + var resourceURL *url.URL + if grant.Resource != nil { + var err error + resourceURL, err = url.Parse(*grant.Resource) + if err != nil { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + + errResp := errors.NewOAuthError(errors.ErrInvalidRequest, "Invalid resource URL", "") + json.NewEncoder(w).Encode(errResp.ToResponseStruct()) + return + } + } + + // Swap refresh token + tokens, err := provider.ExchangeRefreshToken(client, grant.RefreshToken, scopes, resourceURL) + + if err != nil { + w.Header().Set("Content-Type", "application/json") + + if strings.Contains(strings.ToLower(err.Error()), "invalid") { + w.WriteHeader(http.StatusInternalServerError) + errResp := errors.NewOAuthError(errors.ErrInvalidGrant, err.Error(), "") + json.NewEncoder(w).Encode(errResp.ToResponseStruct()) + return + } + + switch { + case err == errors.ErrInvalidParams || err == errors.ErrMissingParams: + w.WriteHeader(http.StatusBadRequest) + errResp := errors.NewOAuthError(errors.ErrInvalidRequest, err.Error(), "") + json.NewEncoder(w).Encode(errResp.ToResponseStruct()) + case err == errors.ErrInvalidJSONRPCParams: + w.WriteHeader(http.StatusBadRequest) + errResp := errors.NewOAuthError(errors.ErrInvalidGrant, err.Error(), "") + json.NewEncoder(w).Encode(errResp.ToResponseStruct()) + default: + w.WriteHeader(http.StatusInternalServerError) + errResp := errors.NewOAuthError(errors.ErrServerError, "Internal Server Error", "") + json.NewEncoder(w).Encode(errResp.ToResponseStruct()) + } + return + } + + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + json.NewEncoder(w).Encode(tokens) +} diff --git a/internal/auth/server/handler/token_test.go b/internal/auth/server/handler/token_test.go new file mode 100644 index 0000000..369f256 --- /dev/null +++ b/internal/auth/server/handler/token_test.go @@ -0,0 +1,697 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package handler + +import ( + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/time/rate" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth/server" + "trpc.group/trpc-go/trpc-mcp-go/internal/errors" +) + +// postFormWithBasicAuth sends a POST request with x-www-form-urlencoded body and HTTP Basic auth +func postFormWithBasicAuth(t *testing.T, h http.Handler, path string, form url.Values, clientID, clientSecret string) *httptest.ResponseRecorder { + t.Helper() + req := httptest.NewRequest(http.MethodPost, path, strings.NewReader(form.Encode())) + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + req.SetBasicAuth(clientID, clientSecret) + + rr := httptest.NewRecorder() + h.ServeHTTP(rr, req) + return rr +} + +// postForm sends a POST request with x-www-form-urlencoded body (no client auth) +func postForm(t *testing.T, h http.Handler, path string, form url.Values) *httptest.ResponseRecorder { + t.Helper() + req := httptest.NewRequest(http.MethodPost, path, strings.NewReader(form.Encode())) + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + rr := httptest.NewRecorder() + h.ServeHTTP(rr, req) + return rr +} + +// postFormWithOrigin sends a POST request with x-www-form-urlencoded body, Basic auth, and Origin header for CORS tests +func postFormWithOrigin(t *testing.T, h http.Handler, path string, form url.Values, clientID, clientSecret, origin string) *httptest.ResponseRecorder { + t.Helper() + req := httptest.NewRequest(http.MethodPost, path, strings.NewReader(form.Encode())) + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + req.Header.Set("Origin", origin) + req.SetBasicAuth(clientID, clientSecret) + + rr := httptest.NewRecorder() + h.ServeHTTP(rr, req) + return rr +} + +// enhancedMockOAuthClientsStore is a simple in-memory clients store used in tests +type enhancedMockOAuthClientsStore struct { + clients map[string]*auth.OAuthClientInformationFull // client_id -> client record +} + +// GetClient returns the client by id or an error if not found. +func (m *enhancedMockOAuthClientsStore) GetClient(clientID string) (*auth.OAuthClientInformationFull, error) { + client, exists := m.clients[clientID] + if !exists { + return nil, fmt.Errorf("client not found") + } + return client, nil +} + +// enhancedMockOAuthServerProvider simulates an OAuth server provider with toggles for different paths +type enhancedMockOAuthServerProvider struct { + clientStore *enhancedMockOAuthClientsStore // backing store for clients + skipLocalPkceValidation bool // when true, PKCE is not validated locally + shouldReturnIdToken bool // when true, adds id_token to token response + shouldFailCodeChallenge bool // when true, ChallengeForAuthorizationCode fails + shouldFailCodeExchange bool // when true, ExchangeAuthorizationCode fails with invalid_grant + shouldFailRefreshExchange bool // when true, ExchangeRefreshToken fails with invalid_grant + supportedScopes []string // list of supported scopes (for tests that depend on scope echo) +} + +// GetSkipLocalPkceValidation exposes whether local PKCE verification should be skipped +func (m *enhancedMockOAuthServerProvider) GetSkipLocalPkceValidation() bool { + return m.skipLocalPkceValidation +} + +// ClientsStore returns a thin adapter around the in-memory store to satisfy the interface +func (m *enhancedMockOAuthServerProvider) ClientsStore() *server.OAuthClientsStore { + return server.NewOAuthClientStore(m.clientStore.GetClient) +} + +// Authorize simulates authorization success by redirecting with code and echoing state +func (m *enhancedMockOAuthServerProvider) Authorize(client auth.OAuthClientInformationFull, params server.AuthorizationParams, res http.ResponseWriter, req *http.Request) error { + // Compose a 302 redirect with code + state + res.Header().Set("Location", "https://redirect-uri.com?code=valid-code&state="+params.State) + res.WriteHeader(http.StatusFound) + return nil +} + +// ChallengeForAuthorizationCode returns a fixed S256 challenge for "valid-code" and errors otherwise +func (m *enhancedMockOAuthServerProvider) ChallengeForAuthorizationCode( + client auth.OAuthClientInformationFull, + authorizationCode string, +) (string, error) { + if m.shouldFailCodeChallenge { + return "", errors.ErrInvalidGrant + } + switch authorizationCode { + case "valid-code": + // Matches code_verifier = "valid-verifier" + return "A_DCKa0ei4rJGhNfKEbwNpiuHzQP7skGQPZ4CBTkJdQ", nil + case "expired-code": + return "", fmt.Errorf("authorization code has expired") + case "invalid-code": + return "", fmt.Errorf("authorization code is invalid") + default: + return "", fmt.Errorf("unknown authorization code") + } +} + +// ExchangeAuthorizationCode returns mock tokens for "valid-code" and errors for others +func (m *enhancedMockOAuthServerProvider) ExchangeAuthorizationCode( + client auth.OAuthClientInformationFull, + authorizationCode string, + codeVerifier *string, + redirectUri *string, + resource *url.URL, +) (*auth.OAuthTokens, error) { + // Simulate upstream failure toggle + if m.shouldFailCodeExchange { + return nil, errors.ErrInvalidGrant + } + + switch authorizationCode { + case "valid-code": + expiresIn := int64(3600) + refreshToken := "mock-refresh-token" + + tokens := &auth.OAuthTokens{ + AccessToken: "mock-access-token", + TokenType: "bearer", + ExpiresIn: &expiresIn, + RefreshToken: &refreshToken, + } + // Optionally attach an ID token for OIDC scenarios + if m.shouldReturnIdToken { + idToken := "mock-id-token" + tokens.IDToken = &idToken + } + return tokens, nil + case "expired-code": + return nil, fmt.Errorf("authorization code has expired") + case "invalid-code": + return nil, fmt.Errorf("authorization code is invalid") + default: + return nil, errors.ErrInvalidGrant + } +} + +// ExchangeRefreshToken returns a new access/refresh token pair for a valid refresh token +func (m *enhancedMockOAuthServerProvider) ExchangeRefreshToken( + client auth.OAuthClientInformationFull, + refreshToken string, + scopes []string, + resource *url.URL, +) (*auth.OAuthTokens, error) { + if m.shouldFailRefreshExchange { + return nil, errors.ErrInvalidGrant + } + + switch refreshToken { + case "valid-refresh-token": + expiresIn := int64(3600) + newRefreshToken := "new-mock-refresh-token" + + tokens := &auth.OAuthTokens{ + AccessToken: "new-mock-access-token", + TokenType: "bearer", + ExpiresIn: &expiresIn, + RefreshToken: &newRefreshToken, + } + if len(scopes) > 0 { + scope := strings.Join(scopes, " ") + tokens.Scope = &scope + } + return tokens, nil + case "invalid-refresh-token": + return nil, fmt.Errorf("refresh token is invalid") + default: + return nil, errors.ErrInvalidGrant + } +} + +// VerifyAccessToken returns a stubbed AuthInfo when token == "valid-token" +func (m *enhancedMockOAuthServerProvider) VerifyAccessToken(token string) (*server.AuthInfo, error) { + if token == "valid-token" { + return &server.AuthInfo{ + ClientID: "valid-client-id", + Scopes: []string{"read", "write"}, + }, nil + } + return nil, fmt.Errorf("invalid token") +} + +// RevokeToken is a no-op in this mock; revocation success is implied +func (m *enhancedMockOAuthServerProvider) RevokeToken(client auth.OAuthClientInformationFull, request auth.OAuthTokenRevocationRequest) error { + return nil +} + +// createMockClient builds a confidential client (client_secret_basic) with a fixed secret +func createMockClient(id string) *auth.OAuthClientInformationFull { + return &auth.OAuthClientInformationFull{ + OAuthClientInformation: auth.OAuthClientInformation{ + ClientID: id, + ClientSecret: "valid-secret", + }, + OAuthClientMetadata: auth.OAuthClientMetadata{ + RedirectURIs: []string{"https://example.com/callback"}, + TokenEndpointAuthMethod: "client_secret_basic", + }, + } +} + +// createEnhancedMockProvider wires the in-memory clients store and returns a preconfigured provider +func createEnhancedMockProvider() *enhancedMockOAuthServerProvider { + clients := make(map[string]*auth.OAuthClientInformationFull) + clients["valid-client"] = createMockClient("valid-client") + store := &enhancedMockOAuthClientsStore{clients: clients} + + return &enhancedMockOAuthServerProvider{ + clientStore: store, + supportedScopes: []string{"read", "write", "profile", "email"}, + } +} + +func TestToken_RequiresPostMethod(t *testing.T) { + provider := createEnhancedMockProvider() + handler := TokenHandler(TokenHandlerOptions{ + Provider: provider, + RateLimit: rate.NewLimiter(rate.Every(15*time.Minute/50), 50), + }) + + req := httptest.NewRequest(http.MethodGet, "/token", nil) + rr := httptest.NewRecorder() + handler.ServeHTTP(rr, req) + + assert.Equal(t, http.StatusMethodNotAllowed, rr.Code) + assert.Equal(t, "POST", rr.Header().Get("Allow")) +} + +func TestToken_RequiresGrantType(t *testing.T) { + provider := createEnhancedMockProvider() + handler := TokenHandler(TokenHandlerOptions{ + Provider: provider, + RateLimit: rate.NewLimiter(rate.Every(15*time.Minute/50), 50), + }) + + form := url.Values{} + // Missing grant_type + + rr := postFormWithBasicAuth(t, handler, "/token", form, "valid-client", "valid-secret") + + assert.Equal(t, http.StatusBadRequest, rr.Code) + var errResp map[string]interface{} + require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &errResp)) + assert.Equal(t, "invalid_request", errResp["error"]) +} + +func TestToken_RejectsUnsupportedGrantTypes(t *testing.T) { + provider := createEnhancedMockProvider() + handler := TokenHandler(TokenHandlerOptions{ + Provider: provider, + RateLimit: rate.NewLimiter(rate.Every(15*time.Minute/50), 50), + }) + + form := url.Values{ + "grant_type": {"password"}, // Unsupported grant type + } + + rr := postFormWithBasicAuth(t, handler, "/token", form, "valid-client", "valid-secret") + + assert.Equal(t, http.StatusBadRequest, rr.Code) + var errResp map[string]interface{} + require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &errResp)) + assert.Equal(t, "unsupported_grant_type", errResp["error"]) + assert.Equal(t, "The grant type is not supported by this authorization server.", errResp["error_description"]) +} + +func TestToken_RequiresValidClientCredentials_CurrentBehavior(t *testing.T) { + provider := createEnhancedMockProvider() + handler := TokenHandler(TokenHandlerOptions{ + Provider: provider, + RateLimit: rate.NewLimiter(rate.Every(15*time.Minute/50), 50), + }) + + form := url.Values{ + "grant_type": {"authorization_code"}, + } + + rr := postFormWithBasicAuth(t, handler, "/token", form, "invalid-client", "wrong-secret") + + // HTTP status should be 401 + assert.Equal(t, http.StatusUnauthorized, rr.Code) + + var errResp map[string]interface{} + require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &errResp)) + + assert.Equal(t, "invalid_client", errResp["error"]) + assert.Contains(t, errResp["error_description"].(string), "client") +} + +func TestToken_AcceptsValidClientCredentials(t *testing.T) { + provider := createEnhancedMockProvider() + handler := TokenHandler(TokenHandlerOptions{ + Provider: provider, + RateLimit: rate.NewLimiter(rate.Every(15*time.Minute/50), 50), + }) + + form := url.Values{ + "grant_type": {"authorization_code"}, + "code": {"valid-code"}, + "code_verifier": {"valid-verifier"}, + } + + rr := postFormWithBasicAuth(t, handler, "/token", form, "valid-client", "valid-secret") + + assert.Equal(t, http.StatusOK, rr.Code) +} + +func TestToken_AuthorizationCode_RequiresCodeParameter(t *testing.T) { + provider := createEnhancedMockProvider() + handler := TokenHandler(TokenHandlerOptions{ + Provider: provider, + RateLimit: rate.NewLimiter(rate.Every(15*time.Minute/50), 50), + }) + + form := url.Values{ + "grant_type": {"authorization_code"}, + // Missing code + "code_verifier": {"valid-verifier"}, + } + + rr := postFormWithBasicAuth(t, handler, "/token", form, "valid-client", "valid-secret") + + assert.Equal(t, http.StatusBadRequest, rr.Code) + var errResp map[string]interface{} + require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &errResp)) + assert.Equal(t, "invalid_request", errResp["error"]) +} + +func TestToken_AuthorizationCode_RequiresCodeVerifierParameter(t *testing.T) { + provider := createEnhancedMockProvider() + handler := TokenHandler(TokenHandlerOptions{ + Provider: provider, + RateLimit: rate.NewLimiter(rate.Every(15*time.Minute/50), 50), + }) + + form := url.Values{ + "grant_type": {"authorization_code"}, + "code": {"valid-code"}, + // Missing code_verifier + } + + rr := postFormWithBasicAuth(t, handler, "/token", form, "valid-client", "valid-secret") + + assert.Equal(t, http.StatusBadRequest, rr.Code) + var errResp map[string]interface{} + require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &errResp)) + assert.Equal(t, "invalid_request", errResp["error"]) +} + +func TestToken_AuthorizationCode_VerifiesPKCEChallenge(t *testing.T) { + provider := createEnhancedMockProvider() + provider.shouldFailCodeChallenge = false // Ensure challenge retrieval works + + handler := TokenHandler(TokenHandlerOptions{ + Provider: provider, + RateLimit: rate.NewLimiter(rate.Every(15*time.Minute/50), 50), + }) + + form := url.Values{ + "grant_type": {"authorization_code"}, + "code": {"valid-code"}, + "code_verifier": {"invalid-verifier"}, // This won't match the challenge + } + + rr := postFormWithBasicAuth(t, handler, "/token", form, "valid-client", "valid-secret") + + assert.Equal(t, http.StatusBadRequest, rr.Code) + var errResp map[string]interface{} + require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &errResp)) + assert.Equal(t, "invalid_grant", errResp["error"]) + assert.Contains(t, errResp["error_description"], "code_verifier") +} + +func TestToken_AuthorizationCode_RejectsExpiredCode(t *testing.T) { + provider := createEnhancedMockProvider() + handler := TokenHandler(TokenHandlerOptions{ + Provider: provider, + RateLimit: rate.NewLimiter(rate.Every(15*time.Minute/50), 50), + }) + + form := url.Values{ + "grant_type": {"authorization_code"}, + "code": {"expired-code"}, + "code_verifier": {"valid-verifier"}, + } + + rr := postFormWithBasicAuth(t, handler, "/token", form, "valid-client", "valid-secret") + + assert.Equal(t, http.StatusBadRequest, rr.Code) + var errResp map[string]interface{} + require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &errResp)) + assert.Equal(t, "invalid_grant", errResp["error"]) +} + +func TestToken_AuthorizationCode_RejectsInvalidCode(t *testing.T) { + provider := createEnhancedMockProvider() + handler := TokenHandler(TokenHandlerOptions{ + Provider: provider, + RateLimit: rate.NewLimiter(rate.Every(15*time.Minute/50), 50), + }) + + form := url.Values{ + "grant_type": {"authorization_code"}, + "code": {"invalid-code"}, + "code_verifier": {"valid-verifier"}, + } + + rr := postFormWithBasicAuth(t, handler, "/token", form, "valid-client", "valid-secret") + + assert.Equal(t, http.StatusBadRequest, rr.Code) + var errResp map[string]interface{} + require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &errResp)) + assert.Equal(t, "invalid_grant", errResp["error"]) +} + +func TestToken_AuthorizationCode_ReturnsTokensForValidExchange(t *testing.T) { + provider := createEnhancedMockProvider() + handler := TokenHandler(TokenHandlerOptions{ + Provider: provider, + RateLimit: rate.NewLimiter(rate.Every(15*time.Minute/50), 50), + }) + + form := url.Values{ + "grant_type": {"authorization_code"}, + "code": {"valid-code"}, + "code_verifier": {"valid-verifier"}, + "resource": {"https://api.example.com/resource"}, + } + + rr := postFormWithBasicAuth(t, handler, "/token", form, "valid-client", "valid-secret") + + assert.Equal(t, http.StatusOK, rr.Code) + + var tokens map[string]interface{} + require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &tokens)) + assert.Equal(t, "mock-access-token", tokens["access_token"]) + assert.Equal(t, "bearer", tokens["token_type"]) + assert.Equal(t, float64(3600), tokens["expires_in"]) + assert.Equal(t, "mock-refresh-token", tokens["refresh_token"]) +} + +func TestToken_AuthorizationCode_ReturnsIdTokenWhenProvided(t *testing.T) { + provider := createEnhancedMockProvider() + provider.shouldReturnIdToken = true + + handler := TokenHandler(TokenHandlerOptions{ + Provider: provider, + RateLimit: rate.NewLimiter(rate.Every(15*time.Minute/50), 50), + }) + + form := url.Values{ + "grant_type": {"authorization_code"}, + "code": {"valid-code"}, + "code_verifier": {"valid-verifier"}, + } + + rr := postFormWithBasicAuth(t, handler, "/token", form, "valid-client", "valid-secret") + + assert.Equal(t, http.StatusOK, rr.Code) + + var tokens map[string]interface{} + require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &tokens)) + assert.Equal(t, "mock-id-token", tokens["id_token"]) +} + +func TestToken_RefreshToken_RequiresRefreshTokenParameter(t *testing.T) { + provider := createEnhancedMockProvider() + handler := TokenHandler(TokenHandlerOptions{ + Provider: provider, + RateLimit: rate.NewLimiter(rate.Every(15*time.Minute/50), 50), + }) + + form := url.Values{ + "grant_type": {"refresh_token"}, + // Missing refresh_token + } + + rr := postFormWithBasicAuth(t, handler, "/token", form, "valid-client", "valid-secret") + + assert.Equal(t, http.StatusBadRequest, rr.Code) + var errResp map[string]interface{} + require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &errResp)) + assert.Equal(t, "invalid_request", errResp["error"]) +} + +func TestToken_RefreshToken_RejectsInvalidRefreshToken(t *testing.T) { + provider := createEnhancedMockProvider() + handler := TokenHandler(TokenHandlerOptions{ + Provider: provider, + RateLimit: rate.NewLimiter(rate.Every(15*time.Minute/50), 50), + }) + + form := url.Values{ + "grant_type": {"refresh_token"}, + "refresh_token": {"invalid-refresh-token"}, + } + + rr := postFormWithBasicAuth(t, handler, "/token", form, "valid-client", "valid-secret") + + assert.Equal(t, http.StatusInternalServerError, rr.Code) + var errResp map[string]interface{} + require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &errResp)) + assert.Equal(t, "invalid_grant", errResp["error"]) +} + +func TestToken_RefreshToken_ReturnsNewTokensForValidRefresh(t *testing.T) { + provider := createEnhancedMockProvider() + handler := TokenHandler(TokenHandlerOptions{ + Provider: provider, + RateLimit: rate.NewLimiter(rate.Every(15*time.Minute/50), 50), + }) + + form := url.Values{ + "grant_type": {"refresh_token"}, + "refresh_token": {"valid-refresh-token"}, + "resource": {"https://api.example.com/resource"}, + } + + rr := postFormWithBasicAuth(t, handler, "/token", form, "valid-client", "valid-secret") + + assert.Equal(t, http.StatusOK, rr.Code) + + var tokens map[string]interface{} + require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &tokens)) + assert.Equal(t, "new-mock-access-token", tokens["access_token"]) + assert.Equal(t, "bearer", tokens["token_type"]) + assert.Equal(t, float64(3600), tokens["expires_in"]) + assert.Equal(t, "new-mock-refresh-token", tokens["refresh_token"]) +} + +func TestToken_RefreshToken_RespectsRequestedScopes(t *testing.T) { + provider := createEnhancedMockProvider() + handler := TokenHandler(TokenHandlerOptions{ + Provider: provider, + RateLimit: rate.NewLimiter(rate.Every(15*time.Minute/50), 50), + }) + + form := url.Values{ + "grant_type": {"refresh_token"}, + "refresh_token": {"valid-refresh-token"}, + "scope": {"profile email"}, + } + + rr := postFormWithBasicAuth(t, handler, "/token", form, "valid-client", "valid-secret") + + assert.Equal(t, http.StatusOK, rr.Code) + + var tokens map[string]interface{} + require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &tokens)) + assert.Equal(t, "profile email", tokens["scope"]) +} + +func TestToken_IncludesCORSHeaders(t *testing.T) { + provider := createEnhancedMockProvider() + handler := TokenHandler(TokenHandlerOptions{ + Provider: provider, + RateLimit: rate.NewLimiter(rate.Every(15*time.Minute/50), 50), + }) + + form := url.Values{ + "grant_type": {"authorization_code"}, + "code": {"valid-code"}, + "code_verifier": {"valid-verifier"}, + } + + rr := postFormWithOrigin(t, handler, "/token", form, "valid-client", "valid-secret", "https://example.com") + + assert.Equal(t, http.StatusOK, rr.Code) + + // Check CORS headers + assert.Contains(t, rr.Header().Get("Access-Control-Allow-Origin"), "*") +} + +func TestToken_RateLimiting(t *testing.T) { + provider := createEnhancedMockProvider() + handler := TokenHandler(TokenHandlerOptions{ + Provider: provider, + RateLimit: rate.NewLimiter(rate.Every(15*time.Minute/50), 1), // Only 1 request allowed + }) + + form := url.Values{ + "grant_type": {"authorization_code"}, + "code": {"valid-code"}, + "code_verifier": {"valid-verifier"}, + } + + // First request should succeed + rr1 := postFormWithBasicAuth(t, handler, "/token", form, "valid-client", "valid-secret") + assert.Equal(t, http.StatusOK, rr1.Code) + + // Second request should be rate-limited + rr2 := postFormWithBasicAuth(t, handler, "/token", form, "valid-client", "valid-secret") + assert.Equal(t, http.StatusTooManyRequests, rr2.Code) +} + +func TestToken_SetsCacheControlHeaders(t *testing.T) { + provider := createEnhancedMockProvider() + handler := TokenHandler(TokenHandlerOptions{ + Provider: provider, + RateLimit: rate.NewLimiter(rate.Every(15*time.Minute/50), 50), + }) + + form := url.Values{ + "grant_type": {"authorization_code"}, + "code": {"valid-code"}, + "code_verifier": {"valid-verifier"}, + } + + rr := postFormWithBasicAuth(t, handler, "/token", form, "valid-client", "valid-secret") + + assert.Equal(t, "no-store", rr.Header().Get("Cache-Control")) + assert.Equal(t, "application/json", rr.Header().Get("Content-Type")) +} + +func TestToken_RejectsOPTIONSMethod(t *testing.T) { + provider := createEnhancedMockProvider() + handler := TokenHandler(TokenHandlerOptions{ + Provider: provider, + RateLimit: rate.NewLimiter(rate.Every(15*time.Minute/50), 50), + }) + + req := httptest.NewRequest(http.MethodOptions, "/token", nil) + rr := httptest.NewRecorder() + handler.ServeHTTP(rr, req) + + assert.Equal(t, http.StatusMethodNotAllowed, rr.Code) +} + +func TestToken_ValidatesResourceParameter(t *testing.T) { + provider := createEnhancedMockProvider() + handler := TokenHandler(TokenHandlerOptions{ + Provider: provider, + RateLimit: rate.NewLimiter(rate.Every(15*time.Minute/50), 50), + }) + + form := url.Values{ + "grant_type": {"authorization_code"}, + "code": {"valid-code"}, + "code_verifier": {"valid-verifier"}, + "resource": {"invalid-url"}, // Invalid URL + } + + rr := postFormWithBasicAuth(t, handler, "/token", form, "valid-client", "valid-secret") + + assert.Equal(t, http.StatusBadRequest, rr.Code) + var errResp map[string]interface{} + require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &errResp)) + assert.Equal(t, "invalid_request", errResp["error"]) + assert.Contains(t, errResp["error_description"], "resource") +} + +func TestToken_SkipLocalPKCEValidation(t *testing.T) { + provider := createEnhancedMockProvider() + provider.skipLocalPkceValidation = true + + handler := TokenHandler(TokenHandlerOptions{ + Provider: provider, + RateLimit: rate.NewLimiter(rate.Every(15*time.Minute/50), 50), + }) + + form := url.Values{ + "grant_type": {"authorization_code"}, + "code": {"valid-code"}, + "code_verifier": {"any-verifier"}, // Should be passed through without local validation + } + + rr := postFormWithBasicAuth(t, handler, "/token", form, "valid-client", "valid-secret") + + assert.Equal(t, http.StatusOK, rr.Code) +} diff --git a/internal/auth/server/http.go b/internal/auth/server/http.go new file mode 100644 index 0000000..18d59f0 --- /dev/null +++ b/internal/auth/server/http.go @@ -0,0 +1,134 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package server + +import ( + "context" + "errors" + "fmt" + "net/http" + + oauthErrors "trpc.group/trpc-go/trpc-mcp-go/internal/errors" +) + +type ctxKey int +type ctxKeyScope int + +const ( + ctxKeyAuthInfo ctxKey = iota + ctxKeyAuthErr + ctxKeyRequiredScope ctxKeyScope = 1 +) + +// WithAuthInfo writes authentication information into the context +func WithAuthInfo(ctx context.Context, info *AuthInfo) context.Context { + // If no info provided, just return the original context + if info == nil { + return ctx + } + // Store authentication info in the context with a private key + return context.WithValue(ctx, ctxKeyAuthInfo, info) +} + +// GetAuthInfo retrieves authentication information from the context +func GetAuthInfo(ctx context.Context) (*AuthInfo, bool) { + // Extract value from context + v := ctx.Value(ctxKeyAuthInfo) + // Return false if not set + if v == nil { + return nil, false + } + // Type assert to *AuthInfo + info, ok := v.(*AuthInfo) + return info, ok && info != nil +} + +// WithAuthErr stores an authentication error into the context +func WithAuthErr(ctx context.Context, err error) context.Context { + // No error means no need to wrap context + if err == nil { + return ctx + } + // Store error in context for later retrieval + return context.WithValue(ctx, ctxKeyAuthErr, err) +} + +// GetAuthErr retrieves an authentication error from the context +func GetAuthErr(ctx context.Context) error { + // Extract error from context + v := ctx.Value(ctxKeyAuthErr) + // If not set, return nil (no error) + if v == nil { + return nil + } + // Type assert to error + if err, ok := v.(error); ok { + return err + } + // Fallback: return generic error if type is invalid + return errors.New("auth err") +} + +// WriteAuthChallenge writes a Bearer authentication challenge response header +// and sends the specified HTTP status code +func WriteAuthChallenge(w http.ResponseWriter, status int, code, desc, scope string) { + // Build WWW-Authenticate header per RFC 6750 + val := fmt.Sprintf(`Bearer realm="mcp", error="%s", error_description="%s"`, code, desc) + // Append required scope if present + if scope != "" { + val += fmt.Sprintf(`, scope="%s"`, scope) + } + // Add headers to prevent caching + w.Header().Set("WWW-Authenticate", val) + w.Header().Set("Cache-Control", "no-store") + w.Header().Set("Pragma", "no-cache") + // Return error with HTTP status + http.Error(w, http.StatusText(status), status) +} + +// DetermineAuthError maps OAuth error types to HTTP status codes +// and standardized error code/description strings +func DetermineAuthError(err error) (int, string, string) { + // Match specific OAuth error cases + switch { + case errors.Is(err, oauthErrors.ErrInsufficientScope): + // Missing required scope: 403 Forbidden + return http.StatusForbidden, "insufficient_scope", "Token lacks required scope" + case errors.Is(err, oauthErrors.ErrInvalidRequest): + // Invalid/malformed request: 400 Bad Request + return http.StatusBadRequest, "invalid_request", "Missing or malformed authorization" + case errors.Is(err, oauthErrors.ErrInvalidClient): + // Invalid client authentication: 401 Unauthorized + return http.StatusUnauthorized, "invalid_client", "Client authentication failed" + case errors.Is(err, oauthErrors.ErrInvalidToken): + // Invalid/expired token: 401 Unauthorized + return http.StatusUnauthorized, "invalid_token", "The access token is invalid or expired" + default: + // Default: classify as invalid_token for security + return http.StatusUnauthorized, "invalid_token", "Token verification failed" + } +} + +// WithRequiredScope stores the minimum required scope into the context +func WithRequiredScope(ctx context.Context, scope string) context.Context { + // No scope provided, no modification + if scope == "" { + return ctx + } + // Store required scope for later validation + return context.WithValue(ctx, ctxKeyRequiredScope, scope) +} + +// GetRequiredScope retrieves the required scope from the context +func GetRequiredScope(ctx context.Context) (string, bool) { + // Extract scope value + v := ctx.Value(ctxKeyRequiredScope) + // Assert type to string + s, ok := v.(string) + // Return non-empty scope if valid + return s, ok && s != "" +} diff --git a/internal/auth/server/http_test.go b/internal/auth/server/http_test.go new file mode 100644 index 0000000..e6ac2d2 --- /dev/null +++ b/internal/auth/server/http_test.go @@ -0,0 +1,90 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package server + +import ( + "net/url" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestAuthInfoFields(t *testing.T) { + // Setting authInfo data + token := "mock-token" + clientID := "client-123" + scopes := []string{"read", "write"} + expiresAt := int64(1609459200) // 2021-01-01 00:00:00 + resource, err := url.Parse("https://example.com/resource") + require.NoError(t, err) + extra := map[string]interface{}{"key1": "value1"} + + authInfo := &AuthInfo{ + Token: token, + ClientID: clientID, + Scopes: scopes, + ExpiresAt: &expiresAt, + Resource: resource, + Extra: extra, + } + + // Verify that the stored data is correct + assert.Equal(t, token, authInfo.Token) + assert.Equal(t, clientID, authInfo.ClientID) + assert.ElementsMatch(t, scopes, authInfo.Scopes) + assert.Equal(t, expiresAt, *authInfo.ExpiresAt) + assert.Equal(t, resource, authInfo.Resource) + assert.Equal(t, extra, authInfo.Extra) +} + +func TestAuthInfoWithNilExpiresAt(t *testing.T) { + authInfo := &AuthInfo{ + Token: "mock-token", + ClientID: "client-123", + Scopes: []string{"read", "write"}, + ExpiresAt: nil, // nil expiresAt + } + + // Verify that ExpiresAt is nil + assert.Nil(t, authInfo.ExpiresAt) +} + +func TestAuthInfoResourceValidation(t *testing.T) { + validURL, err := url.Parse("https://example.com/resource") + require.NoError(t, err) + invalidURL, err := url.Parse("https://example.com/invalid-resource") + require.NoError(t, err) + + authInfo := &AuthInfo{ + Token: "mock-token", + ClientID: "client-123", + Scopes: []string{"read"}, + Resource: validURL, + } + + // Verify that resource match + assert.Equal(t, validURL.String(), authInfo.Resource.String()) + + // Setting Resource to an invalid URL + authInfo.Resource = invalidURL + assert.Equal(t, invalidURL.String(), authInfo.Resource.String()) +} + +func TestAuthInfoExtraData(t *testing.T) { + extraData := map[string]interface{}{"key1": "value1", "key2": 1234} + + authInfo := &AuthInfo{ + Token: "mock-token", + ClientID: "client-123", + Scopes: []string{"read"}, + Extra: extraData, + } + + // Verify that Extra data is stored correctly + assert.Equal(t, extraData, authInfo.Extra) +} diff --git a/internal/auth/server/middleware/allowedMethods.go b/internal/auth/server/middleware/allowedMethods.go new file mode 100644 index 0000000..ee1d1d2 --- /dev/null +++ b/internal/auth/server/middleware/allowedMethods.go @@ -0,0 +1,48 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package middleware + +import ( + "encoding/json" + "fmt" + "net/http" + "strings" + + "trpc.group/trpc-go/trpc-mcp-go/internal/errors" +) + +// AllowedMethods returns a middleware that permits only the provided HTTP methods +// If the request method is not allowed it responds with 405 Method Not Allowed +// The response includes an Allow header and a JSON OAuth error body +func AllowedMethods(methods []string) func(http.Handler) http.Handler { + return func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Allow request to proceed when method matches one of the allowed methods + for _, method := range methods { + if r.Method == method { + next.ServeHTTP(w, r) + return + } + } + + // Build 405 response with Allow header listing permitted methods + w.Header().Set("Allow", strings.Join(methods, ", ")) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusMethodNotAllowed) + + // Create an OAuth-style error payload for method not allowed + oauthErr := errors.NewOAuthError( + errors.ErrMethodNotAllowed, + fmt.Sprintf("The method %s is not allowed for this endpoint", r.Method), + "", // Optional error URI + ) + + // Encode the error as JSON response body + _ = json.NewEncoder(w).Encode(oauthErr.ToResponseStruct()) + }) + } +} diff --git a/internal/auth/server/middleware/allowedMethods_test.go b/internal/auth/server/middleware/allowedMethods_test.go new file mode 100644 index 0000000..44a5c49 --- /dev/null +++ b/internal/auth/server/middleware/allowedMethods_test.go @@ -0,0 +1,130 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package middleware + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +func TestAllowedMethods(t *testing.T) { + // Create a test handler + createTestHandler := func() http.Handler { + mux := http.NewServeMux() + // Define /test route that only supports GET + mux.Handle("/test", AllowedMethods([]string{"GET"})(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("GET success")) + }))) + return mux + } + + // Case 1 allows specified HTTP method + t.Run("allows specified HTTP method", func(t *testing.T) { + handler := createTestHandler() + req := httptest.NewRequest(http.MethodGet, "/test", nil) + rr := httptest.NewRecorder() + handler.ServeHTTP(rr, req) + + if rr.Code != http.StatusOK { + t.Errorf("expected status %d, got %d", http.StatusOK, rr.Code) + } + if body := rr.Body.String(); body != "GET success" { + t.Errorf("expected body %q, got %q", "GET success", body) + } + }) + + // Case 2 returns 405 for unsupported methods + t.Run("returns 405 for unspecified HTTP methods", func(t *testing.T) { + methods := []string{"POST", "PUT", "DELETE", "PATCH"} + + for _, method := range methods { + t.Run(method, func(t *testing.T) { + handler := createTestHandler() + req := httptest.NewRequest(method, "/test", nil) + rr := httptest.NewRecorder() + handler.ServeHTTP(rr, req) + + if rr.Code != http.StatusMethodNotAllowed { + t.Errorf("expected status %d, got %d", http.StatusMethodNotAllowed, rr.Code) + } + + var response map[string]string + if err := json.NewDecoder(rr.Body).Decode(&response); err != nil { + t.Fatalf("failed to decode response: %v", err) + } + + expected := map[string]string{ + "error": "method not allowed", + "error_description": "The method " + method + " is not allowed for this endpoint", + } + if response["error"] == expected["error"] && response["error_description"] == expected["error_description"] { + t.Errorf("expected response %v, got %v", expected, response) + } + }) + } + }) + + // Case 3 checks Allow header + t.Run("includes Allow header with specified methods", func(t *testing.T) { + handler := createTestHandler() + req := httptest.NewRequest(http.MethodPost, "/test", nil) + rr := httptest.NewRecorder() + handler.ServeHTTP(rr, req) + + if allow := rr.Header().Get("Allow"); allow != "GET" { + t.Errorf("expected Allow header %q, got %q", "GET", allow) + } + }) + + // Case 4 supports multiple allowed methods + t.Run("works with multiple allowed methods", func(t *testing.T) { + // Define /multi route supporting GET and POST + mux := http.NewServeMux() + mux.Handle("/multi", AllowedMethods([]string{"GET", "POST"})(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodGet { + _, _ = w.Write([]byte("GET")) + } else if r.Method == http.MethodPost { + _, _ = w.Write([]byte("POST")) + } + }))) + + // Allowed methods + for _, method := range []string{http.MethodGet, http.MethodPost} { + t.Run(method, func(t *testing.T) { + req := httptest.NewRequest(method, "/multi", nil) + rr := httptest.NewRecorder() + mux.ServeHTTP(rr, req) + + if rr.Code != http.StatusOK { + t.Errorf("expected status %d, got %d", http.StatusOK, rr.Code) + } + expectedBody := strings.ToUpper(method) + if body := rr.Body.String(); body != expectedBody { + t.Errorf("expected body %q, got %q", expectedBody, body) + } + }) + } + + // Unsupported method PUT + t.Run("PUT", func(t *testing.T) { + req := httptest.NewRequest(http.MethodPut, "/multi", nil) + rr := httptest.NewRecorder() + mux.ServeHTTP(rr, req) + + if rr.Code != http.StatusMethodNotAllowed { + t.Errorf("expected status %d, got %d", http.StatusMethodNotAllowed, rr.Code) + } + if allow := rr.Header().Get("Allow"); allow != "GET, POST" { + t.Errorf("expected Allow header %q, got %q", "GET, POST", allow) + } + }) + }) +} diff --git a/internal/auth/server/middleware/audit.go b/internal/auth/server/middleware/audit.go new file mode 100644 index 0000000..9abd5d3 --- /dev/null +++ b/internal/auth/server/middleware/audit.go @@ -0,0 +1,768 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package middleware + +import ( + "bytes" + "context" + "crypto/rand" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "fmt" + "io" + "net/http" + "regexp" + "strings" + "time" + + "go.uber.org/zap" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth/server" +) + +// AuditLevel defines audit log verbosity levels +type AuditLevel int + +const ( + AuditLevelNone AuditLevel = iota + AuditLevelBasic + AuditLevelDetailed + AuditLevelFull +) + +// AuditEvent represents an OAuth 2.1 operation audit record +type AuditEvent struct { + EventID string `json:"event_id"` + Timestamp time.Time `json:"timestamp"` + EventType string `json:"event_type"` + AuditLevel AuditLevel `json:"audit_level"` + Method string `json:"method"` + Path string `json:"path"` + QueryParams map[string]string `json:"query_params,omitempty"` + Headers map[string]string `json:"headers,omitempty"` + RemoteAddr string `json:"remote_addr"` + UserAgent string `json:"user_agent"` + RequestID string `json:"request_id,omitempty"` + ClientID string `json:"client_id,omitempty"` + Subject string `json:"subject,omitempty"` + Scopes []string `json:"scopes,omitempty"` + GrantType string `json:"grant_type,omitempty"` + ResponseType string `json:"response_type,omitempty"` + RedirectURI string `json:"redirect_uri,omitempty"` + Resource string `json:"resource,omitempty"` + StatusCode int `json:"status_code"` + ResponseTime time.Duration `json:"response_time"` + ErrorCode string `json:"error_code,omitempty"` + ErrorMessage string `json:"error_message,omitempty"` + TokenHash string `json:"token_hash,omitempty"` + CodeHash string `json:"code_hash,omitempty"` + IPHash string `json:"ip_hash,omitempty"` + Metadata map[string]interface{} `json:"metadata,omitempty"` + RiskLevel string `json:"risk_level,omitempty"` + RiskFactors []string `json:"risk_factors,omitempty"` + RequestBody string `json:"request_body,omitempty"` + ResponseBody string `json:"response_body,omitempty"` +} + +// AuditLogger defines an interface for emitting audit logs +type AuditLogger interface { + LogEvent(event AuditEvent) error + LogError(event AuditEvent, err error) error +} + +// DefaultAuditLogger provides a zap based implementation of AuditLogger +type DefaultAuditLogger struct { + logger *zap.Logger +} + +// NewAuditLogger creates a DefaultAuditLogger using the provided zap logger or sensible defaults +func NewAuditLogger(logger *zap.Logger) *DefaultAuditLogger { + // Build a production logger by default and fall back to development if needed + if logger == nil { + var err error + logger, err = zap.NewProduction() + if err != nil { + logger, _ = zap.NewDevelopment() + } + } + return &DefaultAuditLogger{logger: logger} +} + +// GetZapLogger exposes the underlying zap logger for advanced usage +func (l *DefaultAuditLogger) GetZapLogger() *zap.Logger { + return l.logger +} + +// LogEvent writes a structured audit event at info level +func (l *DefaultAuditLogger) LogEvent(event AuditEvent) error { + // Guard against uninitialized logger + if l.logger == nil { + return fmt.Errorf("zap logger not initialized") + } + + // Marshal full event payload for a single structured field + data, err := json.Marshal(event) + if err != nil { + return fmt.Errorf("failed to marshal audit event: %w", err) + } + + // Emit event with a compact summary sub-structure for quick filtering + l.logger.Info("[AUDIT]", + zap.ByteString("event", data), + zap.Any("audit", struct { + Method string + Path string + StatusCode int + ResponseTime time.Duration + ClientID string + Subject string + Scopes []string + RiskLevel string + }{ + Method: event.Method, + Path: event.Path, + StatusCode: event.StatusCode, + ResponseTime: event.ResponseTime, + ClientID: event.ClientID, + Subject: event.Subject, + Scopes: event.Scopes, + RiskLevel: event.RiskLevel, + }), + ) + return nil +} + +// LogError writes an audit event including the provided error message +func (l *DefaultAuditLogger) LogError(event AuditEvent, err error) error { + // Attach error message then delegate to LogEvent + event.ErrorMessage = err.Error() + return l.LogEvent(event) +} + +// AuditMiddlewareOptions configures what and how the middleware audits +type AuditMiddlewareOptions struct { + Logger AuditLogger + Level AuditLevel + HashSensitiveData bool + IncludeRequestBody bool + IncludeResponseBody bool + RiskAssessor func(AuditEvent) (string, []string) + MetadataExtractor func(*http.Request) map[string]interface{} + EndpointPatterns []string + ExcludePatterns []string + SensitiveKeys []string +} + +// DefaultAuditMiddlewareOptions returns a sane default configuration for OAuth endpoints +func DefaultAuditMiddlewareOptions() *AuditMiddlewareOptions { + // Default to detailed level with hashing and common OAuth endpoint patterns + return &AuditMiddlewareOptions{ + Logger: NewAuditLogger(nil), + Level: AuditLevelDetailed, + HashSensitiveData: true, + EndpointPatterns: []string{ + "/oauth2/authorize", + "/oauth2/token", + "/oauth2/revoke", + "/oauth2/register", + "/oauth2/metadata", + }, + SensitiveKeys: []string{"client_secret", "code_verifier", "password", "authorization", "cookie", "x-api-key"}, + } +} + +// AuditOptionsBuilder helps compose AuditMiddlewareOptions with a fluent API +type AuditOptionsBuilder struct { + options *AuditMiddlewareOptions +} + +// NewAuditOptionsBuilder creates a builder initialized with default options +func NewAuditOptionsBuilder() *AuditOptionsBuilder { + return &AuditOptionsBuilder{options: DefaultAuditMiddlewareOptions()} +} + +// WithLogger sets a custom zap logger for audit output +func (b *AuditOptionsBuilder) WithLogger(logger *zap.Logger) *AuditOptionsBuilder { + b.options.Logger = NewAuditLogger(logger) + return b +} + +// WithLevel sets the audit verbosity level +func (b *AuditOptionsBuilder) WithLevel(level AuditLevel) *AuditOptionsBuilder { + b.options.Level = level + return b +} + +// WithHashSensitiveData toggles hashing of sensitive fields before logging +func (b *AuditOptionsBuilder) WithHashSensitiveData(hash bool) *AuditOptionsBuilder { + b.options.HashSensitiveData = hash + return b +} + +// WithRequestBody toggles inclusion of request body in audit events +func (b *AuditOptionsBuilder) WithRequestBody(include bool) *AuditOptionsBuilder { + b.options.IncludeRequestBody = include + return b +} + +// WithResponseBody toggles inclusion of response body in audit events +func (b *AuditOptionsBuilder) WithResponseBody(include bool) *AuditOptionsBuilder { + b.options.IncludeResponseBody = include + return b +} + +// WithRiskAssessor sets a custom risk assessment function +func (b *AuditOptionsBuilder) WithRiskAssessor(assessor func(AuditEvent) (string, []string)) *AuditOptionsBuilder { + b.options.RiskAssessor = assessor + return b +} + +// WithMetadataExtractor sets a function to extract extra metadata from requests +func (b *AuditOptionsBuilder) WithMetadataExtractor(extractor func(*http.Request) map[string]interface{}) *AuditOptionsBuilder { + b.options.MetadataExtractor = extractor + return b +} + +// WithEndpointPatterns sets regex patterns for endpoints to include in auditing +func (b *AuditOptionsBuilder) WithEndpointPatterns(patterns []string) *AuditOptionsBuilder { + b.options.EndpointPatterns = patterns + return b +} + +// WithExcludePatterns sets regex patterns for endpoints to exclude from auditing +func (b *AuditOptionsBuilder) WithExcludePatterns(patterns []string) *AuditOptionsBuilder { + b.options.ExcludePatterns = patterns + return b +} + +// WithSensitiveKeys sets keys that should be redacted in headers and queries +func (b *AuditOptionsBuilder) WithSensitiveKeys(keys []string) *AuditOptionsBuilder { + b.options.SensitiveKeys = keys + return b +} + +// Build finalizes and returns the configured options +func (b *AuditOptionsBuilder) Build() *AuditMiddlewareOptions { + return b.options +} + +// AuditMiddleware returns an HTTP middleware that emits audit events based on the provided options +func AuditMiddleware(options *AuditMiddlewareOptions) func(http.Handler) http.Handler { + // Initialize default options and logger as needed + if options == nil { + options = DefaultAuditMiddlewareOptions() + } + if options.Logger == nil { + options.Logger = NewAuditLogger(nil) + } + // Validate configuration early and fail fast for programmer errors + if err := validateOptions(options); err != nil { + panic(fmt.Sprintf("invalid audit middleware options: %v", err)) + } + + // Wrap the next handler with auditing behavior + return func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Skip auditing if path does not match include/exclude rules + if !shouldAuditPath(r.URL.Path, options.EndpointPatterns, options.ExcludePatterns) { + next.ServeHTTP(w, r) + return + } + + // Detect Server-Sent Events and avoid capturing streaming bodies + acceptHeader := r.Header.Get("Accept") + isSSE := strings.Contains(acceptHeader, "text/event-stream") + + // Initialize event and wrap writer for status and body capture + event, wrappedWriter := initializeAuditEvent(w, r, options) + + // Disable response capture for SSE to prevent interference with streaming + if isSSE { + wrappedWriter.captured = false + } + + // Ensure event is logged even if downstream panics or early returns + defer logAuditEvent(event, wrappedWriter, options) + + // Continue to next handler with wrapped writer + next.ServeHTTP(wrappedWriter, r) + }) + } +} + +// validateOptions ensures the options contain at least one include or exclude pattern +func validateOptions(options *AuditMiddlewareOptions) error { + // Must define what to include or exclude to avoid auditing everything by accident + if len(options.EndpointPatterns) == 0 && len(options.ExcludePatterns) == 0 { + return fmt.Errorf("at least one endpoint pattern or exclude pattern must be specified") + } + return nil +} + +// auditResponseWriter wraps ResponseWriter to capture status and response body +type auditResponseWriter struct { + http.ResponseWriter + statusCode int + body []byte + captured bool +} + +// WriteHeader intercepts status codes for auditing +func (w *auditResponseWriter) WriteHeader(code int) { + w.statusCode = code + w.ResponseWriter.WriteHeader(code) +} + +// Write intercepts response body bytes when capture is enabled +func (w *auditResponseWriter) Write(b []byte) (int, error) { + // Default status to 200 OK if not set + if w.statusCode == 0 { + w.statusCode = http.StatusOK + } + // Append to buffer only when capture is enabled + if w.captured || w.body != nil { + w.body = append(w.body, b...) + } + return w.ResponseWriter.Write(b) +} + +// Flush forwards flush calls for streaming responses +func (w *auditResponseWriter) Flush() { + if flusher, ok := w.ResponseWriter.(http.Flusher); ok { + flusher.Flush() + } +} + +// Unwrap returns the underlying ResponseWriter +func (w *auditResponseWriter) Unwrap() http.ResponseWriter { + return w.ResponseWriter +} + +// OAuthInfo carries OAuth 2.1 specific request attributes extracted for auditing +type OAuthInfo struct { + ClientID string + Subject string + Scopes []string + GrantType string + ResponseType string + RedirectURI string + Resource string + Token string + Code string +} + +// extractOAuthInfo pulls OAuth related fields from URL query, form body, headers, and context +func extractOAuthInfo(r *http.Request) OAuthInfo { + info := OAuthInfo{} + + // Extract from query parameters + if r.URL != nil { + query := r.URL.Query() + info.ClientID = query.Get("client_id") + info.ResponseType = query.Get("response_type") + info.RedirectURI = query.Get("redirect_uri") + info.Resource = query.Get("resource") + if scope := query.Get("scope"); scope != "" { + info.Scopes = strings.Split(scope, " ") + } + } + + // Extract from form body when present + if err := r.ParseForm(); err == nil { + if info.GrantType == "" { + info.GrantType = r.FormValue("grant_type") + } + info.Code = r.FormValue("code") + if scope := r.FormValue("scope"); scope != "" && len(info.Scopes) == 0 { + info.Scopes = strings.Split(scope, " ") + } + } + + // Extract bearer token from Authorization header + if auth := r.Header.Get("Authorization"); strings.HasPrefix(auth, "Bearer ") { + info.Token = strings.TrimPrefix(auth, "Bearer ") + } + + if authInfo, ok := server.GetAuthInfo(r.Context()); ok { + // Extract subject from the Extra claims + if authInfo.Extra != nil { + if sub, ok := authInfo.Extra["sub"].(string); ok { + info.Subject = sub + } + } + + // Use scopes from authInfo if not already populated + if len(info.Scopes) == 0 && len(authInfo.Scopes) > 0 { + info.Scopes = authInfo.Scopes + } + + // Extract client_id from Extra claims first + if cid, ok := authInfo.Extra["client_id"].(string); ok && info.ClientID == "" { + info.ClientID = cid + } + + // Fallback to AuthInfo.ClientID field if Extra doesn't contain client_id + if info.ClientID == "" && authInfo.ClientID != "" { + info.ClientID = authInfo.ClientID + } + } + return info +} + +// shouldAuditPath checks include and exclude regex patterns to decide auditing +func shouldAuditPath(path string, includePatterns, excludePatterns []string) bool { + // Exclude takes precedence when matched + for _, pattern := range excludePatterns { + if matched, _ := regexp.MatchString(pattern, path); matched { + return false + } + } + // If no include patterns set then audit all non excluded paths + if len(includePatterns) == 0 { + return true + } + // Audit when any include pattern matches + for _, pattern := range includePatterns { + if matched, _ := regexp.MatchString(pattern, path); matched { + return true + } + } + return false +} + +// determineEventType maps path and method to a coarse event category +func determineEventType(path, method string) string { + switch { + case strings.Contains(path, "/authorize"): + return "oauth_authorization" + case strings.Contains(path, "/token"): + return "oauth_token" + case strings.Contains(path, "/revoke"): + return "oauth_revocation" + case strings.Contains(path, "/register"): + return "oauth_registration" + case strings.Contains(path, "/metadata"): + return "oauth_metadata" + default: + return "oauth_request" + } +} + +// generateEventID builds a unique event identifier based on time and random suffix +func generateEventID() string { + return fmt.Sprintf("audit_%d_%s", time.Now().UnixNano(), randomString(8)) +} + +// randomString generates a pseudo random lowercase alphanumeric string +func randomString(length int) string { + const charset = "abcdefghijklmnopqrstuvwxyz0123456789" + b := make([]byte, length) + + // Try cryptographic randomness first and fall back to time based selection + if _, err := rand.Read(b); err != nil { + for i := range b { + b[i] = charset[time.Now().UnixNano()%int64(len(charset))] + } + } else { + for i := range b { + b[i] = charset[int(b[i])%len(charset)] + } + } + return string(b) +} + +// sanitizeMap redacts configured sensitive keys and normalizes values to a single string +func sanitizeMap[T string | []string](data map[string]T, sensitiveKeys []string) map[string]string { + sanitized := make(map[string]string) + + // Iterate keys and redact any that match configured sensitive keys + for key, value := range data { + isSensitive := false + for _, sensitiveKey := range sensitiveKeys { + if strings.EqualFold(key, sensitiveKey) { + isSensitive = true + break + } + } + if isSensitive { + sanitized[key] = "[REDACTED]" + } else { + // Normalize to first value for slice and direct string otherwise + switch v := any(value).(type) { + case string: + sanitized[key] = v + case []string: + if len(v) > 0 { + sanitized[key] = v[0] + } + } + } + } + return sanitized +} + +// sanitizeQueryParams applies redaction and normalization to URL query parameters +func sanitizeQueryParams(query map[string][]string, sensitiveKeys []string) map[string]string { + return sanitizeMap(query, sensitiveKeys) +} + +// sanitizeHeaders applies redaction and normalization to HTTP headers +func sanitizeHeaders(headers map[string][]string, sensitiveKeys []string) map[string]string { + return sanitizeMap(headers, sensitiveKeys) +} + +// hashSensitiveData returns a hex encoded SHA256 hash for a sensitive string +func hashSensitiveData(data string) string { + if data == "" { + return "" + } + hash := sha256.Sum256([]byte(data)) + return hex.EncodeToString(hash[:]) +} + +// defaultRiskAssessment computes a simple risk score and contributing factors +func defaultRiskAssessment(event AuditEvent) (string, []string) { + var riskFactors []string + riskLevel := "low" + + // Client side error increases risk + if event.StatusCode >= 400 { + riskFactors = append(riskFactors, "client_error") + } + + // Server side error increases risk more + if event.StatusCode >= 500 { + riskFactors = append(riskFactors, "server_error") + riskLevel = "medium" + } + + // Slow responses may indicate issues + if event.ResponseTime > 5*time.Second { + riskFactors = append(riskFactors, "slow_response") + riskLevel = "medium" + } + + // Missing client identifier is suspicious + if event.ClientID == "" { + riskFactors = append(riskFactors, "missing_client_id") + riskLevel = "high" + } + + // Specific endpoint categories slightly elevate risk + if strings.Contains(event.Path, "/revoke") { + riskFactors = append(riskFactors, "token_revocation") + riskLevel = "medium" + } + if strings.Contains(event.Path, "/register") { + riskFactors = append(riskFactors, "client_registration") + riskLevel = "medium" + } + return riskLevel, riskFactors +} + +// determineErrorCode maps HTTP status codes to OAuth style error codes +func determineErrorCode(statusCode int) string { + switch { + case statusCode == 400: + return "invalid_request" + case statusCode == 401: + return "invalid_token" + case statusCode == 403: + return "insufficient_scope" + case statusCode == 404: + return "not_found" + case statusCode == 429: + return "too_many_requests" + case statusCode >= 500: + return "server_error" + default: + return "unknown_error" + } +} + +// determineErrorMessage extracts error text from a JSON error response or falls back to status text +func determineErrorMessage(statusCode int, body []byte) string { + if len(body) == 0 { + return "" + } + var errorResponse struct { + Error string `json:"error"` + ErrorDescription string `json:"error_description"` + } + // Try decode standard OAuth error response + if err := json.Unmarshal(body, &errorResponse); err == nil { + if errorResponse.ErrorDescription != "" { + return errorResponse.ErrorDescription + } + if errorResponse.Error != "" { + return errorResponse.Error + } + } + // Fallback to generic status text when payload is not structured + return http.StatusText(statusCode) +} + +// extractSubject reads the subject claim from AuthInfo.Extra +func extractSubject(authInfo server.AuthInfo) string { + if authInfo.Extra != nil { + if sub, ok := authInfo.Extra["sub"].(string); ok { + return sub + } + } + return "" +} + +// GetAuthInfo extracts AuthInfo from the request context +func GetAuthInfo(ctx context.Context) (server.AuthInfo, bool) { + // The context key is expected to be provided by upstream auth middleware + if authInfo, ok := ctx.Value(AuthInfoKey).(server.AuthInfo); ok { + return authInfo, true + } + return server.AuthInfo{}, false +} + +// initializeAuditEvent constructs an AuditEvent and wraps the ResponseWriter for capture +func initializeAuditEvent(w http.ResponseWriter, r *http.Request, options *AuditMiddlewareOptions) (AuditEvent, *auditResponseWriter) { + start := time.Now() + + // Optionally read and restore the request body for logging + var reqBody []byte + if (options.Level >= AuditLevelFull || options.IncludeRequestBody) && r.Body != nil { + reqBody, _ = io.ReadAll(r.Body) + _ = r.Body.Close() + r.Body = io.NopCloser(bytes.NewBuffer(reqBody)) + } + + // Configure wrapped writer to capture response body when enabled + wrappedWriter := &auditResponseWriter{ + ResponseWriter: w, + captured: (options.Level >= AuditLevelFull || options.IncludeResponseBody), + } + + // Gather OAuth specific attributes for context + oauthInfo := extractOAuthInfo(r) + + // Seed the event with request metadata and extracted OAuth fields + event := AuditEvent{ + EventID: generateEventID(), + Timestamp: start, + EventType: determineEventType(r.URL.Path, r.Method), + AuditLevel: options.Level, + Method: r.Method, + Path: r.URL.Path, + RemoteAddr: r.RemoteAddr, + UserAgent: r.UserAgent(), + RequestID: r.Header.Get("X-Request-ID"), + ClientID: oauthInfo.ClientID, + Subject: oauthInfo.Subject, + Scopes: oauthInfo.Scopes, + GrantType: oauthInfo.GrantType, + ResponseType: oauthInfo.ResponseType, + RedirectURI: oauthInfo.RedirectURI, + Resource: oauthInfo.Resource, + Metadata: make(map[string]interface{}), + } + + // At detailed level and above include sanitized query params and headers + if options.Level >= AuditLevelDetailed { + event.QueryParams = sanitizeQueryParams(r.URL.Query(), options.SensitiveKeys) + event.Headers = sanitizeHeaders(r.Header, options.SensitiveKeys) + } + + // Hash tokens and IP address when configured to avoid leaking PII + if options.HashSensitiveData { + event.TokenHash = hashSensitiveData(oauthInfo.Token) + event.CodeHash = hashSensitiveData(oauthInfo.Code) + event.IPHash = hashSensitiveData(r.RemoteAddr) + } + + // Include request body when configured + if (options.Level >= AuditLevelFull || options.IncludeRequestBody) && len(reqBody) > 0 { + event.RequestBody = string(reqBody) + } + + // Extract custom metadata when a provider is supplied + if options.MetadataExtractor != nil { + event.Metadata = options.MetadataExtractor(r) + } + + // Assess risk using custom function or defaults + if options.RiskAssessor != nil { + event.RiskLevel, event.RiskFactors = options.RiskAssessor(event) + } else { + event.RiskLevel, event.RiskFactors = defaultRiskAssessment(event) + } + + return event, wrappedWriter +} + +// logAuditEvent finalizes timing and status then emits the audit event via the configured logger +func logAuditEvent(event AuditEvent, w *auditResponseWriter, options *AuditMiddlewareOptions) { + // Compute latency and attach final status code + event.ResponseTime = time.Since(event.Timestamp) + event.StatusCode = w.statusCode + + // Optionally include captured response body + if w.captured && len(w.body) > 0 && (options.Level >= AuditLevelFull || options.IncludeResponseBody) { + event.ResponseBody = string(w.body) + } + + // Derive error details for non successful responses + if event.StatusCode >= 400 { + event.ErrorCode = determineErrorCode(event.StatusCode) + event.ErrorMessage = determineErrorMessage(event.StatusCode, w.body) + } + + // Emit the event and log any failure to stdout as a last resort + if err := options.Logger.LogEvent(event); err != nil { + fmt.Printf("[AUDIT ERROR] Failed to log audit event: %v\n", err) + } +} + +// WithOAuthAudit returns an OAuth specific audit middleware using provided options +func WithOAuthAudit(options *AuditMiddlewareOptions) func(http.Handler) http.Handler { + return AuditMiddleware(options) +} + +// WithBasicAudit returns a middleware configured for basic auditing +func WithBasicAudit() func(http.Handler) http.Handler { + return AuditMiddleware(NewAuditOptionsBuilder(). + WithLevel(AuditLevelBasic). + Build()) +} + +// WithDetailedAudit returns a middleware configured for detailed auditing +func WithDetailedAudit() func(http.Handler) http.Handler { + return AuditMiddleware(NewAuditOptionsBuilder(). + WithLevel(AuditLevelDetailed). + Build()) +} + +// WithFullAudit returns a middleware configured for full auditing including bodies +func WithFullAudit() func(http.Handler) http.Handler { + return AuditMiddleware(NewAuditOptionsBuilder(). + WithLevel(AuditLevelFull). + WithRequestBody(true). + WithResponseBody(true). + Build()) +} + +// WithZapLogger returns options pre configured with a custom zap logger +func WithZapLogger(logger *zap.Logger) *AuditMiddlewareOptions { + return NewAuditOptionsBuilder(). + WithLogger(logger). + Build() +} + +// WithCustomZapLogger returns options configured with a custom zap logger and core toggles +func WithCustomZapLogger(logger *zap.Logger, level AuditLevel, hashSensitive bool) *AuditMiddlewareOptions { + return NewAuditOptionsBuilder(). + WithLogger(logger). + WithLevel(level). + WithHashSensitiveData(hashSensitive). + Build() +} diff --git a/internal/auth/server/middleware/audit_test.go b/internal/auth/server/middleware/audit_test.go new file mode 100644 index 0000000..d36ec83 --- /dev/null +++ b/internal/auth/server/middleware/audit_test.go @@ -0,0 +1,979 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package middleware + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "go.uber.org/zap/zaptest" + + "trpc.group/trpc-go/trpc-mcp-go/internal/auth/server" +) + +// contains checks whether a slice of strings contains a specific item +func contains(slice []string, item string) bool { + for _, s := range slice { + if s == item { + return true + } + } + return false +} + +// captureLogger is a mock implementation of AuditLogger +// It captures the last logged AuditEvent for inspection in tests +type captureLogger struct { + last AuditEvent +} + +// LogEvent stores the provided event in captureLogger +func (c *captureLogger) LogEvent(e AuditEvent) error { + c.last = e + return nil +} + +// LogError stores the event along with the error message in captureLogger +func (c *captureLogger) LogError(e AuditEvent, err error) error { + e.ErrorMessage = err.Error() + c.last = e + return nil +} + +func TestAuditLevelConstants(t *testing.T) { + // Test audit level constant value + if AuditLevelNone != 0 { + t.Errorf("Expected AuditLevelNone to be 0, got %d", AuditLevelNone) + } + if AuditLevelBasic != 1 { + t.Errorf("Expected AuditLevelBasic to be 1, got %d", AuditLevelBasic) + } + if AuditLevelDetailed != 2 { + t.Errorf("Expected AuditLevelDetailed to be 2, got %d", AuditLevelDetailed) + } + if AuditLevelFull != 3 { + t.Errorf("Expected AuditLevelFull to be 3, got %d", AuditLevelFull) + } +} + +func TestNewAuditLogger(t *testing.T) { + // Test create default logger + logger := NewAuditLogger(nil) + if logger == nil { + t.Fatal("Expected logger to be created") + } + + // Test to get the underlying zap logger + zapLogger := logger.GetZapLogger() + if zapLogger == nil { + t.Fatal("Expected underlying zap logger to exist") + } + + // Testing using a custom zap logger + testLogger := zaptest.NewLogger(t) + customLogger := NewAuditLogger(testLogger) + if customLogger == nil { + t.Fatal("Expected custom logger to be created") + } + + if customLogger.GetZapLogger() != testLogger { + t.Fatal("Expected custom zap logger to be used") + } +} + +func TestDefaultAuditLoggerLogEvent(t *testing.T) { + // Creating a test logger + testLogger := zaptest.NewLogger(t) + auditLogger := NewAuditLogger(testLogger) + + // Creating a test event + event := AuditEvent{ + EventID: "test_123", + Timestamp: time.Now(), + EventType: "test_event", + AuditLevel: AuditLevelBasic, + Method: "GET", + Path: "/test", + StatusCode: 200, + ResponseTime: 100 * time.Millisecond, + ClientID: "test_client", + Subject: "test_user", + Scopes: []string{"read", "write"}, + RiskLevel: "low", + RiskFactors: []string{"normal"}, + } + + // Test logging + err := auditLogger.LogEvent(event) + if err != nil { + t.Errorf("Expected no error, got %v", err) + } +} + +func TestDefaultAuditLoggerLogError(t *testing.T) { + testLogger := zaptest.NewLogger(t) + auditLogger := NewAuditLogger(testLogger) + + event := AuditEvent{ + EventID: "test_123", + Timestamp: time.Now(), + Method: "GET", + Path: "/test", + } + + testErr := &http.MaxBytesError{} + err := auditLogger.LogError(event, testErr) + if err != nil { + t.Errorf("Expected no error, got %v", err) + } + + // Note: Since the event is passed by value, the ErrorMessage is not modified. + // This test mainly verifies that the LogError method does not return an error. +} + +func TestDefaultAuditMiddlewareOptions(t *testing.T) { + options := DefaultAuditMiddlewareOptions() + + // Verify default values + if options.Logger == nil { + t.Error("Expected logger to be set") + } + if options.Level != AuditLevelDetailed { + t.Errorf("Expected level to be Detailed, got %v", options.Level) + } + if !options.HashSensitiveData { + t.Error("Expected HashSensitiveData to be true") + } + if len(options.EndpointPatterns) == 0 { + t.Error("Expected endpoint patterns to be set") + } + if len(options.SensitiveKeys) == 0 { + t.Error("Expected sensitive keys to be set") + } +} + +func TestAuditOptionsBuilder(t *testing.T) { + builder := NewAuditOptionsBuilder() + + // Testing chain calls + options := builder. + WithLevel(AuditLevelFull). + WithHashSensitiveData(false). + WithRequestBody(true). + WithResponseBody(true). + Build() + + if options.Level != AuditLevelFull { + t.Errorf("Expected level to be Full, got %v", options.Level) + } + if options.HashSensitiveData { + t.Error("Expected HashSensitiveData to be false") + } + if !options.IncludeRequestBody { + t.Error("Expected IncludeRequestBody to be true") + } + if !options.IncludeResponseBody { + t.Error("Expected IncludeResponseBody to be true") + } +} + +func TestAuditOptionsBuilderWithCustomFunctions(t *testing.T) { + builder := NewAuditOptionsBuilder() + + // Custom Risk Assessor + customRiskAssessor := func(event AuditEvent) (string, []string) { + return "high", []string{"custom_risk"} + } + + // Custom metadata extractors + customMetadataExtractor := func(r *http.Request) map[string]interface{} { + return map[string]interface{}{ + "custom_field": "custom_value", + } + } + + options := builder. + WithRiskAssessor(customRiskAssessor). + WithMetadataExtractor(customMetadataExtractor). + Build() + + if options.RiskAssessor == nil { + t.Error("Expected RiskAssessor to be set") + } + if options.MetadataExtractor == nil { + t.Error("Expected MetadataExtractor to be set") + } + + // Testing custom functions + event := AuditEvent{} + riskLevel, riskFactors := options.RiskAssessor(event) + if riskLevel != "high" { + t.Errorf("Expected risk level 'high', got %s", riskLevel) + } + if len(riskFactors) != 1 || riskFactors[0] != "custom_risk" { + t.Errorf("Expected risk factors ['custom_risk'], got %v", riskFactors) + } +} + +func TestValidateOptions(t *testing.T) { + // Testing a valid configuration + validOptions := &AuditMiddlewareOptions{ + EndpointPatterns: []string{"/test"}, + } + if err := validateOptions(validOptions); err != nil { + t.Errorf("Expected no error for valid options, got %v", err) + } + + // Testing for invalid configurations + invalidOptions := &AuditMiddlewareOptions{ + EndpointPatterns: []string{}, + ExcludePatterns: []string{}, + } + if err := validateOptions(invalidOptions); err == nil { + t.Error("Expected error for invalid options") + } +} + +func TestShouldAuditPath(t *testing.T) { + tests := []struct { + path string + includePatterns []string + excludePatterns []string + expectedResult bool + description string + }{ + { + path: "/oauth2/authorize", + includePatterns: []string{"/oauth2/.*"}, + excludePatterns: []string{}, + expectedResult: true, + description: "Path matches include pattern", + }, + { + path: "/health", + includePatterns: []string{"/oauth2/.*"}, + excludePatterns: []string{}, + expectedResult: false, + description: "Path doesn't match include pattern", + }, + { + path: "/oauth2/token", + includePatterns: []string{"/oauth2/.*"}, + excludePatterns: []string{"/oauth2/token"}, + expectedResult: false, + description: "Path matches exclude pattern", + }, + { + path: "/any/path", + includePatterns: []string{}, + excludePatterns: []string{}, + expectedResult: true, + description: "No patterns specified, audit all", + }, + } + + for _, tt := range tests { + t.Run(tt.description, func(t *testing.T) { + result := shouldAuditPath(tt.path, tt.includePatterns, tt.excludePatterns) + if result != tt.expectedResult { + t.Errorf("shouldAuditPath(%q, %v, %v) = %v, want %v", + tt.path, tt.includePatterns, tt.excludePatterns, result, tt.expectedResult) + } + }) + } +} + +func TestDetermineEventType(t *testing.T) { + tests := []struct { + path string + method string + expected string + }{ + {"/oauth2/authorize", "GET", "oauth_authorization"}, + {"/oauth2/token", "POST", "oauth_token"}, + {"/oauth2/revoke", "POST", "oauth_revocation"}, + {"/oauth2/register", "POST", "oauth_registration"}, + {"/oauth2/metadata", "GET", "oauth_metadata"}, + {"/unknown/path", "GET", "oauth_request"}, + } + + for _, tt := range tests { + result := determineEventType(tt.path, tt.method) + if result != tt.expected { + t.Errorf("determineEventType(%q, %q) = %q, want %q", + tt.path, tt.method, result, tt.expected) + } + } +} + +func TestGenerateEventID(t *testing.T) { + id1 := generateEventID() + id2 := generateEventID() + + if id1 == id2 { + t.Error("Expected different event IDs") + } + + if !strings.HasPrefix(id1, "audit_") { + t.Errorf("Expected event ID to start with 'audit_', got %s", id1) + } +} + +func TestRandomString(t *testing.T) { + str1 := randomString(10) + str2 := randomString(10) + + if len(str1) != 10 { + t.Errorf("Expected string length 10, got %d", len(str1)) + } + + if str1 == str2 { + t.Error("Expected different random strings") + } +} + +func TestSanitizeMap(t *testing.T) { + // Test query parameter sanitization + queryParams := map[string][]string{ + "client_id": {"test_client"}, + "client_secret": {"secret_value"}, + "scope": {"read write"}, + } + + sensitiveKeys := []string{"client_secret", "password"} + + sanitized := sanitizeQueryParams(queryParams, sensitiveKeys) + + if sanitized["client_id"] != "test_client" { + t.Errorf("Expected client_id to be preserved, got %s", sanitized["client_id"]) + } + + if sanitized["client_secret"] != "[REDACTED]" { + t.Errorf("Expected client_secret to be redacted, got %s", sanitized["client_secret"]) + } + + if sanitized["scope"] != "read write" { + t.Errorf("Expected scope to be preserved, got %s", sanitized["scope"]) + } +} + +func TestSanitizeHeaders(t *testing.T) { + headers := map[string][]string{ + "content-type": {"application/json"}, + "authorization": {"Bearer token123"}, + "user-agent": {"test-agent"}, + } + + sensitiveKeys := []string{"authorization", "cookie"} + + sanitized := sanitizeHeaders(headers, sensitiveKeys) + + if sanitized["content-type"] != "application/json" { + t.Errorf("Expected content-type to be preserved, got %s", sanitized["content-type"]) + } + + if sanitized["authorization"] != "[REDACTED]" { + t.Errorf("Expected authorization to be redacted, got %s", sanitized["authorization"]) + } + + if sanitized["user-agent"] != "test-agent" { + t.Errorf("Expected user-agent to be preserved, got %s", sanitized["user-agent"]) + } +} + +func TestHashSensitiveData(t *testing.T) { + data := "sensitive_data" + hash1 := hashSensitiveData(data) + hash2 := hashSensitiveData(data) + + if hash1 == "" { + t.Error("Expected non-empty hash") + } + + if hash1 != hash2 { + t.Error("Expected same hash for same data") + } + + if hashSensitiveData("") != "" { + t.Error("Expected empty string for empty data") + } +} + +func TestDefaultRiskAssessment(t *testing.T) { + tests := []struct { + name string + event AuditEvent + expectedRisk string + checkFactors func([]string) bool + }{ + { + name: "Normal request", + event: AuditEvent{ + StatusCode: 200, + ResponseTime: 100 * time.Millisecond, + ClientID: "test_client", + }, + expectedRisk: "low", + checkFactors: func(factors []string) bool { + return len(factors) == 0 + }, + }, + { + name: "Client error", + event: AuditEvent{ + StatusCode: 400, + ResponseTime: 100 * time.Millisecond, + ClientID: "test_client", + }, + expectedRisk: "low", + checkFactors: func(factors []string) bool { + return contains(factors, "client_error") + }, + }, + { + name: "Server error", + event: AuditEvent{ + StatusCode: 500, + ResponseTime: 100 * time.Millisecond, + ClientID: "test_client", + }, + expectedRisk: "medium", + checkFactors: func(factors []string) bool { + return contains(factors, "server_error") + }, + }, + { + name: "Slow response", + event: AuditEvent{ + StatusCode: 200, + ResponseTime: 6 * time.Second, + ClientID: "test_client", + }, + expectedRisk: "medium", + checkFactors: func(factors []string) bool { + return contains(factors, "slow_response") + }, + }, + { + name: "Missing client ID", + event: AuditEvent{ + StatusCode: 200, + ResponseTime: 100 * time.Millisecond, + ClientID: "", + }, + expectedRisk: "high", + checkFactors: func(factors []string) bool { + return contains(factors, "missing_client_id") + }, + }, + { + name: "Token revocation", + event: AuditEvent{ + StatusCode: 200, + ResponseTime: 100 * time.Millisecond, + ClientID: "test_client", + Path: "/oauth2/revoke", + }, + expectedRisk: "medium", + checkFactors: func(factors []string) bool { + return contains(factors, "token_revocation") + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + riskLevel, riskFactors := defaultRiskAssessment(tt.event) + if riskLevel != tt.expectedRisk { + t.Errorf("Expected risk level %s, got %s", tt.expectedRisk, riskLevel) + } + if !tt.checkFactors(riskFactors) { + t.Errorf("Risk factors check failed for factors: %v", riskFactors) + } + }) + } +} + +func TestDetermineErrorCode(t *testing.T) { + tests := []struct { + statusCode int + expected string + }{ + {400, "invalid_request"}, + {401, "invalid_token"}, + {403, "insufficient_scope"}, + {404, "not_found"}, + {429, "too_many_requests"}, + {500, "server_error"}, + {999, "server_error"}, // 999 is still considered a server error + } + + for _, tt := range tests { + result := determineErrorCode(tt.statusCode) + if result != tt.expected { + t.Errorf("determineErrorCode(%d) = %s, want %s", tt.statusCode, result, tt.expected) + } + } +} + +func TestDetermineErrorMessage(t *testing.T) { + // Testing JSON error responses + jsonError := `{"error": "invalid_grant", "error_description": "Invalid authorization code"}` + message := determineErrorMessage(400, []byte(jsonError)) + if message != "Invalid authorization code" { + t.Errorf("Expected 'Invalid authorization code', got %s", message) + } + + // Testing responses with only an error field + jsonErrorOnly := `{"error": "invalid_request"}` + message = determineErrorMessage(400, []byte(jsonErrorOnly)) + if message != "invalid_request" { + t.Errorf("Expected 'invalid_request', got %s", message) + } + + // Testing for an empty response body + message = determineErrorMessage(404, []byte{}) + if message != "" { + t.Errorf("Expected empty message for empty body, got %s", message) + } + + // Testing for invalid JSON + invalidJSON := `{invalid json}` + message = determineErrorMessage(500, []byte(invalidJSON)) + if message != "Internal Server Error" { + t.Errorf("Expected 'Internal Server Error', got %s", message) + } +} + +func TestExtractOAuthInfo(t *testing.T) { + // Creating a test request + req := httptest.NewRequest("POST", "/oauth2/token?client_id=test_client&scope=read+write", strings.NewReader("grant_type=authorization_code&code=test_code")) + req.Header.Set("Authorization", "Bearer test_token") + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + + // Set the authentication information in the context + ctx := server.WithAuthInfo(req.Context(), &server.AuthInfo{ + Scopes: []string{"read", "write"}, + Extra: map[string]interface{}{ + "sub": "test_user", + "client_id": "ctx_client_id", + }, + }) + req = req.WithContext(ctx) + + info := extractOAuthInfo(req) + + // Verify the extracted information + if info.ClientID != "test_client" { + t.Errorf("Expected client_id 'test_client', got %s", info.ClientID) + } + + // Note: Since ParseForm may not work in some test environments, we mainly test URL parameters and headers + if info.Token != "test_token" { + t.Errorf("Expected token 'test_token', got %s", info.Token) + } + + if len(info.Scopes) != 2 { + t.Errorf("Expected 2 scopes, got %d", len(info.Scopes)) + } + + if info.Subject != "test_user" { + t.Errorf("Expected subject 'test_user', got %s", info.Subject) + } +} + +func TestExtractSubject(t *testing.T) { + authInfo := server.AuthInfo{ + Extra: map[string]interface{}{ + "sub": "test_user", + "other": "value", + }, + } + + subject := extractSubject(authInfo) + if subject != "test_user" { + t.Errorf("Expected subject 'test_user', got %s", subject) + } + + // Testing without a subject + authInfoNoSub := server.AuthInfo{ + Extra: map[string]interface{}{ + "other": "value", + }, + } + + subject = extractSubject(authInfoNoSub) + if subject != "" { + t.Errorf("Expected empty subject, got %s", subject) + } +} + +func TestAuditMiddlewareBasic(t *testing.T) { + // Create basic audit middleware + middleware := WithBasicAudit() + + // Create a Test Handler + testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + w.Write([]byte("test response")) + }) + + // Package handler + wrappedHandler := middleware(testHandler) + + // Create a test request + req := httptest.NewRequest("GET", "/oauth2/authorize?client_id=test_client", nil) + w := httptest.NewRecorder() + + // Execute Request + wrappedHandler.ServeHTTP(w, req) + + // Validate response + if w.Code != http.StatusOK { + t.Errorf("Expected status 200, got %d", w.Code) + } +} + +func TestAuditMiddlewareDetailed(t *testing.T) { + // Create detailed audit middleware + middleware := WithDetailedAudit() + + // Create a Test Handler + testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + w.Write([]byte("test response")) + }) + + // Package handler + wrappedHandler := middleware(testHandler) + + // Create a test request + req := httptest.NewRequest("POST", "/oauth2/token", strings.NewReader("grant_type=client_credentials")) + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + req.Header.Set("User-Agent", "test-agent") + + w := httptest.NewRecorder() + + // Execute Request + wrappedHandler.ServeHTTP(w, req) + + // Validate response + if w.Code != http.StatusOK { + t.Errorf("Expected status 200, got %d", w.Code) + } +} + +func TestAuditMiddlewareFull(t *testing.T) { + // Create full audit middleware + middleware := WithFullAudit() + + // Create a Test Handler + testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + w.Write([]byte("test response")) + }) + + // Package handler + wrappedHandler := middleware(testHandler) + + // Create a test request + req := httptest.NewRequest("DELETE", "/oauth2/revoke", strings.NewReader("token=test_token")) + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + + w := httptest.NewRecorder() + + // Execute request + wrappedHandler.ServeHTTP(w, req) + + // Validate response + if w.Code != http.StatusOK { + t.Errorf("Expected status 200, got %d", w.Code) + } +} + +func TestAuditMiddlewareWithCustomZapLogger(t *testing.T) { + // Create a custom zap logger + testLogger := zaptest.NewLogger(t) + + // Create audit middleware with a custom logger + options := WithCustomZapLogger(testLogger, AuditLevelDetailed, true) + middleware := AuditMiddleware(options) + + // Create a Test Handler + testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + w.Write([]byte("test response")) + }) + + // Package handler + wrappedHandler := middleware(testHandler) + + // Create a test request + req := httptest.NewRequest("GET", "/oauth2/metadata", nil) + w := httptest.NewRecorder() + + // Execute request + wrappedHandler.ServeHTTP(w, req) + + // Validate response + if w.Code != http.StatusOK { + t.Errorf("Expected status 200, got %d", w.Code) + } +} + +func TestAuditMiddlewareExcludePatterns(t *testing.T) { + // Create custom configured audit middleware + options := NewAuditOptionsBuilder(). + WithEndpointPatterns([]string{"/oauth2/.*"}). + WithExcludePatterns([]string{"/oauth2/health"}). + Build() + + middleware := AuditMiddleware(options) + + // Create a test handler + testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + w.Write([]byte("test response")) + }) + + // Package handler + wrappedHandler := middleware(testHandler) + + // Test paths that should be audited + req1 := httptest.NewRequest("GET", "/oauth2/authorize", nil) + w1 := httptest.NewRecorder() + wrappedHandler.ServeHTTP(w1, req1) + + if w1.Code != http.StatusOK { + t.Errorf("Expected status 200 for audited path, got %d", w1.Code) + } + + // Testing paths that should be excluded + req2 := httptest.NewRequest("GET", "/oauth2/health", nil) + w2 := httptest.NewRecorder() + wrappedHandler.ServeHTTP(w2, req2) + + if w2.Code != http.StatusOK { + t.Errorf("Expected status 200 for excluded path, got %d", w2.Code) + } +} + +func TestAuditMiddlewareErrorHandling(t *testing.T) { + // Create a test logger + testLogger := zaptest.NewLogger(t) + + // Create test audit middleware + options := WithZapLogger(testLogger) + middleware := AuditMiddleware(options) + + // Creating a test handler that returns an error + testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusBadRequest) + errorResponse := map[string]string{ + "error": "invalid_request", + "error_description": "Missing required parameter", + } + json.NewEncoder(w).Encode(errorResponse) + }) + + // Package handler + wrappedHandler := middleware(testHandler) + + // Create a test request + req := httptest.NewRequest("POST", "/oauth2/token", strings.NewReader("")) + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + + w := httptest.NewRecorder() + + // Execute request + wrappedHandler.ServeHTTP(w, req) + + // Validate request + if w.Code != http.StatusBadRequest { + t.Errorf("Expected status 400, got %d", w.Code) + } +} + +func TestAuditMiddlewarePerformance(t *testing.T) { + // Creating audit middleware for performance testing + options := NewAuditOptionsBuilder(). + WithLevel(AuditLevelBasic). + WithHashSensitiveData(false). + Build() + + middleware := AuditMiddleware(options) + + // Create test handler + testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Simulation processing time + time.Sleep(50 * time.Millisecond) + w.WriteHeader(http.StatusOK) + w.Write([]byte("test response")) + }) + + // Package handler + wrappedHandler := middleware(testHandler) + + // Creating a test request + req := httptest.NewRequest("GET", "/oauth2/authorize", nil) + w := httptest.NewRecorder() + + // Execute request + start := time.Now() + wrappedHandler.ServeHTTP(w, req) + duration := time.Since(start) + + // Validate response + if w.Code != http.StatusOK { + t.Errorf("Expected status 200, got %d", w.Code) + } + + // Verify performance + if duration > 200*time.Millisecond { + t.Errorf("Expected reasonable performance, took %v", duration) + } +} + +func TestAuditResponseWriter(t *testing.T) { + // Creating a Test Response Recorder + recorder := httptest.NewRecorder() + + // Creating an Audit Response Writer + auditWriter := &auditResponseWriter{ + ResponseWriter: recorder, + body: make([]byte, 0), + } + + // Test write header + auditWriter.WriteHeader(http.StatusCreated) + if auditWriter.statusCode != http.StatusCreated { + t.Errorf("Expected status code %d, got %d", http.StatusCreated, auditWriter.statusCode) + } + + // Test writing data + testData := []byte("test response") + written, err := auditWriter.Write(testData) + if err != nil { + t.Errorf("Expected no error, got %v", err) + } + if written != len(testData) { + t.Errorf("Expected written bytes %d, got %d", len(testData), written) + } + + // Verification status code is set + if auditWriter.statusCode == 0 { + auditWriter.statusCode = http.StatusOK + } + + // Verify that the response body is captured + if len(auditWriter.body) == 0 { + t.Error("Expected response body to be captured") + } +} + +func BenchmarkAuditMiddleware(b *testing.B) { + // Creating basic audit middleware + middleware := WithBasicAudit() + + // Creating a Test Handler + testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + w.Write([]byte("test response")) + }) + + // Packaging Processor + wrappedHandler := middleware(testHandler) + + // Creating a test request + req := httptest.NewRequest("GET", "/oauth2/authorize", nil) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + w := httptest.NewRecorder() + wrappedHandler.ServeHTTP(w, req) + } +} + +func BenchmarkHashSensitiveData(b *testing.B) { + testData := "sensitive_test_data_that_needs_hashing" + + b.ResetTimer() + for i := 0; i < b.N; i++ { + _ = hashSensitiveData(testData) + } +} + +func BenchmarkSanitizeMap(b *testing.B) { + queryParams := map[string][]string{ + "client_id": {"test_client"}, + "client_secret": {"secret_value"}, + "scope": {"read write"}, + } + sensitiveKeys := []string{"client_secret", "password"} + + b.ResetTimer() + for i := 0; i < b.N; i++ { + _ = sanitizeQueryParams(queryParams, sensitiveKeys) + } +} + +func TestAuditWriter_DefaultStatusOnWrite(t *testing.T) { + rec := httptest.NewRecorder() + aw := &auditResponseWriter{ResponseWriter: rec, body: make([]byte, 0)} + _, _ = aw.Write([]byte("hi")) + if aw.statusCode != http.StatusOK { + t.Fatalf("status should default to 200 on Write, got %d", aw.statusCode) + } +} + +func TestSSE_DisablesCapture(t *testing.T) { + cl := &captureLogger{} + opts := NewAuditOptionsBuilder().Build() + opts.Logger = cl + opts.IncludeResponseBody = true // 即便配置为 true，SSE 也应禁用 + mw := AuditMiddleware(opts) + + h := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // 模拟 SSE 输出 + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + w.(http.Flusher).Flush() + w.Write([]byte("data: ping\n\n")) + }) + + req := httptest.NewRequest("GET", "/oauth2/authorize", nil) + req.Header.Set("Accept", "text/event-stream") + w := httptest.NewRecorder() + mw(h).ServeHTTP(w, req) + + if cl.last.ResponseBody != "" { + t.Fatalf("SSE responses should not be captured") + } +} + +func TestRequestBody_CapturedWhenEnabled(t *testing.T) { + cl := &captureLogger{} + opts := NewAuditOptionsBuilder().Build() + opts.Logger = cl + opts.IncludeRequestBody = true + + mw := AuditMiddleware(opts) + h := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusOK) }) + req := httptest.NewRequest("POST", "/oauth2/token", strings.NewReader("grant_type=client_credentials")) + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + + w := httptest.NewRecorder() + mw(h).ServeHTTP(w, req) + + if cl.last.RequestBody == "" || !strings.Contains(cl.last.RequestBody, "grant_type=client_credentials") { + t.Fatalf("request body should be captured when IncludeRequestBody=true") + } +} diff --git a/internal/auth/server/middleware/bearer_auth.go b/internal/auth/server/middleware/bearer_auth.go new file mode 100644 index 0000000..72be91c --- /dev/null +++ b/internal/auth/server/middleware/bearer_auth.go @@ -0,0 +1,166 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package middleware + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "strings" + "time" + + "trpc.group/trpc-go/trpc-mcp-go/internal/auth/server" + "trpc.group/trpc-go/trpc-mcp-go/internal/errors" +) + +// audienceMatchLocal matches resource against allowed audience values (trim trailing '#') +func audienceMatchLocal(resource string, allowed []string) bool { + resource = strings.TrimSuffix(strings.TrimSpace(resource), "#") + for _, a := range allowed { + if resource == strings.TrimSuffix(strings.TrimSpace(a), "#") { + return true + } + } + return false +} + +// BearerAuthMiddlewareOptions defines configuration for the Bearer auth middleware +type BearerAuthMiddlewareOptions struct { + // Verifier is used to validate the access token + Verifier server.TokenVerifierInterface + + // RequiredScopes lists scopes that must all be present in the token + RequiredScopes []string + + // ResourceMetadataURL is optionally included in the WWW-Authenticate header + ResourceMetadataURL *string + + // Issuer restricts accepted tokens to this issuer (optional) + Issuer string + + // Audience restricts accepted tokens to this audience/resource (optional) + Audience []string +} + +// RequireBearerAuth returns an HTTP middleware that validates Bearer tokens on incoming requests +func RequireBearerAuth(options BearerAuthMiddlewareOptions) func(handler http.Handler) http.Handler { + return func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { + // setErrorResponse writes a JSON OAuth error and appropriate status and headers + setErrorResponse := func(w http.ResponseWriter, err errors.OAuthError, statusCode int) { + // Set WWW-Authenticate only for 401 or 403 to align with TS implementation + if statusCode == http.StatusUnauthorized || statusCode == http.StatusForbidden { + wwwAuthValue := fmt.Sprintf(`Bearer error="%s", error_description="%s"`, err.ErrorCode, err.Message) + if options.ResourceMetadataURL != nil { + wwwAuthValue += fmt.Sprintf(`, resource_metadata="%s"`, *options.ResourceMetadataURL) + } + // Append scope for insufficient_scope + if err.ErrorCode == errors.ErrInsufficientScope.Error() && len(options.RequiredScopes) > 0 { + wwwAuthValue += fmt.Sprintf(`, scope="%s"`, strings.Join(options.RequiredScopes, " ")) + } + w.Header().Set("WWW-Authenticate", wwwAuthValue) + } + // Write JSON body with error details + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(statusCode) + _ = json.NewEncoder(w).Encode(err.ToResponseStruct()) + } + + // Read Authorization header and ensure presence + authHeader := req.Header.Get("Authorization") + if authHeader == "" { + setErrorResponse(w, errors.NewOAuthError(errors.ErrInvalidToken, "Missing Authorization header", ""), http.StatusUnauthorized) + return + } + + // Expect "Bearer " format and extract the token + parts := strings.Split(authHeader, " ") + if len(parts) != 2 || strings.ToLower(parts[0]) != "bearer" || parts[1] == "" { + setErrorResponse(w, errors.NewOAuthError(errors.ErrInvalidToken, "Invalid Authorization header format, expected 'Bearer TOKEN'", ""), http.StatusUnauthorized) + return + } + token := parts[1] + + // Verify token using provided verifier + authInfo, err := options.Verifier.VerifyAccessToken(req.Context(), token) + if err != nil { + // Map verifier error to HTTP status via OAuth error code + if oauthErr, ok := err.(errors.OAuthError); ok { + switch oauthErr.ErrorCode { + case errors.ErrInvalidToken.Error(): + setErrorResponse(w, oauthErr, http.StatusUnauthorized) + case errors.ErrInsufficientScope.Error(): + setErrorResponse(w, oauthErr, http.StatusForbidden) + case errors.ErrServerError.Error(): + setErrorResponse(w, oauthErr, http.StatusInternalServerError) + default: + setErrorResponse(w, oauthErr, http.StatusBadRequest) + } + } else { + // Default unknown errors to invalid_token (401) to avoid leaking internals + invalid := errors.NewOAuthError(errors.ErrInvalidToken, "Invalid access token", "") + setErrorResponse(w, invalid, http.StatusUnauthorized) + } + return + } + + // Optional issuer guarantee + if options.Issuer != "" { + if authInfo.Extra != nil { + if iss, _ := authInfo.Extra["iss"].(string); iss != "" && iss != options.Issuer { + setErrorResponse(w, errors.NewOAuthError(errors.ErrInvalidToken, "Invalid token issuer", ""), http.StatusUnauthorized) + return + } + } + } + + // Optional audience/resource check (RFC 8707 simplified) + if len(options.Audience) > 0 && authInfo.Resource != nil { + if !audienceMatchLocal(authInfo.Resource.String(), options.Audience) { + setErrorResponse(w, errors.NewOAuthError(errors.ErrInvalidToken, "Invalid token audience", ""), http.StatusUnauthorized) + return + } + } + + // Enforce required scopes if configured + if len(options.RequiredScopes) > 0 { + for _, scope := range options.RequiredScopes { + found := false + for _, tokenScope := range authInfo.Scopes { + if tokenScope == scope { + found = true + break + } + } + if !found { + setErrorResponse(w, errors.NewOAuthError(errors.ErrInsufficientScope, "Insufficient scope", ""), http.StatusForbidden) + return + } + } + } + + // Ensure token has an expiration time and is not expired + if authInfo.ExpiresAt == nil || *authInfo.ExpiresAt == 0 { + setErrorResponse(w, errors.NewOAuthError(errors.ErrInvalidToken, "Token has no expiration time", ""), http.StatusUnauthorized) + return + } + if *authInfo.ExpiresAt <= time.Now().Unix() { + setErrorResponse(w, errors.NewOAuthError(errors.ErrInvalidToken, "Token has expired", ""), http.StatusUnauthorized) + return + } + + // Attach validated auth info to the request context under AuthInfoKey (avoid token propagation) + authInfo.Token = "" + ctx := context.WithValue(req.Context(), AuthInfoKey, authInfo) + req = req.WithContext(ctx) + + // Delegate to next handler + next.ServeHTTP(w, req) + }) + } +} diff --git a/internal/auth/server/middleware/bearer_auth_test.go b/internal/auth/server/middleware/bearer_auth_test.go new file mode 100644 index 0000000..6117dfc --- /dev/null +++ b/internal/auth/server/middleware/bearer_auth_test.go @@ -0,0 +1,766 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package middleware + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "testing" + "time" + + srv "trpc.group/trpc-go/trpc-mcp-go/internal/auth/server" + oauth "trpc.group/trpc-go/trpc-mcp-go/internal/errors" +) + +// mockVerifier is a test double for TokenVerifierInterface that records the +// last token and delegates verification to a provided function +type mockVerifier struct { + verify func(ctx context.Context, token string) (srv.AuthInfo, error) + last string +} + +// VerifyAccessToken records the token and forwards verification to the mock function +func (m *mockVerifier) VerifyAccessToken(ctx context.Context, token string) (srv.AuthInfo, error) { + m.last = token + return m.verify(ctx, token) +} + +// runWithMiddleware builds a request, executes the BearerAuth middleware with the +// provided options and Authorization header, and returns the recorder and whether +// the next handler was called +func runWithMiddleware(t *testing.T, options BearerAuthMiddlewareOptions, authHeader string) (rec *httptest.ResponseRecorder, nextCalled bool) { + t.Helper() + nextCalled = false + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + nextCalled = true + // return 200 to assert the middleware allowed the reques + w.WriteHeader(http.StatusOK) + }) + + handler := RequireBearerAuth(options)(next) + + req := httptest.NewRequest(http.MethodGet, "/", nil) + if authHeader != "" { + req.Header.Set("Authorization", authHeader) + } + rec = httptest.NewRecorder() + handler.ServeHTTP(rec, req) + return rec, nextCalled +} + +// decodeOAuthResp parses the OAuth error response body into OAuthErrorResponse +func decodeOAuthResp(t *testing.T, rec *httptest.ResponseRecorder) *oauth.OAuthErrorResponse { + t.Helper() + var body oauth.OAuthErrorResponse + _ = json.NewDecoder(rec.Body).Decode(&body) + return &body +} + +func TestRequireBearerAuth_ValidToken(t *testing.T) { + exp := time.Now().Add(1 * time.Hour).Unix() + valid := srv.AuthInfo{Token: "valid-token", ClientID: "client-123", Scopes: []string{"read", "write"}, ExpiresAt: &exp} + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { return valid, nil }} + + rec, nextCalled := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv}, "Bearer valid-token") + + if mv.last != "valid-token" { + t.Fatalf("expected token 'valid-token', got %q", mv.last) + } + if !nextCalled { + t.Fatalf("expected next to be called") + } + if rec.Code != http.StatusOK { + t.Fatalf("expected status 200, got %d", rec.Code) + } +} + +func TestRequireBearerAuth_ExpiredToken(t *testing.T) { + expiredSeconds := []int{100, 0} + + for _, seconds := range expiredSeconds { + t.Run(fmt.Sprintf("expired_%d_seconds_ago", seconds), func(t *testing.T) { + // Set the expiration time to the current time minus the specified number of seconds + exp := time.Now().Add(time.Duration(-seconds) * time.Second).Unix() + expired := srv.AuthInfo{ + Token: "expired-token", + ClientID: "client-123", + Scopes: []string{"read", "write"}, + ExpiresAt: &exp, + } + mv := &mockVerifier{ + verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { + return expired, nil + }, + } + + rec, nextCalled := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv}, "Bearer expired-token") + + if mv.last != "expired-token" { + t.Fatalf("expected token 'expired-token', got %q", mv.last) + } + if nextCalled { + t.Fatalf("expected next not to be called") + } + if rec.Code != http.StatusUnauthorized { + t.Fatalf("expected 401, got %d", rec.Code) + } + hdr := rec.Header().Get("WWW-Authenticate") + if !strings.Contains(hdr, `error="invalid_token"`) || !strings.Contains(hdr, "Token has expired") { + t.Fatalf("unexpected WWW-Authenticate: %q", hdr) + } + body := decodeOAuthResp(t, rec) + if body.Error != "invalid_token" || body.ErrorDescription != "Token has expired" { + t.Fatalf("unexpected body: %+v", body) + } + }) + } +} + +func TestRequireBearerAuth_NoExpiration(t *testing.T) { + // case1: nil + t.Run("ExpiresAt=nil", func(t *testing.T) { + noexp := srv.AuthInfo{Token: "t1", ClientID: "c", Scopes: []string{"read"}, ExpiresAt: nil} + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { return noexp, nil }} + rec, nextCalled := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv}, "Bearer t1") + if nextCalled || rec.Code != http.StatusUnauthorized { + t.Fatalf("expected 401 without next, got %d next=%v", rec.Code, nextCalled) + } + if hdr := rec.Header().Get("WWW-Authenticate"); !strings.Contains(hdr, `error="invalid_token"`) || !strings.Contains(hdr, "Token has no expiration time") { + t.Fatalf("unexpected WWW-Authenticate: %q", hdr) + } + body := decodeOAuthResp(t, rec) + if body.Error != "invalid_token" || body.ErrorDescription != "Token has no expiration time" { + t.Fatalf("unexpected body: %+v", body) + } + }) + + // case2: 0 + t.Run("ExpiresAt=0", func(t *testing.T) { + zero := int64(0) + noexp2 := srv.AuthInfo{Token: "t2", ClientID: "c", Scopes: []string{"read"}, ExpiresAt: &zero} + mv2 := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { return noexp2, nil }} + rec2, nextCalled2 := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv2}, "Bearer t2") + if nextCalled2 || rec2.Code != http.StatusUnauthorized { + t.Fatalf("expected 401 without next, got %d next=%v", rec2.Code, nextCalled2) + } + if hdr := rec2.Header().Get("WWW-Authenticate"); !strings.Contains(hdr, `error="invalid_token"`) || !strings.Contains(hdr, "Token has no expiration time") { + t.Fatalf("unexpected WWW-Authenticate: %q", hdr) + } + body := decodeOAuthResp(t, rec2) + if body.Error != "invalid_token" || body.ErrorDescription != "Token has no expiration time" { + t.Fatalf("unexpected body: %+v", body) + } + }) +} +func TestRequireBearerAuth_NonExpiredAccepted(t *testing.T) { + exp := time.Now().Add(1 * time.Hour).Unix() + ai := srv.AuthInfo{Token: "valid", ClientID: "c", Scopes: []string{"read"}, ExpiresAt: &exp} + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { return ai, nil }} + rec, nextCalled := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv}, "Bearer valid") + if rec.Code != http.StatusOK || !nextCalled { + t.Fatalf("expected 200 and next called, got %d next=%v", rec.Code, nextCalled) + } +} + +func TestRequireBearerAuth_RequiredScopes(t *testing.T) { + exp := time.Now().Add(1 * time.Hour).Unix() + ai := srv.AuthInfo{Token: "valid", ClientID: "c", Scopes: []string{"read"}, ExpiresAt: &exp} + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { return ai, nil }} + rec, nextCalled := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv, RequiredScopes: []string{"read", "write"}}, "Bearer valid") + + if mv.last != "valid" { + t.Fatalf("expected verifier to be called with token 'valid', got %q", mv.last) + } + if nextCalled || rec.Code != http.StatusForbidden { + t.Fatalf("expected 403 and next not called, got %d next=%v", rec.Code, nextCalled) + } + hdr := rec.Header().Get("WWW-Authenticate") + if !strings.Contains(hdr, `error="insufficient_scope"`) { + t.Fatalf("unexpected WWW-Authenticate: %q", hdr) + } + body := decodeOAuthResp(t, rec) + if body.Error != "insufficient_scope" || body.ErrorDescription != "Insufficient scope" { + t.Fatalf("unexpected body: %+v", body) + } +} + +func TestRequireBearerAuth_AcceptsAllRequiredScopes(t *testing.T) { + exp := time.Now().Add(1 * time.Hour).Unix() + ai := srv.AuthInfo{Token: "valid", ClientID: "c", Scopes: []string{"read", "write", "admin"}, ExpiresAt: &exp} + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { return ai, nil }} + rec, nextCalled := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv, RequiredScopes: []string{"read", "write"}}, "Bearer valid") + if rec.Code != http.StatusOK || !nextCalled { + t.Fatalf("expected 200 and next called, got %d next=%v", rec.Code, nextCalled) + } +} + +func TestRequireBearerAuth_MissingAuthorization(t *testing.T) { + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { return srv.AuthInfo{}, nil }} + rec, nextCalled := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv}, "") + if mv.last != "" { + t.Fatalf("verifier should not be called") + } + if nextCalled || rec.Code != http.StatusUnauthorized { + t.Fatalf("expected 401 without next, got %d next=%v", rec.Code, nextCalled) + } + // 检查完整的 WWW-Authenticate 头 + expectedHeader := `Bearer error="invalid_token", error_description="Missing Authorization header"` + hdr := rec.Header().Get("WWW-Authenticate") + if hdr != expectedHeader { + t.Fatalf("expected WWW-Authenticate: %q, got %q", expectedHeader, hdr) + } + // 检查响应体 + body := decodeOAuthResp(t, rec) + if body.Error != "invalid_token" || body.ErrorDescription != "Missing Authorization header" { + t.Fatalf("expected body error=\"invalid_token\", error_description=\"Missing Authorization header\", got %+v", body) + } +} + +func TestRequireBearerAuth_InvalidAuthorizationFormat(t *testing.T) { + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { return srv.AuthInfo{}, nil }} + rec, nextCalled := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv}, "InvalidFormat") + if mv.last != "" { + t.Fatalf("verifier should not be called") + } + if nextCalled || rec.Code != http.StatusUnauthorized { + t.Fatalf("expected 401 without next, got %d next=%v", rec.Code, nextCalled) + } + // 验证 WWW-Authenticate 头的完整字符串 + expectedHeader := `Bearer error="invalid_token", error_description="Invalid Authorization header format, expected 'Bearer TOKEN'"` + if hdr := rec.Header().Get("WWW-Authenticate"); hdr != expectedHeader { + t.Fatalf("expected WWW-Authenticate: %q, got %q", expectedHeader, hdr) + } + // 验证响应体 + body := decodeOAuthResp(t, rec) + if body.Error != "invalid_token" || body.ErrorDescription != "Invalid Authorization header format, expected 'Bearer TOKEN'" { + t.Fatalf("unexpected body: %+v", body) + } +} + +func TestRequireBearerAuth_VerifierErrors(t *testing.T) { + t.Run("invalid_token -> 401", func(t *testing.T) { + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { + return srv.AuthInfo{}, oauth.NewOAuthError(oauth.ErrInvalidToken, "Token expired", "") + }} + rec, nextCalled := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv}, "Bearer invalid-token") + if mv.last != "invalid-token" { + t.Fatalf("expected token 'invalid-token', got %q", mv.last) + } + if nextCalled || rec.Code != http.StatusUnauthorized { + t.Fatalf("expected 401 without next, got %d next=%v", rec.Code, nextCalled) + } + hdr := rec.Header().Get("WWW-Authenticate") + if !strings.Contains(hdr, `error="invalid_token"`) || !strings.Contains(hdr, "Token expired") { + t.Fatalf("unexpected WWW-Authenticate: %q", hdr) + } + }) + + t.Run("insufficient_scope -> 403", func(t *testing.T) { + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { + return srv.AuthInfo{}, oauth.NewOAuthError(oauth.ErrInsufficientScope, "Required scopes: read, write", "") + }} + rec, nextCalled := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv}, "Bearer valid-token") + if mv.last != "valid-token" { + t.Fatalf("expected token 'valid-token', got %q", mv.last) + } + if nextCalled || rec.Code != http.StatusForbidden { + t.Fatalf("expected 403 without next, got %d next=%v", rec.Code, nextCalled) + } + hdr := rec.Header().Get("WWW-Authenticate") + if !strings.Contains(hdr, `error="insufficient_scope"`) { + t.Fatalf("unexpected WWW-Authenticate: %q", hdr) + } + body := decodeOAuthResp(t, rec) + if body.Error != "insufficient_scope" || body.ErrorDescription != "Required scopes: read, write" { + t.Fatalf("unexpected body: %+v", body) + } + }) + + t.Run("server_error -> 500", func(t *testing.T) { + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { + return srv.AuthInfo{}, oauth.NewOAuthError(oauth.ErrServerError, "Internal server issue", "") + }} + rec, nextCalled := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv}, "Bearer valid-token") + if mv.last != "valid-token" { + t.Fatalf("expected token 'valid-token', got %q", mv.last) + } + if nextCalled || rec.Code != http.StatusInternalServerError { + t.Fatalf("expected 500 without next, got %d next=%v", rec.Code, nextCalled) + } + if hdr := rec.Header().Get("WWW-Authenticate"); hdr != "" { + t.Fatalf("expected no WWW-Authenticate header, got %q", hdr) + } + body := decodeOAuthResp(t, rec) + if body.Error != "server_error" || body.ErrorDescription != "Internal server issue" { + t.Fatalf("unexpected body: %+v", body) + } + }) + + t.Run("generic oauth error -> 400", func(t *testing.T) { + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { + return srv.AuthInfo{}, oauth.NewOAuthError(oauth.ErrInvalidRequest, "Some OAuth error", "") + }} + rec, nextCalled := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv}, "Bearer valid-token") + if nextCalled || rec.Code != http.StatusBadRequest { + t.Fatalf("expected 400 without next, got %d next=%v", rec.Code, nextCalled) + } + if mv.last != "valid-token" { + t.Fatalf("expected token 'valid-token', got %q", mv.last) + } + body := decodeOAuthResp(t, rec) + if body.Error != "invalid_request" || body.ErrorDescription != "Some OAuth error" { + t.Fatalf("unexpected body: %+v", body) + } + }) + + t.Run("unexpected error -> 401", func(t *testing.T) { + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { + return srv.AuthInfo{}, fmt.Errorf("unexpected error") + }} + rec, nextCalled := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv}, "Bearer valid-token") + if mv.last != "valid-token" { + t.Fatalf("expected token 'valid-token', got %q", mv.last) + } + if nextCalled { + t.Fatalf("expected next not to be called") + } + if rec.Code != http.StatusUnauthorized { + t.Fatalf("expected 401, got %d", rec.Code) + } + body := decodeOAuthResp(t, rec) + if body.Error != "invalid_token" || body.ErrorDescription != "Invalid access token" { + t.Fatalf("unexpected body: %+v", body) + } + hdr := rec.Header().Get("WWW-Authenticate") + if !strings.Contains(hdr, `error="invalid_token"`) || !strings.Contains(hdr, "Invalid access token") { + t.Fatalf("unexpected WWW-Authenticate: %q", hdr) + } + }) +} + +func TestRequireBearerAuth_WithResourceMetadata(t *testing.T) { + url := "https://api.example.com/.well-known/oauth-protected-resource" + + t.Run("401 includes resource_metadata when missing header", func(t *testing.T) { + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { return srv.AuthInfo{}, nil }} + rec, _ := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv, ResourceMetadataURL: &url}, "") + if rec.Code != http.StatusUnauthorized { + t.Fatalf("expected 401, got %d", rec.Code) + } + if mv.last != "" { + t.Fatalf("verifier should not be called, got token %q", mv.last) + } + hdr := rec.Header().Get("WWW-Authenticate") + expectedHdr := `Bearer error="invalid_token", error_description="Missing Authorization header", resource_metadata="` + url + `"` + if hdr != expectedHdr { + t.Fatalf("expected WWW-Authenticate: %q, got %q", expectedHdr, hdr) + } + }) + + t.Run("401 includes resource_metadata when verifier returns invalid_token", func(t *testing.T) { + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { + return srv.AuthInfo{}, oauth.NewOAuthError(oauth.ErrInvalidToken, "Token expired", "") + }} + rec, _ := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv, ResourceMetadataURL: &url}, "Bearer bad") + if rec.Code != http.StatusUnauthorized { + t.Fatalf("expected 401, got %d", rec.Code) + } + hdr := rec.Header().Get("WWW-Authenticate") + if !strings.Contains(hdr, `resource_metadata="`+url+`"`) { + t.Fatalf("resource_metadata missing in header: %q", hdr) + } + }) + + t.Run("403 includes resource_metadata for insufficient scope", func(t *testing.T) { + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { + return srv.AuthInfo{}, oauth.NewOAuthError(oauth.ErrInsufficientScope, "Required scopes: admin", "") + }} + rec, _ := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv, ResourceMetadataURL: &url}, "Bearer t") + if rec.Code != http.StatusForbidden { + t.Fatalf("expected 403, got %d", rec.Code) + } + if hdr := rec.Header().Get("WWW-Authenticate"); !strings.Contains(hdr, `resource_metadata="`+url+`"`) { + t.Fatalf("resource_metadata missing in header: %q", hdr) + } + }) + + t.Run("expired token includes resource_metadata", func(t *testing.T) { + exp := time.Now().Add(-100 * time.Second).Unix() + ai := srv.AuthInfo{Token: "expired", ClientID: "c", Scopes: []string{"read", "write"}, ExpiresAt: &exp} + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { return ai, nil }} + rec, _ := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv, ResourceMetadataURL: &url}, "Bearer expired") + if rec.Code != http.StatusUnauthorized { + t.Fatalf("expected 401, got %d", rec.Code) + } + if hdr := rec.Header().Get("WWW-Authenticate"); !strings.Contains(hdr, `resource_metadata="`+url+`"`) { + t.Fatalf("resource_metadata missing in header: %q", hdr) + } + }) + + t.Run("scope check fail includes resource_metadata", func(t *testing.T) { + exp := time.Now().Add(1 * time.Hour).Unix() + ai := srv.AuthInfo{Token: "ok", ClientID: "c", Scopes: []string{"read"}, ExpiresAt: &exp} + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { return ai, nil }} + rec, _ := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv, RequiredScopes: []string{"read", "write"}, ResourceMetadataURL: &url}, "Bearer ok") + if rec.Code != http.StatusForbidden { + t.Fatalf("expected 403, got %d", rec.Code) + } + if hdr := rec.Header().Get("WWW-Authenticate"); !strings.Contains(hdr, `resource_metadata="`+url+`"`) || !strings.Contains(hdr, `scope="read write"`) { + t.Fatalf("resource_metadata or scope missing in header: %q", hdr) + } + }) + + t.Run("server error does not include WWW-Authenticate header", func(t *testing.T) { + url := "https://api.example.com/.well-known/oauth-protected-resource" + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { + return srv.AuthInfo{}, oauth.NewOAuthError(oauth.ErrServerError, "Internal server issue", "") + }} + rec, nextCalled := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv, ResourceMetadataURL: &url}, "Bearer valid-token") + if mv.last != "valid-token" { + t.Fatalf("expected token 'valid-token', got %q", mv.last) + } + if nextCalled { + t.Fatalf("expected next not to be called") + } + if rec.Code != http.StatusInternalServerError { + t.Fatalf("expected 500, got %d", rec.Code) + } + if hdr := rec.Header().Get("WWW-Authenticate"); hdr != "" { + t.Fatalf("expected no WWW-Authenticate header, got %q", hdr) + } + body := decodeOAuthResp(t, rec) + if body.Error != "server_error" || body.ErrorDescription != "Internal server issue" { + t.Fatalf("expected body {error: \"server_error\", error_description: \"Internal server issue\"}, got %+v", body) + } + }) +} + +func TestRequireBearerAuth_IssuerChecks(t *testing.T) { + t.Run("issuer accepted when matches", func(t *testing.T) { + exp := time.Now().Add(1 * time.Hour).Unix() + ai := srv.AuthInfo{Token: "tok", ClientID: "c", Scopes: []string{"read"}, ExpiresAt: &exp, Extra: map[string]interface{}{"iss": "https://issuer.example"}} + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { return ai, nil }} + rec, nextCalled := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv, Issuer: "https://issuer.example"}, "Bearer tok") + if rec.Code != http.StatusOK || !nextCalled { + t.Fatalf("expected 200 and next called, got %d next=%v", rec.Code, nextCalled) + } + }) + + t.Run("issuer rejected when mismatches", func(t *testing.T) { + exp := time.Now().Add(1 * time.Hour).Unix() + ai := srv.AuthInfo{Token: "tok", ClientID: "c", Scopes: []string{"read"}, ExpiresAt: &exp, Extra: map[string]interface{}{"iss": "https://issuer.example"}} + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { return ai, nil }} + rec, nextCalled := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv, Issuer: "https://another-issuer"}, "Bearer tok") + if nextCalled || rec.Code != http.StatusUnauthorized { + t.Fatalf("expected 401 without next, got %d next=%v", rec.Code, nextCalled) + } + hdr := rec.Header().Get("WWW-Authenticate") + if !strings.Contains(hdr, `error="invalid_token"`) || !strings.Contains(hdr, "Invalid token issuer") { + t.Fatalf("unexpected WWW-Authenticate: %q", hdr) + } + body := decodeOAuthResp(t, rec) + if body.Error != "invalid_token" || body.ErrorDescription != "Invalid token issuer" { + t.Fatalf("unexpected body: %+v", body) + } + }) + + t.Run("issuer check skipped when Extra is nil", func(t *testing.T) { + exp := time.Now().Add(1 * time.Hour).Unix() + ai := srv.AuthInfo{Token: "tok", ClientID: "c", Scopes: []string{"read"}, ExpiresAt: &exp, Extra: nil} + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { return ai, nil }} + rec, nextCalled := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv, Issuer: "https://issuer.example"}, "Bearer tok") + if rec.Code != http.StatusOK || !nextCalled { + t.Fatalf("expected 200 and next called, got %d next=%v", rec.Code, nextCalled) + } + }) + + t.Run("issuer check skipped when iss claim is non-string", func(t *testing.T) { + exp := time.Now().Add(1 * time.Hour).Unix() + ai := srv.AuthInfo{Token: "tok", ClientID: "c", Scopes: []string{"read"}, ExpiresAt: &exp, Extra: map[string]interface{}{"iss": 12345}} + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { return ai, nil }} + rec, nextCalled := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv, Issuer: "https://issuer.example"}, "Bearer tok") + if rec.Code != http.StatusOK || !nextCalled { + t.Fatalf("expected 200 and next called, got %d next=%v", rec.Code, nextCalled) + } + }) + + t.Run("issuer check skipped when iss is empty string", func(t *testing.T) { + exp := time.Now().Add(1 * time.Hour).Unix() + ai := srv.AuthInfo{Token: "tok", ClientID: "c", Scopes: []string{"read"}, ExpiresAt: &exp, Extra: map[string]interface{}{"iss": ""}} + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { return ai, nil }} + rec, nextCalled := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv, Issuer: "https://issuer.example"}, "Bearer tok") + if rec.Code != http.StatusOK || !nextCalled { + t.Fatalf("expected 200 and next called, got %d next=%v", rec.Code, nextCalled) + } + }) +} + +func TestRequireBearerAuth_AudienceChecks(t *testing.T) { + t.Run("audience accepted when matches exactly", func(t *testing.T) { + exp := time.Now().Add(1 * time.Hour).Unix() + // token resource: https://api.example.com/mcp# -> middleware trims trailing '#' + u := mustParseURL(t, "https://api.example.com/mcp#") + ai := srv.AuthInfo{Token: "tok", ClientID: "c", Scopes: []string{"read"}, ExpiresAt: &exp, Resource: u} + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { return ai, nil }} + rec, nextCalled := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv, Audience: []string{"https://api.example.com/mcp"}}, "Bearer tok") + if rec.Code != http.StatusOK || !nextCalled { + t.Fatalf("expected 200 and next called, got %d next=%v", rec.Code, nextCalled) + } + }) + + t.Run("audience rejected when mismatched", func(t *testing.T) { + exp := time.Now().Add(1 * time.Hour).Unix() + u := mustParseURL(t, "https://api.example.com/mcp") + ai := srv.AuthInfo{Token: "tok", ClientID: "c", Scopes: []string{"read"}, ExpiresAt: &exp, Resource: u} + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { return ai, nil }} + rec, nextCalled := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv, Audience: []string{"https://other.example.com/mcp"}}, "Bearer tok") + if nextCalled || rec.Code != http.StatusUnauthorized { + t.Fatalf("expected 401 without next, got %d next=%v", rec.Code, nextCalled) + } + hdr := rec.Header().Get("WWW-Authenticate") + if !strings.Contains(hdr, `error="invalid_token"`) || !strings.Contains(hdr, "Invalid token audience") { + t.Fatalf("unexpected WWW-Authenticate: %q", hdr) + } + body := decodeOAuthResp(t, rec) + if body.Error != "invalid_token" || body.ErrorDescription != "Invalid token audience" { + t.Fatalf("unexpected body: %+v", body) + } + }) + + t.Run("audience accepted when in multi-value list", func(t *testing.T) { + exp := time.Now().Add(1 * time.Hour).Unix() + u := mustParseURL(t, "https://api.example.com/mcp#") + ai := srv.AuthInfo{Token: "tok", ClientID: "c", Scopes: []string{"read"}, ExpiresAt: &exp, Resource: u} + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { return ai, nil }} + rec, nextCalled := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv, Audience: []string{"https://other.example.com", "https://api.example.com/mcp"}}, "Bearer tok") + if rec.Code != http.StatusOK || !nextCalled { + t.Fatalf("expected 200 and next called, got %d next=%v", rec.Code, nextCalled) + } + }) + + t.Run("audience accepted when allowed value has trailing hash", func(t *testing.T) { + exp := time.Now().Add(1 * time.Hour).Unix() + u := mustParseURL(t, "https://api.example.com/mcp") + ai := srv.AuthInfo{Token: "tok", ClientID: "c", Scopes: []string{"read"}, ExpiresAt: &exp, Resource: u} + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { return ai, nil }} + rec, nextCalled := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv, Audience: []string{"https://api.example.com/mcp#"}}, "Bearer tok") + if rec.Code != http.StatusOK || !nextCalled { + t.Fatalf("expected 200 and next called, got %d next=%v", rec.Code, nextCalled) + } + }) + + t.Run("skip audience when Resource is nil", func(t *testing.T) { + exp := time.Now().Add(1 * time.Hour).Unix() + ai := srv.AuthInfo{Token: "tok", ClientID: "c", Scopes: []string{"read"}, ExpiresAt: &exp, Resource: nil} + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { return ai, nil }} + rec, nextCalled := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv, Audience: []string{"https://api.example.com/mcp"}}, "Bearer tok") + if rec.Code != http.StatusOK || !nextCalled { + t.Fatalf("expected 200 and next called, got %d next=%v", rec.Code, nextCalled) + } + }) + + t.Run("skip audience when options.Audience is empty", func(t *testing.T) { + exp := time.Now().Add(1 * time.Hour).Unix() + u := mustParseURL(t, "https://api.example.com/mcp") + ai := srv.AuthInfo{Token: "tok", ClientID: "c", Scopes: []string{"read"}, ExpiresAt: &exp, Resource: u} + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { return ai, nil }} + rec, nextCalled := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv, Audience: []string{}}, "Bearer tok") + if rec.Code != http.StatusOK || !nextCalled { + t.Fatalf("expected 200 and next called, got %d next=%v", rec.Code, nextCalled) + } + }) +} + +func TestRequireBearerAuth_ContextInjectionAndTokenCleared(t *testing.T) { + exp := time.Now().Add(1 * time.Hour).Unix() + ai := srv.AuthInfo{Token: "secret-token", ClientID: "c", Scopes: []string{"read"}, ExpiresAt: &exp} + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { return ai, nil }} + + nextCalled := false + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + nextCalled = true + v := r.Context().Value(AuthInfoKey) + info, ok := v.(srv.AuthInfo) + if !ok { + t.Fatalf("auth info not injected in context") + } + if info.Token != "" { + t.Fatalf("expected token to be cleared, got %q", info.Token) + } + if info.ClientID != "c" || len(info.Scopes) != 1 || info.Scopes[0] != "read" { + t.Fatalf("unexpected auth info: %+v", info) + } + w.WriteHeader(http.StatusOK) + }) + + handler := RequireBearerAuth(BearerAuthMiddlewareOptions{Verifier: mv})(next) + req := httptest.NewRequest(http.MethodGet, "/", nil) + req.Header.Set("Authorization", "Bearer secret-token") + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK || !nextCalled { + t.Fatalf("expected 200 and next called, got %d next=%v", rec.Code, nextCalled) + } +} + +func TestRequireBearerAuth_WWWAuthenticateScopeParamOnInsufficientScope(t *testing.T) { + exp := time.Now().Add(1 * time.Hour).Unix() + ai := srv.AuthInfo{Token: "tok", ClientID: "c", Scopes: []string{"read"}, ExpiresAt: &exp} + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { return ai, nil }} + + opts := BearerAuthMiddlewareOptions{Verifier: mv, RequiredScopes: []string{"read", "write"}} + rec, _ := runWithMiddleware(t, opts, "Bearer tok") + if rec.Code != http.StatusForbidden { + t.Fatalf("expected 403, got %d", rec.Code) + } + hdr := rec.Header().Get("WWW-Authenticate") + if !strings.Contains(hdr, `error="insufficient_scope"`) { + t.Fatalf("expected insufficient_scope in header, got %q", hdr) + } + if !strings.Contains(hdr, `scope="read write"`) { + t.Fatalf("expected scope=\"read write\" in header, got %q", hdr) + } +} + +func TestRequireBearerAuth_WWWAuthenticateHeaderCombos(t *testing.T) { + t.Run("invalid_token 401 header has no scope param", func(t *testing.T) { + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { + return srv.AuthInfo{}, oauth.NewOAuthError(oauth.ErrInvalidToken, "Bad token", "") + }} + rec, _ := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv, RequiredScopes: []string{"read", "write"}}, "Bearer bad") + if rec.Code != http.StatusUnauthorized { + t.Fatalf("expected 401, got %d", rec.Code) + } + hdr := rec.Header().Get("WWW-Authenticate") + if strings.Contains(hdr, "scope=") { + t.Fatalf("unexpected scope param in header: %q", hdr) + } + }) + + t.Run("400 invalid_request should not set WWW-Authenticate header", func(t *testing.T) { + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { + return srv.AuthInfo{}, oauth.NewOAuthError(oauth.ErrInvalidRequest, "Bad req", "") + }} + rec, _ := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv}, "Bearer t") + if rec.Code != http.StatusBadRequest { + t.Fatalf("expected 400, got %d", rec.Code) + } + // For 400, middleware should NOT set WWW-Authenticate header + if hdr := rec.Header().Get("WWW-Authenticate"); hdr != "" { + t.Fatalf("expected no WWW-Authenticate header, got %q", hdr) + } + }) +} + +func TestRequireBearerAuth_BearerPrefixCaseInsensitive(t *testing.T) { + exp := time.Now().Add(1 * time.Hour).Unix() + ai := srv.AuthInfo{Token: "tok", ClientID: "c", Scopes: []string{"read"}, ExpiresAt: &exp} + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { return ai, nil }} + + // Use mixed-case prefix "BeArEr" + rec, nextCalled := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv}, "BeArEr tok") + if mv.last != "tok" { + t.Fatalf("expected verifier to receive token 'tok', got %q", mv.last) + } + if rec.Code != http.StatusOK || !nextCalled { + t.Fatalf("expected 200 and next called, got %d next=%v", rec.Code, nextCalled) + } +} + +func TestRequireBearerAuth_IssuerEdgeCases(t *testing.T) { + t.Run("issuer check skipped when Extra=nil", func(t *testing.T) { + exp := time.Now().Add(1 * time.Hour).Unix() + ai := srv.AuthInfo{Token: "tok", ClientID: "c", Scopes: []string{"read"}, ExpiresAt: &exp, Extra: nil} + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { return ai, nil }} + rec, nextCalled := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv, Issuer: "https://issuer.example"}, "Bearer tok") + if rec.Code != http.StatusOK || !nextCalled { + t.Fatalf("expected 200 and next called, got %d next=%v", rec.Code, nextCalled) + } + }) + + t.Run("issuer check skipped when iss is not a string", func(t *testing.T) { + exp := time.Now().Add(1 * time.Hour).Unix() + ai := srv.AuthInfo{Token: "tok", ClientID: "c", Scopes: []string{"read"}, ExpiresAt: &exp, Extra: map[string]interface{}{"iss": 123}} + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { return ai, nil }} + rec, nextCalled := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv, Issuer: "https://issuer.example"}, "Bearer tok") + if rec.Code != http.StatusOK || !nextCalled { + t.Fatalf("expected 200 and next called, got %d next=%v", rec.Code, nextCalled) + } + }) + + t.Run("issuer check skipped when iss is empty string", func(t *testing.T) { + exp := time.Now().Add(1 * time.Hour).Unix() + ai := srv.AuthInfo{Token: "tok", ClientID: "c", Scopes: []string{"read"}, ExpiresAt: &exp, Extra: map[string]interface{}{"iss": ""}} + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { return ai, nil }} + rec, nextCalled := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv, Issuer: "https://issuer.example"}, "Bearer tok") + if rec.Code != http.StatusOK || !nextCalled { + t.Fatalf("expected 200 and next called, got %d next=%v", rec.Code, nextCalled) + } + }) +} + +func TestRequireBearerAuth_NoWWWAuthenticateOn400(t *testing.T) { + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { + return srv.AuthInfo{}, oauth.NewOAuthError(oauth.ErrInvalidRequest, "invalid input", "") + }} + rec, nextCalled := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv}, "Bearer any") + if nextCalled || rec.Code != http.StatusBadRequest { + t.Fatalf("expected 400 without next, got %d next=%v", rec.Code, nextCalled) + } + if hdr := rec.Header().Get("WWW-Authenticate"); hdr != "" { + t.Fatalf("expected no WWW-Authenticate header, got %q", hdr) + } +} + +func TestRequireBearerAuth_InsufficientScopeFromVerifierIncludesScopeParam(t *testing.T) { + exp := time.Now().Add(1 * time.Hour).Unix() + ai := srv.AuthInfo{Token: "tok", ClientID: "c", Scopes: []string{"read"}, ExpiresAt: &exp} + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { + return ai, oauth.NewOAuthError(oauth.ErrInsufficientScope, "need read write", "") + }} + rec, _ := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv, RequiredScopes: []string{"read", "write"}}, "Bearer tok") + if rec.Code != http.StatusForbidden { + t.Fatalf("expected 403, got %d", rec.Code) + } + hdr := rec.Header().Get("WWW-Authenticate") + if !strings.Contains(hdr, `error="insufficient_scope"`) { + t.Fatalf("expected insufficient_scope in header, got %q", hdr) + } + if !strings.Contains(hdr, `scope="read write"`) { + t.Fatalf("expected scope=\"read write\" in header, got %q", hdr) + } +} + +func TestRequireBearerAuth_NoWWWAuthenticateOn400WithMetadata(t *testing.T) { + urlStr := "https://api.example.com/.well-known/oauth-protected-resource" + mv := &mockVerifier{verify: func(ctx context.Context, token string) (srv.AuthInfo, error) { + return srv.AuthInfo{}, oauth.NewOAuthError(oauth.ErrInvalidRequest, "invalid input", "") + }} + rec, nextCalled := runWithMiddleware(t, BearerAuthMiddlewareOptions{Verifier: mv, ResourceMetadataURL: &urlStr}, "Bearer any") + if nextCalled || rec.Code != http.StatusBadRequest { + t.Fatalf("expected 400 without next, got %d next=%v", rec.Code, nextCalled) + } + if hdr := rec.Header().Get("WWW-Authenticate"); hdr != "" { + t.Fatalf("expected no WWW-Authenticate header, got %q", hdr) + } +} + +// mustParseURL is a small helper for building *url.URL in tests +func mustParseURL(t *testing.T, raw string) *url.URL { + t.Helper() + u, err := url.Parse(raw) + if err != nil { + t.Fatalf("parse url: %v", err) + } + return u +} diff --git a/internal/auth/server/middleware/client_auth.go b/internal/auth/server/middleware/client_auth.go new file mode 100644 index 0000000..1e559f4 --- /dev/null +++ b/internal/auth/server/middleware/client_auth.go @@ -0,0 +1,199 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package middleware + +import ( + "bytes" + "context" + "encoding/base64" + "encoding/json" + "io" + "net/http" + "net/url" + "strings" + "time" + + "trpc.group/trpc-go/trpc-mcp-go/internal/auth" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth/server" + "trpc.group/trpc-go/trpc-mcp-go/internal/errors" +) + +// ClientAuthenticationMiddlewareOptions contains options for client authentication middleware +type ClientAuthenticationMiddlewareOptions struct { + // ClientsStore is a store used to read information about registered OAuth clients + ClientsStore server.OAuthClientsStoreInterface + // Optional: When grant_type=refresh_token and client_id is not provided, try to parse/reverse-check client_id from refresh_token + ResolveClientIDFromRefreshToken func(refreshToken string) (clientID string, ok bool) +} + +// ClientAuthenticatedRequest represents the request schema for client authentication +type ClientAuthenticatedRequest struct { + ClientID string `json:"client_id"` + ClientSecret string `json:"client_secret,omitempty"` +} + +// clientInfoKeyType used to identify the context key storing OAuthClientInformationFull +type clientInfoKeyType struct{} + +// validateClientRequest validates the client authentication request +func validateClientRequest(req *ClientAuthenticatedRequest) error { + if req.ClientID == "" { + return errors.NewOAuthError(errors.ErrInvalidRequest, "client_id is required", "") + } + return nil +} + +// AuthenticateClient returns an HTTP middleware function for client authentication +func AuthenticateClient(options ClientAuthenticationMiddlewareOptions) func(http.Handler) http.Handler { + return func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + setErrorResponse := func(w http.ResponseWriter, err errors.OAuthError, clientID string) { + var statusCode int + switch err.ErrorCode { + case errors.ErrInvalidClient.Error(): + statusCode = http.StatusUnauthorized + case errors.ErrInvalidRequest.Error(): + statusCode = http.StatusBadRequest + case errors.ErrServerError.Error(): + statusCode = http.StatusInternalServerError + default: + statusCode = http.StatusBadRequest + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(statusCode) + _ = json.NewEncoder(w).Encode(err.ToResponseStruct()) + } + + var reqData ClientAuthenticatedRequest + var clientID string + var bodyBytes []byte + + // Priority: Basic Auth first + if authz := r.Header.Get("Authorization"); strings.HasPrefix(strings.ToLower(authz), "basic ") { + enc := strings.TrimSpace(authz[len("Basic "):]) + raw, decErr := base64.StdEncoding.DecodeString(enc) + if decErr != nil { + setErrorResponse(w, errors.NewOAuthError(errors.ErrInvalidClient, "malformed basic credentials", ""), "") + return + } + parts := strings.SplitN(string(raw), ":", 2) + if len(parts) != 2 { + setErrorResponse(w, errors.NewOAuthError(errors.ErrInvalidClient, "malformed basic credentials", ""), "") + return + } + reqData.ClientID, reqData.ClientSecret = parts[0], parts[1] + clientID = reqData.ClientID + } else { + // Non-Basic: buffer and restore Body, support form or JSON + bodyBytes, _ = io.ReadAll(r.Body) + _ = r.Body.Close() + r.Body = io.NopCloser(bytes.NewReader(bodyBytes)) + + ct := strings.ToLower(r.Header.Get("Content-Type")) + switch { + case strings.HasPrefix(ct, "application/x-www-form-urlencoded"): + formVals, _ := url.ParseQuery(string(bodyBytes)) + reqData.ClientID = formVals.Get("client_id") + reqData.ClientSecret = formVals.Get("client_secret") + clientID = reqData.ClientID + case strings.HasPrefix(ct, "application/json"): + if err := json.Unmarshal(bodyBytes, &reqData); err != nil { + setErrorResponse(w, errors.NewOAuthError(errors.ErrInvalidRequest, "Invalid request body", ""), "") + return + } + clientID = reqData.ClientID + default: + // Unknown type: maintain compatibility behavior, treat as JSON decode error + setErrorResponse(w, errors.NewOAuthError(errors.ErrInvalidRequest, "Invalid request body", ""), "") + return + } + } + + // Only try to fall back when client_id is not obtained, and it is a form or JSON, and grant_type=refresh_token + if reqData.ClientID == "" { + ct := strings.ToLower(r.Header.Get("Content-Type")) + var grantType, refreshToken string + + switch { + case strings.HasPrefix(ct, "application/x-www-form-urlencoded"): + formVals, _ := url.ParseQuery(string(bodyBytes)) + grantType = formVals.Get("grant_type") + refreshToken = formVals.Get("refresh_token") + + case strings.HasPrefix(ct, "application/json"): + type raw struct { + GrantType string `json:"grant_type"` + RefreshToken string `json:"refresh_token"` + } + var v raw + _ = json.Unmarshal(bodyBytes, &v) + grantType = v.GrantType + refreshToken = v.RefreshToken + } + + if strings.EqualFold(grantType, "refresh_token") && refreshToken != "" && options.ResolveClientIDFromRefreshToken != nil { + if cid, ok := options.ResolveClientIDFromRefreshToken(refreshToken); ok && cid != "" { + reqData.ClientID = cid + clientID = cid + } + } + } + + // Validate client_id + if err := validateClientRequest(&reqData); err != nil { + if oauthErr, ok := err.(errors.OAuthError); ok { + setErrorResponse(w, oauthErr, clientID) + } else { + setErrorResponse(w, errors.NewOAuthError(errors.ErrInvalidRequest, "Invalid client_id", ""), clientID) + } + return + } + + // Read client and validate secret/expiration + client, err := options.ClientsStore.GetClient(reqData.ClientID) + if err != nil { + setErrorResponse(w, errors.NewOAuthError(errors.ErrInvalidClient, "invalid client credentials", ""), clientID) + return + } + if client == nil { + setErrorResponse(w, errors.NewOAuthError(errors.ErrInvalidClient, "invalid client credentials", ""), clientID) + return + } + if client.ClientSecret != "" { + if reqData.ClientSecret == "" { + setErrorResponse(w, errors.NewOAuthError(errors.ErrInvalidClient, "Client secret is required", ""), clientID) + return + } + if client.ClientSecret != reqData.ClientSecret { + setErrorResponse(w, errors.NewOAuthError(errors.ErrInvalidClient, "Invalid client_secret", ""), clientID) + return + } + if client.ClientSecretExpiresAt != nil { + now := time.Now().Unix() + if *client.ClientSecretExpiresAt != 0 && *client.ClientSecretExpiresAt < now { + setErrorResponse(w, errors.NewOAuthError(errors.ErrInvalidClient, "Client secret has expired", ""), clientID) + return + } + } + } + + ctx := context.WithValue(r.Context(), clientInfoKeyType{}, client) + next.ServeHTTP(w, r.WithContext(ctx)) + }) + } +} + +// GetAuthenticatedClient retrieves the authenticated client from HTTP request context +func GetAuthenticatedClient(r *http.Request) (*auth.OAuthClientInformationFull, bool) { + client := r.Context().Value(clientInfoKeyType{}) + if client == nil { + return nil, false + } + + authenticatedClient, ok := client.(*auth.OAuthClientInformationFull) + return authenticatedClient, ok +} diff --git a/internal/auth/server/middleware/client_auth_test.go b/internal/auth/server/middleware/client_auth_test.go new file mode 100644 index 0000000..b1201ec --- /dev/null +++ b/internal/auth/server/middleware/client_auth_test.go @@ -0,0 +1,208 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package middleware + +import ( + "bytes" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "testing" + "time" + + "trpc.group/trpc-go/trpc-mcp-go/internal/auth" + srv "trpc.group/trpc-go/trpc-mcp-go/internal/auth/server" + oauth "trpc.group/trpc-go/trpc-mcp-go/internal/errors" +) + +// mockClientsStore is a test double for OAuthClientsStoreInterface that lets tests +// control lookups via a provided function +type mockClientsStore struct { + get func(clientID string) (*auth.OAuthClientInformationFull, error) +} + +// GetClient returns the mocked client info for the given client ID +func (m *mockClientsStore) GetClient(clientID string) (*auth.OAuthClientInformationFull, error) { + return m.get(clientID) +} + +// RegisterClient indicates dynamic client registration is not supported in this mock +func (m *mockClientsStore) RegisterClient(client auth.OAuthClientInformationFull) (*auth.OAuthClientInformationFull, error) { + return nil, fmt.Errorf("dynamic client registration is not supported") +} + +// runClientAuth executes a request through the client-authentication middleware and +// returns the recorder and whether the next handler was called +func runClientAuth(t *testing.T, store srv.OAuthClientsStoreInterface, body interface{}, contentType string) (rec *httptest.ResponseRecorder, nextCalled bool) { + t.Helper() + + handlerCalled := false + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + handlerCalled = true + cli, ok := GetAuthenticatedClient(r) + if !ok { + t.Fatalf("authenticated client not found in context") + } + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(struct { + Success bool `json:"success"` + Client *auth.OAuthClientInformationFull `json:"client"` + }{Success: true, Client: cli}) + }) + + middleware := AuthenticateClient(ClientAuthenticationMiddlewareOptions{ClientsStore: store}) + h := middleware(next) + + var b []byte + switch v := body.(type) { + case []byte: + b = v + default: + var err error + b, err = json.Marshal(v) + if err != nil { + t.Fatalf("failed to marshal body: %v", err) + } + } + + req := httptest.NewRequest(http.MethodPost, "/protected", bytes.NewReader(b)) + if contentType != "" { + req.Header.Set("Content-Type", contentType) + } else { + req.Header.Set("Content-Type", "application/json") + } + rec = httptest.NewRecorder() + h.ServeHTTP(rec, req) + return rec, handlerCalled +} + +// decodeOAuthError parses an OAuthErrorResponse from the test recorder body +func decodeOAuthError(t *testing.T, rec *httptest.ResponseRecorder) *oauth.OAuthErrorResponse { + t.Helper() + var v oauth.OAuthErrorResponse + _ = json.NewDecoder(rec.Body).Decode(&v) + return &v +} + +func TestAuthenticateClient_ValidCredentials(t *testing.T) { + store := &mockClientsStore{get: func(clientID string) (*auth.OAuthClientInformationFull, error) { + if clientID == "valid-client" { + return &auth.OAuthClientInformationFull{OAuthClientMetadata: auth.OAuthClientMetadata{RedirectURIs: []string{"https://example.com/callback"}}, OAuthClientInformation: auth.OAuthClientInformation{ClientID: "valid-client", ClientSecret: "valid-secret"}}, nil + } + return nil, nil + }} + + rec, nextCalled := runClientAuth(t, store, map[string]interface{}{"client_id": "valid-client", "client_secret": "valid-secret"}, "application/json") + if rec.Code != http.StatusOK || !nextCalled { + t.Fatalf("expected 200 and next called, got %d next=%v", rec.Code, nextCalled) + } + + var body struct { + Success bool `json:"success"` + Client struct { + ClientID string `json:"client_id"` + } `json:"client"` + } + _ = json.NewDecoder(rec.Body).Decode(&body) + if !body.Success || body.Client.ClientID != "valid-client" { + t.Fatalf("unexpected body: %+v", body) + } +} + +func TestAuthenticateClient_InvalidClientID(t *testing.T) { + store := &mockClientsStore{get: func(clientID string) (*auth.OAuthClientInformationFull, error) { return nil, nil }} + rec, nextCalled := runClientAuth(t, store, map[string]interface{}{"client_id": "non-existent-client", "client_secret": "some-secret"}, "application/json") + if nextCalled || rec.Code != http.StatusUnauthorized { + t.Fatalf("expected 401 without next, got %d next=%v", rec.Code, nextCalled) + } + body := decodeOAuthError(t, rec) + if body.Error != "invalid_client" || body.ErrorDescription != "invalid client credentials" { + t.Fatalf("unexpected body: %+v", body) + } +} + +func TestAuthenticateClient_InvalidClientSecret(t *testing.T) { + store := &mockClientsStore{get: func(clientID string) (*auth.OAuthClientInformationFull, error) { + if clientID == "valid-client" { + return &auth.OAuthClientInformationFull{OAuthClientInformation: auth.OAuthClientInformation{ClientID: "valid-client", ClientSecret: "valid-secret"}}, nil + } + return nil, nil + }} + rec, nextCalled := runClientAuth(t, store, map[string]interface{}{"client_id": "valid-client", "client_secret": "wrong-secret"}, "application/json") + if nextCalled || rec.Code != http.StatusUnauthorized { + t.Fatalf("expected 401 without next, got %d next=%v", rec.Code, nextCalled) + } + body := decodeOAuthError(t, rec) + if body.Error != "invalid_client" || body.ErrorDescription != "Invalid client_secret" { + t.Fatalf("unexpected body: %+v", body) + } +} + +func TestAuthenticateClient_MissingClientID(t *testing.T) { + store := &mockClientsStore{get: func(clientID string) (*auth.OAuthClientInformationFull, error) { return nil, nil }} + rec, nextCalled := runClientAuth(t, store, map[string]interface{}{"client_secret": "valid-secret"}, "application/json") + if nextCalled || rec.Code != http.StatusBadRequest { + t.Fatalf("expected 400 without next, got %d next=%v", rec.Code, nextCalled) + } + body := decodeOAuthError(t, rec) + if body.Error != "invalid_request" { + t.Fatalf("unexpected body: %+v", body) + } +} + +func TestAuthenticateClient_AllowsMissingSecretIfClientHasNone(t *testing.T) { + store := &mockClientsStore{get: func(clientID string) (*auth.OAuthClientInformationFull, error) { + if clientID == "expired-client" { + return &auth.OAuthClientInformationFull{OAuthClientInformation: auth.OAuthClientInformation{ClientID: "expired-client"}}, nil + } + return nil, nil + }} + rec, nextCalled := runClientAuth(t, store, map[string]interface{}{"client_id": "expired-client"}, "application/json") + if rec.Code != http.StatusOK || !nextCalled { + t.Fatalf("expected 200 and next called, got %d next=%v", rec.Code, nextCalled) + } +} + +func TestAuthenticateClient_RejectsExpiredSecret(t *testing.T) { + past := time.Now().Add(-1 * time.Hour).Unix() + store := &mockClientsStore{get: func(clientID string) (*auth.OAuthClientInformationFull, error) { + if clientID == "client-with-expired-secret" { + return &auth.OAuthClientInformationFull{OAuthClientInformation: auth.OAuthClientInformation{ClientID: "client-with-expired-secret", ClientSecret: "expired-secret", ClientSecretExpiresAt: &past}}, nil + } + return nil, nil + }} + rec, nextCalled := runClientAuth(t, store, map[string]interface{}{"client_id": "client-with-expired-secret", "client_secret": "expired-secret"}, "application/json") + if nextCalled || rec.Code != http.StatusUnauthorized { + t.Fatalf("expected 401 without next, got %d next=%v", rec.Code, nextCalled) + } + body := decodeOAuthError(t, rec) + if body.Error != "invalid_client" || body.ErrorDescription != "Client secret has expired" { + t.Fatalf("unexpected body: %+v", body) + } +} + +func TestAuthenticateClient_MalformedRequestBody(t *testing.T) { + store := &mockClientsStore{get: func(clientID string) (*auth.OAuthClientInformationFull, error) { return nil, nil }} + rec, nextCalled := runClientAuth(t, store, []byte("not-json-format"), "application/json") + if nextCalled || rec.Code != http.StatusBadRequest { + t.Fatalf("expected 400 without next, got %d next=%v", rec.Code, nextCalled) + } +} + +func TestAuthenticateClient_IgnoresExtraFields(t *testing.T) { + store := &mockClientsStore{get: func(clientID string) (*auth.OAuthClientInformationFull, error) { + if clientID == "valid-client" { + return &auth.OAuthClientInformationFull{OAuthClientInformation: auth.OAuthClientInformation{ClientID: "valid-client", ClientSecret: "valid-secret"}}, nil + } + return nil, nil + }} + rec, nextCalled := runClientAuth(t, store, map[string]interface{}{"client_id": "valid-client", "client_secret": "valid-secret", "extra_field": "ignored"}, "application/json") + if rec.Code != http.StatusOK || !nextCalled { + t.Fatalf("expected 200 and next called, got %d next=%v", rec.Code, nextCalled) + } +} diff --git a/internal/auth/server/middleware/context.go b/internal/auth/server/middleware/context.go new file mode 100644 index 0000000..6eb3d2f --- /dev/null +++ b/internal/auth/server/middleware/context.go @@ -0,0 +1,14 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package middleware + +// authInfoKeyType is an unexported empty struct used as a context key to prevent collisions with other packages +type authInfoKeyType struct{} + +// AuthInfoKey is the context key for storing and retrieving authentication information on requests +// Use context.WithValue(ctx, AuthInfoKey, authInfo) to attach and ctx.Value(AuthInfoKey) to read +var AuthInfoKey = authInfoKeyType{} diff --git a/internal/auth/server/middleware/security.go b/internal/auth/server/middleware/security.go new file mode 100644 index 0000000..bc4015e --- /dev/null +++ b/internal/auth/server/middleware/security.go @@ -0,0 +1,249 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package middleware + +import ( + "encoding/json" + "fmt" + "net/http" + "strings" + + "golang.org/x/time/rate" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth/server" + "trpc.group/trpc-go/trpc-mcp-go/internal/errors" +) + +// SecurityMiddlewareOption holds pluggable security dependencies for middleware +type SecurityMiddlewareOption struct { + verifier server.TokenVerifier // Token verifier used by middleware that need direct verification +} + +// Authorizer defines a unified authorization decision interface +// Implementations should return nil when access is allowed and an OAuthError when denied +type Authorizer interface { + Authorize(authInfo server.AuthInfo, resource string, action string) error +} + +// ScopePermissionMapper maps OAuth scopes to internal permissions +// The returned slice represents permissions granted by the provided scopes +type ScopePermissionMapper interface { + MapScopes(scopes []string) []string +} + +// DefaultScopeMapper is a simple mapper using a static scope→permissions table +type DefaultScopeMapper struct { + Mapping map[string][]string // Per-scope permission list +} + +// MapScopes expands scopes into a flattened permission list using Mapping +func (m *DefaultScopeMapper) MapScopes(scopes []string) []string { + // Accumulate permissions granted by each scope + var perms []string + for _, scope := range scopes { + if mapped, ok := m.Mapping[scope]; ok { + perms = append(perms, mapped...) + } + } + return perms +} + +// PolicyAuthorizer authorizes by checking if the required permission exists after scope mapping +type PolicyAuthorizer struct { + ScopeMapper ScopePermissionMapper // Pluggable scope→permission mapper +} + +// Authorize checks whether authInfo scopes grant the required {resource}:{action} permission +func (a *PolicyAuthorizer) Authorize(authInfo server.AuthInfo, resource string, action string) error { + // Convert scopes to internal permissions + perms := a.ScopeMapper.MapScopes(authInfo.Scopes) + + // Build required permission string like urn:mcp:workspace:xyz:read + required := fmt.Sprintf("%s:%s", resource, action) + + // Return success when permission is present + for _, p := range perms { + if p == required { + return nil + } + } + + // Otherwise return standardized insufficient_scope error + return errors.NewOAuthError(errors.ErrInsufficientScope, + fmt.Sprintf("Missing permission %s", required), "") +} + +// responseWriterWithStatus wraps http.ResponseWriter to capture the final status code +type responseWriterWithStatus struct { + http.ResponseWriter + statusCode int +} + +// WriteHeader intercepts WriteHeader calls to store the status code +func (rw *responseWriterWithStatus) WriteHeader(code int) { + rw.statusCode = code + rw.ResponseWriter.WriteHeader(code) +} + +// CorsMiddleware applies permissive CORS headers similar to express default behavior +// It returns 204 for OPTIONS preflight while forwarding non-preflight requests downstream +func CorsMiddleware(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Read Origin header to detect cross origin requests + origin := r.Header.Get("Origin") + if origin == "" { + // Not a CORS request so proceed without CORS headers + next.ServeHTTP(w, r) + return + } + + // Set basic CORS headers + w.Header().Set("Access-Control-Allow-Origin", "*") + w.Header().Set("Access-Control-Allow-Methods", "GET,HEAD,PUT,PATCH,POST,DELETE") + + // Handle preflight with 204 and zero content length + if r.Method == http.MethodOptions { + w.Header().Set("Content-Length", "0") + w.WriteHeader(http.StatusNoContent) + return + } + + // Forward actual request + next.ServeHTTP(w, r) + }) +} + +// RateLimitMiddleware applies a token bucket limiter to incoming requests +// When the limiter denies a request a 429 JSON OAuth error is returned +func RateLimitMiddleware(limiter *rate.Limiter) func(http.Handler) http.Handler { + return func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Short circuit when the limiter does not allow the request + if !limiter.Allow() { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusTooManyRequests) + + // Build standardized OAuth error payload + tooManyRequestsError := errors.NewOAuthError( + errors.ErrTooManyRequests, + "You have exceeded the rate limit for token revocation requests", + "", + ) + _ = json.NewEncoder(w).Encode(tooManyRequestsError.ToResponseStruct()) + return + } + + // Continue to next handler + next.ServeHTTP(w, r) + }) + } +} + +// ContentTypeValidationMiddleware validates the Content-Type header against an allowlist +// When allowJSONFallback is true application/json is accepted in addition to allowedTypes[0] +func ContentTypeValidationMiddleware(allowedTypes []string, allowJSONFallback bool) func(http.Handler) http.Handler { + return func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + contentType := r.Header.Get("Content-Type") + + // Content-Type header is required for these endpoints + if contentType == "" { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + + invalidReqError := errors.NewOAuthError( + errors.ErrInvalidRequest, + "Content-Type header is required", + "", + ) + _ = json.NewEncoder(w).Encode(invalidReqError.ToResponseStruct()) + return + } + + // Check prefix match to allow charset parameters + var isValid bool + for _, allowedType := range allowedTypes { + if strings.HasPrefix(contentType, allowedType) { + isValid = true + break + } + } + + // Optionally accept JSON when configured + if !isValid && allowJSONFallback && strings.HasPrefix(contentType, "application/json") { + isValid = true + } + + // Reject unsupported content types with a helpful message + if !isValid { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + + errorMsg := fmt.Sprintf("Content-Type must be one of: %s", strings.Join(allowedTypes, ", ")) + if allowJSONFallback && len(allowedTypes) > 0 { + errorMsg = fmt.Sprintf("Content-Type must be %s (preferred) or application/json", allowedTypes[0]) + } + + invalidReqError := errors.NewOAuthError( + errors.ErrInvalidRequest, + errorMsg, + "", + ) + _ = json.NewEncoder(w).Encode(invalidReqError.ToResponseStruct()) + return + } + + // Forward to the next handler + next.ServeHTTP(w, r) + }) + } +} + +// URLEncodedValidationMiddleware enforces application/x-www-form-urlencoded for RFC 7009 style endpoints +func URLEncodedValidationMiddleware(allowJSONFallback bool) func(http.Handler) http.Handler { + return ContentTypeValidationMiddleware([]string{"application/x-www-form-urlencoded"}, allowJSONFallback) +} + +// JSONValidationMiddleware enforces application/json for endpoints that only accept JSON +func JSONValidationMiddleware() func(http.Handler) http.Handler { + return ContentTypeValidationMiddleware([]string{"application/json"}, false) +} + +// AuthorizationMiddleware performs authorization using the provided Authorizer for a {resource, action} pair +// It requires a validated AuthInfo in context and returns 401 or 403 with OAuth style error when denied +func AuthorizationMiddleware(authorizer Authorizer, resource string, action string) func(http.Handler) http.Handler { + return func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Read validated auth info placed in context by upstream auth middleware + authInfo, ok := GetAuthInfo(r.Context()) + if !ok { + w.Header().Set("WWW-Authenticate", `Bearer error="invalid_token", error_description="No authentication info found"`) + http.Error(w, "Unauthorized", http.StatusUnauthorized) + return + } + + // Evaluate authorization decision for the requested resource and action + err := authorizer.Authorize(authInfo, resource, action) + if err != nil { + // Return standardized insufficient_scope response + w.Header().Set("Content-Type", "application/json") + w.Header().Set("WWW-Authenticate", `Bearer error="insufficient_scope"`) + w.WriteHeader(http.StatusForbidden) + _ = json.NewEncoder(w).Encode(err.(errors.OAuthError).ToResponseStruct()) + + // Optionally extract subject for audit or side effects + _ = extractSubject(authInfo) + return + } + + // Optionally extract subject for audit or side effects + _ = extractSubject(authInfo) + + // Authorized so continue to next handler + next.ServeHTTP(w, r) + }) + } +} diff --git a/internal/auth/server/middleware/security_test.go b/internal/auth/server/middleware/security_test.go new file mode 100644 index 0000000..261e8cd --- /dev/null +++ b/internal/auth/server/middleware/security_test.go @@ -0,0 +1,233 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package middleware + +import ( + "bytes" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "golang.org/x/time/rate" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth/server" +) + +// okHandler returns a simple HTTP handler that always responds with status 200 OK and the body "ok" +func okHandler() http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`ok`)) + }) +} + +// do executes the given handler with a constructed HTTP request +func do(handler http.Handler, method, path string, body io.Reader, headers map[string]string) *httptest.ResponseRecorder { + req := httptest.NewRequest(method, path, body) + for k, v := range headers { + req.Header.Set(k, v) + } + rr := httptest.NewRecorder() + handler.ServeHTTP(rr, req) + return rr +} + +func TestCorsMiddleware_NoOrigin_PassThrough(t *testing.T) { + h := CorsMiddleware(okHandler()) + rr := do(h, http.MethodGet, "/", nil, nil) + + if rr.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", rr.Code) + } + if v := rr.Header().Get("Access-Control-Allow-Origin"); v != "" { + t.Fatalf("expected no CORS header, got %q", v) + } +} + +func TestCorsMiddleware_WithOrigin_SetsHeaders(t *testing.T) { + h := CorsMiddleware(okHandler()) + rr := do(h, http.MethodGet, "/", nil, map[string]string{ + "Origin": "http://example.com", + }) + + if rr.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", rr.Code) + } + if v := rr.Header().Get("Access-Control-Allow-Origin"); v != "*" { + t.Fatalf("expected '*', got %q", v) + } + if v := rr.Header().Get("Access-Control-Allow-Methods"); !strings.Contains(v, "GET") { + t.Fatalf("expected allow methods set, got %q", v) + } +} + +func TestCorsMiddleware_Options_Preflight(t *testing.T) { + h := CorsMiddleware(okHandler()) + rr := do(h, http.MethodOptions, "/", nil, map[string]string{ + "Origin": "http://example.com", + }) + + if rr.Code != http.StatusNoContent { + t.Fatalf("expected 204, got %d", rr.Code) + } + if v := rr.Header().Get("Content-Length"); v != "0" { + t.Fatalf("expected Content-Length 0, got %q", v) + } +} + +func TestRateLimitMiddleware_Allow(t *testing.T) { + lim := rate.NewLimiter(rate.Inf, 0) // always allow + h := RateLimitMiddleware(lim)(okHandler()) + + rr := do(h, http.MethodPost, "/token", nil, nil) + if rr.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", rr.Code) + } +} + +func TestRateLimitMiddleware_Reject(t *testing.T) { + lim := rate.NewLimiter(0, 0) // always reject + h := RateLimitMiddleware(lim)(okHandler()) + + rr := do(h, http.MethodPost, "/token", nil, nil) + if rr.Code != http.StatusTooManyRequests { + t.Fatalf("expected 429, got %d", rr.Code) + } + if !strings.Contains(strings.ToLower(rr.Body.String()), "too_many_requests") { + t.Fatalf("expected body to contain too_many_requests, got %q", rr.Body.String()) + } +} + +func TestContentTypeValidation_MissingHeader(t *testing.T) { + h := ContentTypeValidationMiddleware([]string{"application/x-www-form-urlencoded"}, false)(okHandler()) + + rr := do(h, http.MethodPost, "/token", bytes.NewBufferString("a=b"), nil) + if rr.Code != http.StatusBadRequest { + t.Fatalf("expected 400, got %d", rr.Code) + } + if !strings.Contains(strings.ToLower(rr.Body.String()), "invalid_request") { + t.Fatalf("expected invalid_request, got %q", rr.Body.String()) + } +} + +func TestContentTypeValidation_Allowed(t *testing.T) { + h := ContentTypeValidationMiddleware([]string{"application/x-www-form-urlencoded"}, false)(okHandler()) + + rr := do(h, http.MethodPost, "/token", bytes.NewBufferString("a=b"), map[string]string{ + "Content-Type": "application/x-www-form-urlencoded; charset=utf-8", + }) + if rr.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", rr.Code) + } +} + +func TestContentTypeValidation_NotAllowed_NoFallback(t *testing.T) { + h := ContentTypeValidationMiddleware([]string{"application/x-www-form-urlencoded"}, false)(okHandler()) + + rr := do(h, http.MethodPost, "/token", bytes.NewBufferString(`{"a":"b"}`), map[string]string{ + "Content-Type": "application/json", + }) + if rr.Code != http.StatusBadRequest { + t.Fatalf("expected 400, got %d", rr.Code) + } +} + +func TestContentTypeValidation_JSONFallback(t *testing.T) { + h := ContentTypeValidationMiddleware([]string{"application/x-www-form-urlencoded"}, true)(okHandler()) + + rr := do(h, http.MethodPost, "/token", bytes.NewBufferString(`{"a":"b"}`), map[string]string{ + "Content-Type": "application/json; charset=utf-8", + }) + if rr.Code != http.StatusOK { + t.Fatalf("expected 200 with JSON fallback, got %d", rr.Code) + } +} + +func TestURLEncodedValidationMiddleware(t *testing.T) { + h := URLEncodedValidationMiddleware(false)(okHandler()) + + rr := do(h, http.MethodPost, "/revoke", bytes.NewBufferString("a=b"), map[string]string{ + "Content-Type": "application/x-www-form-urlencoded", + }) + if rr.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", rr.Code) + } +} + +func TestJSONValidationMiddleware(t *testing.T) { + h := JSONValidationMiddleware()(okHandler()) + + rr := do(h, http.MethodPost, "/register", bytes.NewBufferString(`{"a":"b"}`), map[string]string{ + "Content-Type": "application/json", + }) + if rr.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", rr.Code) + } +} + +type denyAllAuthorizer struct{} + +func (denyAllAuthorizer) Authorize(authInfo server.AuthInfo, resource, action string) error { + // not used in this test since no auth info will be present + return nil +} + +func TestAuthorizationMiddleware_NoAuthInfo_Returns401(t *testing.T) { + h := AuthorizationMiddleware(denyAllAuthorizer{}, "urn:mcp:resource", "read")(okHandler()) + + rr := do(h, http.MethodGet, "/protected", nil, nil) + if rr.Code != http.StatusUnauthorized { + t.Fatalf("expected 401, got %d", rr.Code) + } + if v := rr.Header().Get("WWW-Authenticate"); !strings.Contains(v, "invalid_token") { + t.Fatalf("expected WWW-Authenticate with invalid_token, got %q", v) + } +} + +func TestPolicyAuthorizer_Authorize_Allowed(t *testing.T) { + mapper := &DefaultScopeMapper{ + Mapping: map[string][]string{ + "workspace.read": {"urn:mcp:workspace:xyz:read"}, + }, + } + a := &PolicyAuthorizer{ScopeMapper: mapper} + + authInfo := server.AuthInfo{Scopes: []string{"workspace.read"}} + if err := a.Authorize(authInfo, "urn:mcp:workspace:xyz", "read"); err != nil { + t.Fatalf("expected authorize success, got %v", err) + } +} + +func TestPolicyAuthorizer_Authorize_Denied(t *testing.T) { + mapper := &DefaultScopeMapper{ + Mapping: map[string][]string{ + "workspace.read": {"urn:mcp:workspace:xyz:read"}, + }, + } + a := &PolicyAuthorizer{ScopeMapper: mapper} + + authInfo := server.AuthInfo{Scopes: []string{"workspace.read"}} + if err := a.Authorize(authInfo, "urn:mcp:workspace:xyz", "write"); err == nil { + t.Fatalf("expected authorize deny, got nil error") + } +} + +func TestRateLimitMiddleware_Reject_ReturnsJSON(t *testing.T) { + lim := rate.NewLimiter(0, 0) + h := RateLimitMiddleware(lim)(okHandler()) + + rr := do(h, http.MethodPost, "/token", nil, nil) + if ct := rr.Header().Get("Content-Type"); ct != "application/json" { + t.Fatalf("expected application/json, got %q", ct) + } + var parsed map[string]any + if err := json.Unmarshal(rr.Body.Bytes(), &parsed); err != nil { + t.Fatalf("expected JSON body, got err: %v; body=%q", err, rr.Body.String()) + } +} diff --git a/internal/auth/server/provider.go b/internal/auth/server/provider.go new file mode 100644 index 0000000..78daf2d --- /dev/null +++ b/internal/auth/server/provider.go @@ -0,0 +1,71 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package server + +import ( + "net/http" + "net/url" + + "trpc.group/trpc-go/trpc-mcp-go/internal/auth" +) + +// AuthorizationParams carries parameters for starting an OAuth authorization request +type AuthorizationParams struct { + CodeChallenge string `json:"code_challenge"` // PKCE code challenge from client + RedirectURI string `json:"redirect_uri"` // Redirect URI registered by the client + State string `json:"state"` // Optional opaque value to maintain client state between request and callback + Scopes []string `json:"scopes"` // Optional empty slice means not provided + Resource *url.URL `json:"resource"` // Optional nil means not provided +} + +// OAuthServerProvider defines a complete OAuth 2.1 server interface including client management authorization token exchange verification and revocation +type OAuthServerProvider interface { + + // ClientsStore returns the store used to read registered OAuth client information + ClientsStore() *OAuthClientsStore + + // Authorize starts the authorization flow implemented by this server or by redirecting to another authorization server + // The server must ultimately redirect to the given redirect URI with either a success or an error response per OAuth 2.1 + // On success include query params code and state if provided + // On error include query param error and may include error_description + Authorize(client auth.OAuthClientInformationFull, params AuthorizationParams, res http.ResponseWriter, req *http.Request) error + + // ChallengeForAuthorizationCode returns the codeChallenge that was used when the indicated authorization began + ChallengeForAuthorizationCode(client auth.OAuthClientInformationFull, authorizationCode string) (string, error) + + // ExchangeAuthorizationCode exchanges an authorization code for access tokens + // Validate code and optional PKCE codeVerifier and optional redirectUri and resource + ExchangeAuthorizationCode( + client auth.OAuthClientInformationFull, + authorizationCode string, codeVerifier *string, + redirectUri *string, + resource *url.URL, + ) (*auth.OAuthTokens, error) + + // ExchangeRefreshToken exchanges a refresh token for new access tokens + // Accept optional scopes and optional resource + ExchangeRefreshToken( + client auth.OAuthClientInformationFull, + refreshToken string, + scopes []string, // Optional empty slice if not provided + resource *url.URL, // Optional nil if not provided + ) (*auth.OAuthTokens, error) + + // VerifyAccessToken verifies an access token and returns its associated information + VerifyAccessToken(token string) (*AuthInfo, error) + + // SupportTokenRevocation indicates optional support for token revocation + SupportTokenRevocation +} + +// SupportTokenRevocation defines optional token revocation capability +type SupportTokenRevocation interface { + // RevokeToken revokes an access or refresh token + // If the token is invalid or already revoked this should be a no op + // Optional method + RevokeToken(client auth.OAuthClientInformationFull, request auth.OAuthTokenRevocationRequest) error +} diff --git a/internal/auth/server/providers/proxy.go b/internal/auth/server/providers/proxy.go new file mode 100644 index 0000000..51916dd --- /dev/null +++ b/internal/auth/server/providers/proxy.go @@ -0,0 +1,418 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package providers + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "net/http" + "net/url" + "strings" + + "github.com/go-playground/validator/v10" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth/server" + "trpc.group/trpc-go/trpc-mcp-go/internal/errors" +) + +// ProxyEndpoints defines the OAuth 2.0/2.1 server endpoints used by the proxy. +// It contains the URLs for various OAuth operations. +type ProxyEndpoints struct { + // AuthorizationURL is the URL of the OAuth 2.0/2.1 authorization endpoint. + // This is where users are redirected to authorize the client. + // "https://auth.example.com/authorize" + AuthorizationURL string `json:"authorizationUrl"` + + // TokenURL is the URL of the OAuth 2.0/2.1 token endpoint. + // This is where the client exchanges an authorization code for an access token. + // "https://auth.example.com/token" + TokenURL string `json:"tokenUrl"` + + // RevocationURL is the optional URL of the OAuth 2.0 token revocation endpoint. + // If provided, it's used to revoke access tokens or refresh tokens. + // "https://auth.example.com/revoke" + RevocationURL string `json:"revocationUrl,omitempty"` + + // RegistrationURL is the optional URL of the OAuth 2.0 dynamic client registration endpoint. + // If provided, it allows clients to register with the authorization server dynamically. 。 + // "https://auth.example.com/register" + RegistrationURL string `json:"registrationUrl,omitempty"` +} + +// ProxyOptions defines configuration options for the proxy OAuth server +type ProxyOptions struct { + // Endpoints presents Endpoint configuration for proxy OAuth operations + Endpoints ProxyEndpoints + + // VerifyAccessToken verifies access tokens and return auth info + VerifyAccessToken func(token string) (*server.AuthInfo, error) + + // GetClient fetches client information from the upstream server + GetClient func(clientID string) (*auth.OAuthClientInformationFull, error) + + // Fetch customs fetch implementation used for all network requests, optional + Fetch auth.FetchFunc +} + +// ProxyOAuthServerProvider defines proxy OAuth server provider +type ProxyOAuthServerProvider struct { + // endpoints defines proxy endpoint configuration + endpoints ProxyEndpoints + + // verifyAccessToken verifies access tokens + verifyAccessToken func(token string) (*server.AuthInfo, error) + + // getClient get client's information + getClient func(clientID string) (*auth.OAuthClientInformationFull, error) + + // SkipLocalPkceValidation determines whether to skip local PKCE validation. + // If true, the server will not perform PKCE validation locally and will pass the code_verifier to the upstream server. + // NOTE: This should only be true if the upstream server is performing the actual PKCE validation. + // 可选字段，默认false / Optional field, defaults to false + SkipLocalPkceValidation bool `json:"skipLocalPkceValidation,omitempty"` + + // Custom fetch implementation, optional + fetch auth.FetchFunc +} + +// Authorize handles an OAuth authorization request by constructing the query +// parameters and redirecting the user agent to the configured authorization endpoint. +func (p *ProxyOAuthServerProvider) Authorize(client auth.OAuthClientInformationFull, params server.AuthorizationParams, res http.ResponseWriter, req *http.Request) error { + // Validate the configured authorization endpoint URL + targetURL, err := url.Parse(p.endpoints.AuthorizationURL) + if err != nil { + return fmt.Errorf("invalid authorization URL: %v", err) + } + + // Build required OAuth query parameters + query := url.Values{ + "client_id": {client.ClientID}, + "response_type": {"code"}, + "redirect_uri": {params.RedirectURI}, + "code_challenge": {params.CodeChallenge}, + "code_challenge_method": {"S256"}, + } + + // Add optional parameters when present + if params.State != "" { + query.Set("state", params.State) + } + if len(params.Scopes) > 0 { + query.Set("scope", strings.Join(params.Scopes, " ")) + } + if params.Resource != nil { + query.Set("resource", params.Resource.String()) + } + + // Attach encoded query to the target URL + targetURL.RawQuery = query.Encode() + + // Perform a 302 redirect to the upstream authorization endpoint + http.Redirect(res, req, targetURL.String(), http.StatusFound) + return nil +} + +// VerifyAccessToken proxies token verification to the configured verifier function. +func (p *ProxyOAuthServerProvider) VerifyAccessToken(token string) (*server.AuthInfo, error) { + // Delegate to injected verifier to allow custom verification strategies + return p.verifyAccessToken(token) +} + +// NewProxyOAuthServerProvider creates a new ProxyOAuthServerProvider using the provided options. +// By default SkipLocalPkceValidation is set to true to defer PKCE verification to the upstream server. +func NewProxyOAuthServerProvider(options ProxyOptions) *ProxyOAuthServerProvider { + // Populate provider with endpoints, dependency functions, and optional fetch + provider := &ProxyOAuthServerProvider{ + endpoints: options.Endpoints, + verifyAccessToken: options.VerifyAccessToken, + getClient: options.GetClient, + fetch: options.Fetch, + SkipLocalPkceValidation: true, + } + // Return the ready to use provider + return provider +} + +// doFetch executes an HTTP request using the custom fetch function if provided, +// otherwise falls back to the default HTTP client. +func (p *ProxyOAuthServerProvider) doFetch(req *http.Request) (*http.Response, error) { + // Prefer custom fetch to allow callers to add auth, retries, or instrumentation + if p.fetch != nil { + return p.fetch(req.URL.String(), req) + } + // Fallback to a vanilla http.Client + client := &http.Client{} + return client.Do(req) +} + +// RevokeToken sends a token revocation request to the configured revocation endpoint. +// If the revocation endpoint is not configured, an error is returned. +func (p *ProxyOAuthServerProvider) RevokeToken(client auth.OAuthClientInformationFull, request auth.OAuthTokenRevocationRequest) error { + // Ensure the revocation endpoint exists + if p.endpoints.RevocationURL == "" { + return fmt.Errorf("no revocation endpoint configured") + } + + // Build form-encoded body with required parameters + params := url.Values{ + "token": {request.Token}, + "client_id": {client.ClientID}, + } + + // Include client_secret when available for confidential clients + if client.ClientSecret != "" { + params.Set("client_secret", client.ClientSecret) + } + + // Optionally include token_type_hint to help the AS + if request.TokenTypeHint != "" { + params.Set("token_type_hint", request.TokenTypeHint) + } + + // Create POST request with application/x-www-form-urlencoded payload + req, err := http.NewRequest("POST", p.endpoints.RevocationURL, strings.NewReader(params.Encode())) + if err != nil { + return fmt.Errorf("create request failed: %v", err) + } + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + + // Execute request via custom fetch or default client + resp, err := p.doFetch(req) + if err != nil { + return err + } + defer resp.Body.Close() + + // Expect 200 OK per RFC 7009 + if resp.StatusCode != http.StatusOK { + return errors.NewOAuthError(errors.ErrServerError, fmt.Sprintf("Token revocation failed: %v", resp.StatusCode), "") + } + + // No body parsing required for successful revocation + return nil +} + +// ClientsStore returns an OAuthClientsStore wired for lookup and optional dynamic client registration +// depending on whether a registration endpoint is configured. +func (p *ProxyOAuthServerProvider) ClientsStore() *server.OAuthClientsStore { + var store *server.OAuthClientsStore + + // If registration URL is configured, enable dynamic client registration proxy + if p.endpoints.RegistrationURL != "" { + // Define registration function that forwards the registration request upstream + registerClient := func(client auth.OAuthClientInformationFull) (*auth.OAuthClientInformationFull, error) { + // Serialize client metadata to JSON request body + body, err := json.Marshal(client) + if err != nil { + // TODO: add logging for serialization error + return nil, fmt.Errorf("failed to marshal client: %v", err) + } + + // Create HTTP POST to upstream registration endpoint + req, err := http.NewRequest("POST", p.endpoints.RegistrationURL, bytes.NewReader(body)) + if err != nil { + // TODO: add logging for request creation error + return nil, fmt.Errorf("failed to create request: %v", err) + } + req.Header.Set("Content-Type", "application/json") + + // Execute registration request + resp, err := p.doFetch(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + // Expect 200 OK with client registration response + if resp.StatusCode != http.StatusOK { + return nil, errors.NewOAuthError(errors.ErrServerError, fmt.Errorf("client registration failed: %v", resp.StatusCode).Error(), "") + } + + // Decode response JSON into full client record + var data auth.OAuthClientInformationFull + if err := json.NewDecoder(resp.Body).Decode(&data); err != nil { + // TODO: add logging for decode error + return nil, fmt.Errorf("failed to decode response: %v", err) + } + + // Return registered client info to caller + return &data, nil + } + + // Build a store that supports both lookup and registration + store = server.NewOAuthClientStoreSupportDynamicRegistration(p.getClient, registerClient) + } else { + // Build a lookup-only store when registration is not supported + store = server.NewOAuthClientStore(p.getClient) + } + + return store +} + +// ChallengeForAuthorizationCode returns the PKCE code_challenge for a previously initiated authorization. +// In a proxy setup this is not stored locally and we defer validation to the upstream server. +func (p *ProxyOAuthServerProvider) ChallengeForAuthorizationCode(client auth.OAuthClientInformationFull, authorizationCode string) (string, error) { + // No local storage of code_challenge in proxy mode + // Upstream AS validates code_verifier against its stored challenge + return "", nil +} + +// ExchangeAuthorizationCode exchanges an authorization code for tokens by forwarding +// the request to the upstream token endpoint and returning the parsed response. +func (p *ProxyOAuthServerProvider) ExchangeAuthorizationCode(client auth.OAuthClientInformationFull, authorizationCode string, codeVerifier *string, redirectUri *string, resource *url.URL) (*auth.OAuthTokens, error) { + // Ensure a token endpoint is configured + if p.endpoints.TokenURL == "" { + return nil, fmt.Errorf("no token endpoint configured") + } + + // Build form parameters required by the authorization_code grant + params := url.Values{ + "grant_type": {"authorization_code"}, + "client_id": {client.ClientID}, + "code": {authorizationCode}, + } + + // Include client_secret for confidential clients + if client.ClientSecret != "" { + params.Set("client_secret", client.ClientSecret) + } + + // Forward PKCE code_verifier when provided + if codeVerifier != nil { + params.Set("code_verifier", *codeVerifier) + } + + // Include redirect_uri when provided to satisfy AS validation + if redirectUri != nil { + params.Set("redirect_uri", *redirectUri) + } + + // Forward resource indicator when present + if resource != nil { + params.Set("resource", resource.String()) + } + + // Create POST request to token endpoint with form-encoded body + req, err := http.NewRequest("POST", p.endpoints.TokenURL, bytes.NewReader([]byte(params.Encode()))) + if err != nil { + return nil, fmt.Errorf("failed to create request: %v", err) + } + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + + // Send request via fetch helper + resp, err := p.doFetch(req) + if err != nil { + return nil, errors.NewOAuthError(errors.ErrServerError, fmt.Sprintf("token exchange failed: %v", err), "") + } + defer resp.Body.Close() + + // Expect 200 OK for successful token response + if resp.StatusCode != http.StatusOK { + return nil, errors.NewOAuthError(errors.ErrServerError, fmt.Sprintf("token exchange failed: %v", resp.StatusCode), "") + } + + // Decode token response JSON into OAuthTokens + var data auth.OAuthTokens + if err := json.NewDecoder(resp.Body).Decode(&data); err != nil { + return nil, fmt.Errorf("failed to decode response: %v", err) + } + + // Return parsed tokens to caller + return &data, nil +} + +// ExchangeRefreshToken exchanges a refresh token for a new access token by calling +// the upstream token endpoint and validating the response payload. +func (p *ProxyOAuthServerProvider) ExchangeRefreshToken( + client auth.OAuthClientInformationFull, + refreshToken string, + scopes []string, // Optional empty slice if not provided + resource *url.URL, // Optional nil if not provided +) (*auth.OAuthTokens, error) { + // Assemble form parameters for the refresh_token grant + params := url.Values{ + "grant_type": {"refresh_token"}, + "client_id": {client.ClientID}, + "refresh_token": {refreshToken}, + } + + // Include client_secret for confidential clients + if client.ClientSecret != "" { + params.Set("client_secret", client.ClientSecret) + } + + // Optionally narrow or expand scopes as requested + if len(scopes) > 0 { + params.Set("scope", strings.Join(scopes, " ")) + } + + // Forward resource indicator when present + if resource != nil { + params.Set("resource", resource.String()) + } + + // Create a context bound POST request to the token endpoint + req, err := http.NewRequestWithContext(context.Background(), http.MethodPost, p.endpoints.TokenURL, bytes.NewBufferString(params.Encode())) + if err != nil { + return nil, fmt.Errorf("failed to create request: %v", err) + } + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + + // Choose custom fetch when provided, otherwise use default client + fetch := p.fetch + if fetch == nil { + fetch = func(url string, req *http.Request) (*http.Response, error) { + return http.DefaultClient.Do(req) + } + } + + // Execute the HTTP request + resp, err := fetch(p.endpoints.TokenURL, req) + if err != nil { + return nil, fmt.Errorf("failed to send request: %v", err) + } + defer resp.Body.Close() + + // Expect a successful status code from the token endpoint + if resp.StatusCode != http.StatusOK { + return nil, errors.NewOAuthError(errors.ErrServerError, fmt.Sprintf("token refresh failed: %v", resp.StatusCode), "") + } + + // Decode the token response body + var data auth.OAuthTokens + if err := json.NewDecoder(resp.Body).Decode(&data); err != nil { + return nil, fmt.Errorf("failed to decode response: %v", err) + } + + // Validate the decoded structure using validator/v10 + if err := validateOAuthTokens(&data); err != nil { + return nil, fmt.Errorf("validation failed: %v", err) + } + + // Return validated tokens + return &data, nil +} + +// validateOAuthTokens validates the OAuthTokens struct using github.com/go-playground/validator. +// This can be extended with custom field validators if needed. +func validateOAuthTokens(tokens *auth.OAuthTokens) error { + // Initialize a new validator instance and run struct validation + validate := validator.New() + if err := validate.Struct(tokens); err != nil { + return fmt.Errorf("validation errors: %v", err) + } + return nil +} + +// GetSkipLocalPkceValidation exposes whether local PKCE verification should be skipped. +// Token handlers can use this to decide if code_verifier must be validated locally or forwarded. +func (p *ProxyOAuthServerProvider) GetSkipLocalPkceValidation() bool { + // Return the current setting as provided during construction + return p.SkipLocalPkceValidation +} diff --git a/internal/auth/server/providers/proxy_test.go b/internal/auth/server/providers/proxy_test.go new file mode 100644 index 0000000..2d0781e --- /dev/null +++ b/internal/auth/server/providers/proxy_test.go @@ -0,0 +1,572 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package providers + +import ( + "errors" + "io" + "net/http" + "net/http/httptest" + "net/url" + "os" + "strings" + "testing" + "time" + + "trpc.group/trpc-go/trpc-mcp-go/internal/auth" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth/server" + oauthErrors "trpc.group/trpc-go/trpc-mcp-go/internal/errors" +) + +// shared test fixtures for proxy provider tests +var ( + // validClient models a typical confidential client usable across tests + validClient = auth.OAuthClientInformationFull{ + OAuthClientInformation: auth.OAuthClientInformation{ + ClientID: "test-client", + ClientSecret: "test-secret", + }, + OAuthClientMetadata: auth.OAuthClientMetadata{ + RedirectURIs: []string{"https://example.com/callback"}, + }, + } + + // baseOptions is the baseline ProxyOptions used by tests, with hooks injected in TestMain + baseOptions = ProxyOptions{ + Endpoints: ProxyEndpoints{ + AuthorizationURL: "https://auth.example.com/authorize", + TokenURL: "https://auth.example.com/token", + RevocationURL: "https://auth.example.com/revoke", + RegistrationURL: "https://auth.example.com/register", + }, + VerifyAccessToken: nil, + GetClient: nil, + Fetch: nil, + } + + // mock token payload values reused across assertions + RefreshToken = "new-refresh-token" + ExpiresIn = int64(3600) + mockTokenResponse = auth.OAuthTokens{ + AccessToken: "new-access-token", + TokenType: "Bearer", + ExpiresIn: &ExpiresIn, + RefreshToken: &RefreshToken, + } + + // mockFetch is an overridable HTTP transport used to intercept outbound requests in tests + mockFetch func(url string, req *http.Request) (*http.Response, error) +) + +// TestMain wires per-suite hooks for VerifyAccessToken, GetClient and fetch before running tests +func TestMain(m *testing.M) { + // set up VerifyAccessToken behavior + baseOptions.VerifyAccessToken = func(token string) (*server.AuthInfo, error) { + if token == "valid-token" { + ExpiresAt := time.Now().Unix() + 3600 + return &server.AuthInfo{ + Token: token, + ClientID: "test-client", + Scopes: []string{"read", "write"}, + ExpiresAt: &ExpiresAt, + }, nil + } + if token == "token-with-insufficient-scope" { + return nil, oauthErrors.NewOAuthError(oauthErrors.ErrInsufficientScope, "Required scopes: read, write", "") + } + if token == "valid-token-unexpected" { + return nil, errors.New("unexpected error") + } + return nil, oauthErrors.NewOAuthError(oauthErrors.ErrInvalidToken, "Invalid token", "") + } + + // set up client lookup behavior + baseOptions.GetClient = func(clientID string) (*auth.OAuthClientInformationFull, error) { + if clientID == "test-client" { + return &validClient, nil + } + return nil, nil + } + + // run tests + code := m.Run() + + // cleanup + mockFetch = nil + os.Exit(code) +} + +func TestProxyOAuthServerProvider(t *testing.T) { + provider := NewProxyOAuthServerProvider(baseOptions) + + // Mock codeVerifier and redirectURI + codeVerifier := "test-verifier" + redirectURI := "https://example.com/callback" + + t.Run("Authorization", func(t *testing.T) { + t.Run("Redirects to authorization endpoint with correct parameters", func(t *testing.T) { + rr := httptest.NewRecorder() + req := httptest.NewRequest("GET", "/", nil) + resource, _ := url.Parse("https://api.example.com/resource") + err := provider.Authorize(validClient, server.AuthorizationParams{ + RedirectURI: "https://example.com/callback", + CodeChallenge: "test-challenge", + State: "test-state", + Scopes: []string{"read", "write"}, + Resource: resource, + }, rr, req) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + // Verify the status code and Location header + if rr.Code != http.StatusFound { + t.Errorf("expected status code %d, got %d", http.StatusFound, rr.Code) + } + + gotURL := rr.Header().Get("Location") + // Debug output + t.Logf("got redirect URL: %s", gotURL) + expectedURL, _ := url.Parse("https://auth.example.com/authorize") + q := expectedURL.Query() + q.Set("client_id", "test-client") + q.Set("response_type", "code") + q.Set("redirect_uri", "https://example.com/callback") + q.Set("code_challenge", "test-challenge") + q.Set("code_challenge_method", "S256") + q.Set("state", "test-state") + q.Set("scope", "read write") + q.Set("resource", "https://api.example.com/resource") + expectedURL.RawQuery = q.Encode() + + if gotURL != expectedURL.String() { + t.Errorf("expected redirect URL %s, got %s", expectedURL.String(), gotURL) + } + }) + }) + + t.Run("Token Exchange", func(t *testing.T) { + t.Run("Exchanges authorization code for tokens", func(t *testing.T) { + mockFetch = func(url string, req *http.Request) (*http.Response, error) { + body, _ := io.ReadAll(req.Body) + t.Logf("request body: %s", string(body)) // 调试请求体 + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(`{"access_token":"new-access-token","token_type":"Bearer","expires_in":3600,"refresh_token":"new-refresh-token"}`)), + }, nil + } + provider.fetch = mockFetch + + tokens, err := provider.ExchangeAuthorizationCode(validClient, "test-code", &codeVerifier, nil, nil) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + // Debug output + t.Logf("tokens: %+v", tokens) + if tokens.AccessToken != mockTokenResponse.AccessToken { + t.Errorf("expected access_token %s, got %s", mockTokenResponse.AccessToken, tokens.AccessToken) + } + if tokens.TokenType != mockTokenResponse.TokenType { + t.Errorf("expected token_type %s, got %s", mockTokenResponse.TokenType, tokens.TokenType) + } + if tokens.ExpiresIn == nil || *tokens.ExpiresIn != *mockTokenResponse.ExpiresIn { + t.Errorf("expected expires_in %d, got %v", *mockTokenResponse.ExpiresIn, tokens.ExpiresIn) + } + if tokens.RefreshToken == nil || *tokens.RefreshToken != *mockTokenResponse.RefreshToken { + t.Errorf("expected refresh_token %s, got %v", *mockTokenResponse.RefreshToken, tokens.RefreshToken) + } + }) + + t.Run("Includes redirect_uri in token request when provided", func(t *testing.T) { + var calledBody string + mockFetch = func(url string, req *http.Request) (*http.Response, error) { + body, _ := io.ReadAll(req.Body) + calledBody = string(body) + // Debug output + t.Logf("request body: %s", calledBody) + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(`{"access_token":"new-access-token","token_type":"Bearer","expires_in":3600,"refresh_token":"new-refresh-token"}`)), + }, nil + } + provider.fetch = mockFetch + + _, err := provider.ExchangeAuthorizationCode(validClient, "test-code", &codeVerifier, &redirectURI, nil) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if !strings.Contains(calledBody, "redirect_uri=https%3A%2F%2Fexample.com%2Fcallback") { + t.Errorf("expected redirect_uri in body, got %s", calledBody) + } + }) + + t.Run("Handles token exchange failure", func(t *testing.T) { + mockFetch = func(url string, req *http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusBadRequest, + Body: io.NopCloser(strings.NewReader("")), + }, nil + } + provider.fetch = mockFetch + + _, err := provider.ExchangeAuthorizationCode(validClient, "test-code", &codeVerifier, nil, nil) + // Debug output + t.Logf("error: %v", err) + if err == nil { + t.Fatal("expected error, got nil") + } + var oauthErr oauthErrors.OAuthError + if !errors.As(err, &oauthErr) { + t.Errorf("expected error to be of type oauthErrors.OAuthError, got %T", err) + } + if oauthErr.ErrorCode != oauthErrors.ErrServerError.Error() { + t.Errorf("expected OAuthError with code %s, got %s", oauthErrors.ErrServerError.Error(), oauthErr.ErrorCode) + } + }) + + t.Run("Includes resource parameter in authorization code exchange", func(t *testing.T) { + var calledBody string + mockFetch = func(url string, req *http.Request) (*http.Response, error) { + if url != "https://auth.example.com/token" { + t.Errorf("expected URL %s, got %s", "https://auth.example.com/token", url) + } + if req.Method != http.MethodPost { + t.Errorf("expected method POST, got %s", req.Method) + } + if contentType := req.Header.Get("Content-Type"); contentType != "application/x-www-form-urlencoded" { + t.Errorf("expected Content-Type application/x-www-form-urlencoded, got %s", contentType) + } + body, _ := io.ReadAll(req.Body) + calledBody = string(body) + t.Logf("request body: %s", calledBody) + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(`{"access_token":"new-access-token","token_type":"Bearer","expires_in":3600,"refresh_token":"new-refresh-token"}`)), + }, nil + } + provider.fetch = mockFetch + + resource, _ := url.Parse("https://api.example.com/resource") + codeVerifier := "test-verifier" + redirectURI := "https://example.com/callback" + tokens, err := provider.ExchangeAuthorizationCode(validClient, "test-code", &codeVerifier, &redirectURI, resource) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if !strings.Contains(calledBody, "resource=https%3A%2F%2Fapi.example.com%2Fresource") { + t.Errorf("expected resource parameter in body, got %s", calledBody) + } + if tokens.AccessToken != mockTokenResponse.AccessToken { + t.Errorf("expected access_token %s, got %s", mockTokenResponse.AccessToken, tokens.AccessToken) + } + if tokens.TokenType != mockTokenResponse.TokenType { + t.Errorf("expected token_type %s, got %s", mockTokenResponse.TokenType, tokens.TokenType) + } + if tokens.ExpiresIn == nil || *tokens.ExpiresIn != *mockTokenResponse.ExpiresIn { + t.Errorf("expected expires_in %d, got %v", *mockTokenResponse.ExpiresIn, tokens.ExpiresIn) + } + if tokens.RefreshToken == nil || *tokens.RefreshToken != *mockTokenResponse.RefreshToken { + t.Errorf("expected refresh_token %s, got %v", *mockTokenResponse.RefreshToken, tokens.RefreshToken) + } + }) + + t.Run("Handles authorization code exchange without resource parameter", func(t *testing.T) { + var calledBody string + mockFetch = func(url string, req *http.Request) (*http.Response, error) { + body, _ := io.ReadAll(req.Body) + calledBody = string(body) + t.Logf("request body: %s", calledBody) // 调试输出 + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(`{"access_token":"new-access-token","token_type":"Bearer","expires_in":3600,"refresh_token":"new-refresh-token"}`)), + }, nil + } + provider.fetch = mockFetch + + tokens, err := provider.ExchangeAuthorizationCode(validClient, "test-code", &codeVerifier, nil, nil) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if strings.Contains(calledBody, "resource=") { + t.Errorf("expected no resource parameter in body, got %s", calledBody) + } + if tokens.AccessToken != mockTokenResponse.AccessToken { + t.Errorf("expected access_token %s, got %s", mockTokenResponse.AccessToken, tokens.AccessToken) + } + if tokens.TokenType != mockTokenResponse.TokenType { + t.Errorf("expected token_type %s, got %s", mockTokenResponse.TokenType, tokens.TokenType) + } + if tokens.ExpiresIn == nil || *tokens.ExpiresIn != *mockTokenResponse.ExpiresIn { + t.Errorf("expected expires_in %d, got %v", *mockTokenResponse.ExpiresIn, tokens.ExpiresIn) + } + if tokens.RefreshToken == nil || *tokens.RefreshToken != *mockTokenResponse.RefreshToken { + t.Errorf("expected refresh_token %s, got %v", *mockTokenResponse.RefreshToken, tokens.RefreshToken) + } + }) + + t.Run("Includes resource parameter in refresh token exchange", func(t *testing.T) { + var calledBody string + mockFetch = func(url string, req *http.Request) (*http.Response, error) { + body, _ := io.ReadAll(req.Body) + calledBody = string(body) + // Debug output + t.Logf("request body: %s", calledBody) + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(`{"access_token":"new-access-token","token_type":"Bearer","expires_in":3600,"refresh_token":"new-refresh-token"}`)), + }, nil + } + provider.fetch = mockFetch + + resource, _ := url.Parse("https://api.example.com/resource") + tokens, err := provider.ExchangeRefreshToken(validClient, "test-refresh-token", []string{"read", "write"}, resource) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if !strings.Contains(calledBody, "resource=https%3A%2F%2Fapi.example.com%2Fresource") { + t.Errorf("expected resource in body, got %s", calledBody) + } + if tokens.AccessToken != mockTokenResponse.AccessToken { + t.Errorf("expected access_token %s, got %s", mockTokenResponse.AccessToken, tokens.AccessToken) + } + if tokens.TokenType != mockTokenResponse.TokenType { + t.Errorf("expected token_type %s, got %s", mockTokenResponse.TokenType, tokens.TokenType) + } + if tokens.ExpiresIn == nil || *tokens.ExpiresIn != *mockTokenResponse.ExpiresIn { + t.Errorf("expected expires_in %d, got %v", *mockTokenResponse.ExpiresIn, tokens.ExpiresIn) + } + if tokens.RefreshToken == nil || *tokens.RefreshToken != *mockTokenResponse.RefreshToken { + t.Errorf("expected refresh_token %s, got %v", *mockTokenResponse.RefreshToken, tokens.RefreshToken) + } + }) + }) + + t.Run("Client Registration", func(t *testing.T) { + t.Run("Registers new client", func(t *testing.T) { + mockFetch = func(url string, req *http.Request) (*http.Response, error) { + body, _ := io.ReadAll(req.Body) + // Debug output + t.Logf("register request body: %s", string(body)) + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(`{"client_id":"new-client","client_secret":"new-secret","redirect_uris":["https://new-client.com/callback"]}`)), + }, nil + } + provider.fetch = mockFetch + + newClient := auth.OAuthClientInformationFull{ + OAuthClientInformation: auth.OAuthClientInformation{ + ClientID: "new-client", + }, + OAuthClientMetadata: auth.OAuthClientMetadata{ + RedirectURIs: []string{"https://new-client.com/callback"}, + }, + } + result, err := provider.ClientsStore().RegisterClient(newClient) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if result.ClientID != newClient.ClientID { + t.Errorf("expected client_id %s, got %s", newClient.ClientID, result.ClientID) + } + }) + + t.Run("Handles registration failure", func(t *testing.T) { + mockFetch = func(url string, req *http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusBadRequest, + Body: io.NopCloser(strings.NewReader("")), + }, nil + } + provider.fetch = mockFetch + + newClient := auth.OAuthClientInformationFull{ + OAuthClientInformation: auth.OAuthClientInformation{ + ClientID: "new-client", + }, + OAuthClientMetadata: auth.OAuthClientMetadata{ + RedirectURIs: []string{"https://new-client.com/callback"}, + }, + } + _, err := provider.ClientsStore().RegisterClient(newClient) + t.Logf("error: %v", err) // Debug output + if err == nil { + t.Fatal("expected error, got nil") + } + var serverErr oauthErrors.OAuthError + if !errors.As(err, &serverErr) { + t.Errorf("expected error to be of type oauthErrors.OAuthError, got %T", err) + } + if serverErr.ErrorCode != oauthErrors.ErrServerError.Error() { + t.Errorf("expected OAuthError with code %s, got %s", oauthErrors.ErrServerError.Error(), serverErr.ErrorCode) + } + }) + }) + + t.Run("Token Revocation", func(t *testing.T) { + t.Run("Revokes token", func(t *testing.T) { + mockFetch = func(url string, req *http.Request) (*http.Response, error) { + body, _ := io.ReadAll(req.Body) + t.Logf("revoke request body: %s", string(body)) // 调试输出 + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader("")), + }, nil + } + provider.fetch = mockFetch + + err := provider.RevokeToken(validClient, auth.OAuthTokenRevocationRequest{ + Token: "token-to-revoke", + TokenTypeHint: "access_token", + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + }) + + t.Run("Handles revocation failure", func(t *testing.T) { + mockFetch = func(url string, req *http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusBadRequest, + Body: io.NopCloser(strings.NewReader("")), + }, nil + } + provider.fetch = mockFetch + + err := provider.RevokeToken(validClient, auth.OAuthTokenRevocationRequest{ + Token: "invalid-token", + }) + t.Logf("error: %v", err) // 调试输出 + if err == nil { + t.Fatal("expected error, got nil") + } + var serverErr oauthErrors.OAuthError + if !errors.As(err, &serverErr) { + t.Errorf("expected error to be of type oauthErrors.OAuthError, got %T", err) + } + if serverErr.ErrorCode != oauthErrors.ErrServerError.Error() { + t.Errorf("expected OAuthError with code %s, got %s", oauthErrors.ErrServerError.Error(), serverErr.ErrorCode) + } + }) + }) + + t.Run("Token Verification", func(t *testing.T) { + t.Run("Verifies valid token", func(t *testing.T) { + var calledToken string + options := baseOptions + options.VerifyAccessToken = func(token string) (*server.AuthInfo, error) { + calledToken = token + return baseOptions.VerifyAccessToken(token) + } + provider := NewProxyOAuthServerProvider(options) + + authInfo, err := provider.VerifyAccessToken("valid-token") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if calledToken != "valid-token" { + t.Errorf("expected VerifyAccessToken called with 'valid-token', got %q", calledToken) + } + t.Logf("authInfo: %+v", authInfo) // 调试输出 + if authInfo.ClientID != "test-client" { + t.Errorf("expected clientId test-client, got %s", authInfo.ClientID) + } + }) + + t.Run("Passes through InvalidTokenError", func(t *testing.T) { + var calledToken string + options := baseOptions + options.VerifyAccessToken = func(token string) (*server.AuthInfo, error) { + calledToken = token // 记录调用参数 + if token == "valid-token" { + ExpiresAt := time.Now().Unix() + 3600 + return &server.AuthInfo{ + Token: token, + ClientID: "test-client", + Scopes: []string{"read", "write"}, + ExpiresAt: &ExpiresAt, + }, nil + } + return nil, oauthErrors.NewOAuthError(oauthErrors.ErrInvalidToken, "Invalid token", "") + } + provider := NewProxyOAuthServerProvider(options) + + _, err := provider.VerifyAccessToken("invalid-token") + t.Logf("error: %v", err) + if err == nil { + t.Fatal("expected error, got nil") + } + if calledToken != "invalid-token" { + t.Errorf("expected VerifyAccessToken called with 'invalid-token', got %q", calledToken) + } + var invalidTokenErr oauthErrors.OAuthError + if !errors.As(err, &invalidTokenErr) { + t.Errorf("expected error to be of type oauthErrors.OAuthError, got %T", err) + } + if invalidTokenErr.ErrorCode != oauthErrors.ErrInvalidToken.Error() { + t.Errorf("expected OAuthError with code %s, got %s", oauthErrors.ErrInvalidToken.Error(), invalidTokenErr.ErrorCode) + } + }) + + t.Run("Passes through InsufficientScopeError", func(t *testing.T) { + var calledToken string + options := baseOptions // 复制全局配置 + options.VerifyAccessToken = func(token string) (*server.AuthInfo, error) { + calledToken = token + return nil, oauthErrors.NewOAuthError(oauthErrors.ErrInsufficientScope, "Required scopes: read, write", "") + } + provider := NewProxyOAuthServerProvider(options) + + _, err := provider.VerifyAccessToken("token-with-insufficient-scope") + t.Logf("error: %v", err) + if err == nil { + t.Fatal("expected error, got nil") + } + var insufficientScopeErr oauthErrors.OAuthError + if !errors.As(err, &insufficientScopeErr) { + t.Errorf("expected error to be of type oauthErrors.OAuthError, got %T", err) + } + if insufficientScopeErr.ErrorCode != "insufficient_scope" { + t.Errorf("expected OAuthError with code %s, got %s", "insufficient_scope", insufficientScopeErr.ErrorCode) + } + if insufficientScopeErr.Message != "Required scopes: read, write" { + t.Errorf("expected OAuthError with description %s, got %s", "Required scopes: read, write", insufficientScopeErr.Message) + } + if calledToken != "token-with-insufficient-scope" { + t.Errorf("expected VerifyAccessToken called with %s, got %s", "token-with-insufficient-scope", calledToken) + } + }) + + t.Run("Passes through unexpected errors", func(t *testing.T) { + var calledToken string + // Copy global configuration + options := baseOptions + options.VerifyAccessToken = func(token string) (*server.AuthInfo, error) { + calledToken = token + return baseOptions.VerifyAccessToken(token) + } + provider := NewProxyOAuthServerProvider(options) + + // Invoke VerifyAccessToken + _, err := provider.VerifyAccessToken("valid-token-unexpected") + t.Logf("error: %v", err) // 调试输出 + if err == nil { + t.Fatal("expected error, got nil") + } + var oauthErr oauthErrors.OAuthError + if errors.As(err, &oauthErr) { + t.Errorf("expected error to be non-OAuth error, got oauthErrors.OAuthError") + } + if err.Error() != "unexpected error" { + t.Errorf("expected error message %s, got %s", "unexpected error", err.Error()) + } + if calledToken != "valid-token-unexpected" { + t.Errorf("expected VerifyAccessToken called with %s, got %s", "valid-token-unexpected", calledToken) + } + }) + }) +} diff --git a/internal/auth/server/router/router.go b/internal/auth/server/router/router.go new file mode 100644 index 0000000..6e0ac00 --- /dev/null +++ b/internal/auth/server/router/router.go @@ -0,0 +1,391 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package router + +import ( + "fmt" + "net/http" + "net/url" + + "trpc.group/trpc-go/trpc-mcp-go/internal/auth" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth/server" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth/server/handler" +) + +// AuthRouterOptions holds configuration options for the MCP authentication router. +// It configures how OAuth 2.1 endpoints (/authorize, /token, /revoke, /register) are exposed. +type AuthRouterOptions struct { + // Provider is the OAuth server implementation. + // It manages client registration, authorization codes, tokens, and verification. + Provider server.OAuthServerProvider + + // IssuerUrl is the OAuth issuer identifier (RFC 8414). + // Typically something like "https://auth.example.com". + IssuerUrl *url.URL + + // BaseUrl is the base URL of this service, used to construct endpoint URLs + // such as /authorize, /token, etc. + BaseUrl *url.URL + + // ServiceDocumentationUrl points to human-readable documentation about the service, + // usually an API docs page. + ServiceDocumentationUrl *url.URL + + // ScopesSupported lists all scopes supported by this authorization server, + // for example: ["read", "write"]. + ScopesSupported []string + + // ResourceName is an optional logical name for the protected resource/API. + ResourceName *string + + // AuthorizationOptions configures the /authorize endpoint (validation, rate limiting, etc.). + AuthorizationOptions *handler.AuthorizationHandlerOptions + + // ClientRegistrationOptions configures the /register endpoint for dynamic client registration (RFC 7591). + ClientRegistrationOptions *handler.ClientRegistrationHandlerOptions + + // RevocationOptions configures the /revoke endpoint for token revocation (RFC 7009). + RevocationOptions *handler.RevocationHandlerOptions + + // TokenOptions configures the /token endpoint for issuing tokens (supports auth code flow, PKCE, etc.). + TokenOptions *handler.TokenHandlerOptions +} + +// AuthMetadataOptions holds configuration options for the MCP authentication metadata endpoints. +// It controls what is published via OAuth 2.1 Authorization Server Metadata (RFC 8414). +type AuthMetadataOptions struct { + // OAuthMetadata contains the full OAuth 2.1 Authorization Server Metadata, + // including authorization_endpoint, token_endpoint, scopes_supported, etc. + OAuthMetadata auth.OAuthMetadata + + // ResourceServerUrl points to the protected resource server, + // used by clients to discover where to send API requests. + ResourceServerUrl *url.URL + + // ServiceDocumentationUrl points to human-readable documentation about the service. + ServiceDocumentationUrl *url.URL + + // ScopesSupported lists the scopes supported by this resource server. + ScopesSupported []string + + // ResourceName is an optional logical name for the resource server. + ResourceName *string +} + +// checkIssuerUrl validates the issuer URL according to RFC 8414. +func checkIssuerUrl(issuer *url.URL) error { + // Technically RFC 8414 does not permit a localhost HTTPS exemption, + // but this will be necessary for ease of testing + if issuer.Scheme != "https" && issuer.Hostname() != "localhost" && issuer.Hostname() != "127.0.0.1" { + return fmt.Errorf("issuer URL must be HTTPS") + } + if issuer.Fragment != "" { + return fmt.Errorf("issuer URL must not have a fragment: %s", issuer.String()) + } + if issuer.RawQuery != "" { + return fmt.Errorf("issuer URL must not have a query string: %s", issuer.String()) + } + return nil +} + +// supportsClientRegistration checks if the provider supports dynamic client registration +func supportsClientRegistration(provider server.OAuthServerProvider) bool { + if provider == nil { + return false + } + + clientsStore := provider.ClientsStore() + if clientsStore == nil { + return false + } + // Check if the clients store supports registration + return clientsStore.SupportsRegistration() +} + +// supportsTokenRevocation checks if the provider supports token revocation +func supportsTokenRevocation(provider server.OAuthServerProvider) bool { + if provider == nil { + return false + } + // Use type assertion to check if the provider implements SupportTokenRevocation interface + _, ok := provider.(server.SupportTokenRevocation) + return ok +} + +// CreateOAuthMetadata generates OAuth 2.1 compliant Authorization Server Metadata. +func CreateOAuthMetadata(options struct { + Provider server.OAuthServerProvider + IssuerUrl *url.URL + BaseUrl *url.URL + ServiceDocumentationUrl *url.URL + ScopesSupported []string +}) (auth.OAuthMetadata, error) { + if options.Provider == nil { + return auth.OAuthMetadata{}, fmt.Errorf("provider is required") + } + + issuer := options.IssuerUrl + baseUrl := options.BaseUrl + + // Validate issuer URL + if err := checkIssuerUrl(issuer); err != nil { + return auth.OAuthMetadata{}, err + } + + // Determine base URL for endpoints + var baseUrlForEndpoints *url.URL + if baseUrl != nil { + baseUrlForEndpoints = baseUrl + } else { + baseUrlForEndpoints = issuer + } + + // Required endpoints + authorizationEndpoint := "/authorize" + tokenEndpoint := "/token" + + authEndpointUrl, _ := url.Parse(authorizationEndpoint) + tokenEndpointUrl, _ := url.Parse(tokenEndpoint) + + metadata := auth.OAuthMetadata{ + // Core fields + Issuer: issuer.String(), + AuthorizationEndpoint: baseUrlForEndpoints.ResolveReference(authEndpointUrl).String(), + TokenEndpoint: baseUrlForEndpoints.ResolveReference(tokenEndpointUrl).String(), + + // OAuth 2.1 requires PKCE support + ResponseTypesSupported: []string{"code"}, // OAuth 2.1 removes implicit flow + CodeChallengeMethodsSupported: []string{"S256"}, // OAuth 2.1 requires S256, plain is deprecated + + // Token endpoint authentication methods + TokenEndpointAuthMethodsSupported: []string{"client_secret_post", "client_secret_basic"}, + + // OAuth 2.1 supported grant types + GrantTypesSupported: []string{"authorization_code", "refresh_token"}, + + // Optional fields + ScopesSupported: options.ScopesSupported, + } + + // Add service documentation if provided + if options.ServiceDocumentationUrl != nil { + serviceDoc := options.ServiceDocumentationUrl.String() + metadata.ServiceDocumentation = &serviceDoc + } + + // Check for optional endpoints based on provider capabilities + if supportsTokenRevocation(options.Provider) { + revocationEndpoint := "/revoke" + revEndpointUrl, _ := url.Parse(revocationEndpoint) + revEndpoint := baseUrlForEndpoints.ResolveReference(revEndpointUrl).String() + metadata.RevocationEndpoint = &revEndpoint + metadata.RevocationEndpointAuthMethodsSupported = []string{"client_secret_post", "client_secret_basic"} + } + + if supportsClientRegistration(options.Provider) { + registrationEndpoint := "/register" + regEndpointUrl, _ := url.Parse(registrationEndpoint) + regEndpoint := baseUrlForEndpoints.ResolveReference(regEndpointUrl).String() + metadata.RegistrationEndpoint = ®Endpoint + } + + return metadata, nil +} + +// McpAuthRouter sets up OAuth 2.1 compliant MCP authorization server endpoints +func McpAuthRouter(mux *http.ServeMux, options AuthRouterOptions) error { + // Create OAuth metadata with error handling + oauthMetadata, err := CreateOAuthMetadata(struct { + Provider server.OAuthServerProvider + IssuerUrl *url.URL + BaseUrl *url.URL + ServiceDocumentationUrl *url.URL + ScopesSupported []string + }{ + Provider: options.Provider, + IssuerUrl: options.IssuerUrl, + BaseUrl: options.BaseUrl, + ServiceDocumentationUrl: options.ServiceDocumentationUrl, + ScopesSupported: options.ScopesSupported, + }) + if err != nil { + return fmt.Errorf("failed to create OAuth metadata: %w", err) + } + + // Authorization endpoint (GET only for OAuth 2.1) + authorizationURL, _ := url.Parse(oauthMetadata.AuthorizationEndpoint) + authzOptions := handler.AuthorizationHandlerOptions{ + Provider: options.Provider, + } + if options.AuthorizationOptions != nil && options.AuthorizationOptions.RateLimit != nil { + authzOptions.RateLimit = options.AuthorizationOptions.RateLimit + } + mux.Handle(authorizationURL.Path, methodRestrictedHandler("GET", handler.AuthorizationHandler(authzOptions))) + + // Token endpoint (POST only for OAuth 2.1) + tokenURL, _ := url.Parse(oauthMetadata.TokenEndpoint) + tokenOptions := handler.TokenHandlerOptions{Provider: options.Provider} + if options.TokenOptions != nil { + if options.TokenOptions.RateLimit != nil { + tokenOptions.RateLimit = options.TokenOptions.RateLimit + } + } + mux.Handle(tokenURL.Path, methodRestrictedHandler("POST", handler.TokenHandler(tokenOptions))) + + // Metadata endpoints + issuerURL, _ := url.Parse(oauthMetadata.Issuer) + resourceURL := options.BaseUrl + if resourceURL == nil { + resourceURL = issuerURL + } + if err := McpAuthMetadataRouter(mux, AuthMetadataOptions{ + OAuthMetadata: oauthMetadata, + ResourceServerUrl: resourceURL, + ServiceDocumentationUrl: options.ServiceDocumentationUrl, + ScopesSupported: options.ScopesSupported, + ResourceName: options.ResourceName, + }); err != nil { + return fmt.Errorf("failed to setup metadata router: %w", err) + } + + // Dynamic client registration (optional, POST only) + if oauthMetadata.RegistrationEndpoint != nil { + // Ensure ClientsStore() is not nil before mounting /register + if clientsStore := options.Provider.ClientsStore(); clientsStore != nil { + registrationURL, _ := url.Parse(*oauthMetadata.RegistrationEndpoint) + regOpts := handler.ClientRegistrationHandlerOptions{ + ClientsStore: clientsStore, + } + if options.ClientRegistrationOptions != nil { + regOpts = *options.ClientRegistrationOptions + regOpts.ClientsStore = clientsStore + } else { + // OAuth 2.1 recommended rate limiting for client registration + regOpts.RateLimit = &handler.RegisterRateLimitConfig{ + WindowMs: 60000, + Max: 10, + } + } + mux.Handle(registrationURL.Path, methodRestrictedHandler("POST", handler.ClientRegistrationHandler(regOpts))) + } + } + + // Token revocation endpoint (optional, POST only) + if oauthMetadata.RevocationEndpoint != nil { + revocationURL, _ := url.Parse(*oauthMetadata.RevocationEndpoint) + + revOpts := handler.RevocationHandlerOptions{ + Provider: options.Provider, + } + if options.RevocationOptions != nil && options.RevocationOptions.RateLimit != nil { + revOpts.RateLimit = options.RevocationOptions.RateLimit + } + + mux.Handle(revocationURL.Path, methodRestrictedHandler("POST", handler.RevocationHandler(revOpts))) + } + + return nil +} + +// McpAuthMetadataRouter sets up OAuth 2.1 compliant metadata endpoints +func McpAuthMetadataRouter(mux *http.ServeMux, options AuthMetadataOptions) error { + issuerURL, _ := url.Parse(options.OAuthMetadata.Issuer) + if err := checkIssuerUrl(issuerURL); err != nil { + return fmt.Errorf("invalid issuer URL in metadata: %w", err) + } + + // Create protected resource metadata + protectedResourceMetadata := auth.OAuthProtectedResourceMetadata{ + Resource: options.ResourceServerUrl.String(), + AuthorizationServers: []string{ + options.OAuthMetadata.Issuer, + }, + ScopesSupported: options.ScopesSupported, + } + + // Add optional fields + if options.ResourceName != nil { + protectedResourceMetadata.ResourceName = options.ResourceName + } + + if options.ServiceDocumentationUrl != nil { + resourceDoc := options.ServiceDocumentationUrl.String() + protectedResourceMetadata.ResourceDocumentation = &resourceDoc + } + + // Protected resource metadata endpoint (GET only) + mux.Handle("/.well-known/oauth-protected-resource", + methodRestrictedHandler("GET", handler.MetadataHandler(protectedResourceMetadata))) + + // Authorization server metadata endpoint (GET only, for backward compatibility) + mux.Handle("/.well-known/oauth-authorization-server", + methodRestrictedHandler("GET", handler.MetadataHandler(options.OAuthMetadata))) + + return nil +} + +// GetOAuthProtectedResourceMetadataUrl constructs the OAuth 2.0 Protected Resource Metadata URL from a given server URL +func GetOAuthProtectedResourceMetadataUrl(serverUrl *url.URL) string { + metadataUrl, _ := url.Parse("/.well-known/oauth-protected-resource") + return serverUrl.ResolveReference(metadataUrl).String() +} + +// InstallMCPAuthRoutes convenience function to simplify OAuth 2.1 compliant route installation +func InstallMCPAuthRoutes( + mux *http.ServeMux, + issuerBaseURL string, + resourceServerURL string, + provider server.OAuthServerProvider, + scopesSupported []string, + resourceName *string, + serviceDocURL *string, +) error { + issuerURL, err := url.Parse(issuerBaseURL) + if err != nil { + return fmt.Errorf("invalid issuer URL: %w", err) + } + + var baseURL *url.URL + if resourceServerURL != "" { + baseURL, err = url.Parse(resourceServerURL) + if err != nil { + return fmt.Errorf("invalid resource server URL: %w", err) + } + } + + var serviceDocumentationUrl *url.URL + if serviceDocURL != nil { + serviceDocumentationUrl, err = url.Parse(*serviceDocURL) + if err != nil { + return fmt.Errorf("invalid service documentation URL: %w", err) + } + } + + options := AuthRouterOptions{ + Provider: provider, + IssuerUrl: issuerURL, + BaseUrl: baseURL, + ServiceDocumentationUrl: serviceDocumentationUrl, + ScopesSupported: scopesSupported, + ResourceName: resourceName, + } + + return McpAuthRouter(mux, options) +} + +// methodRestrictedHandler returns an HTTP handler that restricts requests +// to the specified HTTP method. If the request method does not match +func methodRestrictedHandler(allowedMethod string, h http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != allowedMethod { + w.Header().Set("Allow", allowedMethod) + http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) + return + } + h.ServeHTTP(w, r) + }) +} diff --git a/internal/auth/server/router/router_test.go b/internal/auth/server/router/router_test.go new file mode 100644 index 0000000..3a17147 --- /dev/null +++ b/internal/auth/server/router/router_test.go @@ -0,0 +1,399 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package router + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "testing" + "time" + + "trpc.group/trpc-go/trpc-mcp-go/internal/auth" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth/server" +) + +// fullProvider implements a fuller OAuthServerProvider used to exercise all router endpoints in tests +type fullProvider struct{} + +// ClientsStore returns a store that supports dynamic registration for tests +func (p *fullProvider) ClientsStore() *server.OAuthClientsStore { + return server.NewOAuthClientStoreSupportDynamicRegistration( + func(clientID string) (*auth.OAuthClientInformationFull, error) { + return &auth.OAuthClientInformationFull{ + OAuthClientInformation: auth.OAuthClientInformation{ + ClientID: clientID, + }, + }, nil + }, + func(client auth.OAuthClientInformationFull) (*auth.OAuthClientInformationFull, error) { + return &client, nil + }, + ) +} + +// Authorize simulates a successful authorization response by redirecting with a mock code +func (p *fullProvider) Authorize(client auth.OAuthClientInformationFull, params server.AuthorizationParams, w http.ResponseWriter, r *http.Request) error { + u, _ := url.Parse(params.RedirectURI) + q := u.Query() + q.Set("code", "mock_auth_code") + if params.State != "" { + q.Set("state", params.State) + } + u.RawQuery = q.Encode() + http.Redirect(w, r, u.String(), http.StatusFound) + return nil +} + +// ChallengeForAuthorizationCode returns a fixed PKCE challenge for testing +func (p *fullProvider) ChallengeForAuthorizationCode(client auth.OAuthClientInformationFull, code string) (string, error) { + return "mock_challenge", nil +} + +// ExchangeAuthorizationCode returns mock tokens for a valid authorization code exchange +func (p *fullProvider) ExchangeAuthorizationCode(client auth.OAuthClientInformationFull, code string, verifier *string, redirect *string, resource *url.URL) (*auth.OAuthTokens, error) { + expires := int64(3600) + rt := "mock_refresh_token" + return &auth.OAuthTokens{ + AccessToken: "mock_access_token", + TokenType: "bearer", + ExpiresIn: &expires, + RefreshToken: &rt, + }, nil +} + +// ExchangeRefreshToken returns mock tokens for a valid refresh token exchange +func (p *fullProvider) ExchangeRefreshToken(client auth.OAuthClientInformationFull, rt string, scopes []string, resource *url.URL) (*auth.OAuthTokens, error) { + expires := int64(3600) + newRT := "new_mock_refresh_token" + return &auth.OAuthTokens{ + AccessToken: "new_mock_access_token", + TokenType: "bearer", + ExpiresIn: &expires, + RefreshToken: &newRT, + }, nil +} + +// VerifyAccessToken validates a token and returns mock auth info or an error +func (p *fullProvider) VerifyAccessToken(token string) (*server.AuthInfo, error) { + if token == "valid_token" { + exp := time.Now().Add(time.Hour).Unix() + return &server.AuthInfo{ + Token: token, + ClientID: "valid-client", + Scopes: []string{"read", "write"}, + ExpiresAt: &exp, + }, nil + } + return nil, ErrInvalid +} + +// RevokeToken implements token revocation to satisfy the interface +func (p *fullProvider) RevokeToken(client auth.OAuthClientInformationFull, req auth.OAuthTokenRevocationRequest) error { + return nil +} + +// minimalProvider implements the minimal surface needed for router tests without dynamic registration +type minimalProvider struct{} + +// ClientsStore returns nil to indicate no dynamic client registration in minimal mode +func (p *minimalProvider) ClientsStore() *server.OAuthClientsStore { return nil } + +// Authorize simulates a basic authorization redirect with a mock code +func (p *minimalProvider) Authorize(client auth.OAuthClientInformationFull, params server.AuthorizationParams, w http.ResponseWriter, r *http.Request) error { + u, _ := url.Parse(params.RedirectURI) + q := u.Query() + q.Set("code", "mock_auth_code") + u.RawQuery = q.Encode() + http.Redirect(w, r, u.String(), http.StatusFound) + return nil +} + +// ChallengeForAuthorizationCode returns a fixed PKCE challenge in minimal mode +func (p *minimalProvider) ChallengeForAuthorizationCode(client auth.OAuthClientInformationFull, code string) (string, error) { + return "mock_challenge", nil +} + +// ExchangeAuthorizationCode returns a basic mock access token for code exchange +func (p *minimalProvider) ExchangeAuthorizationCode(client auth.OAuthClientInformationFull, code string, verifier *string, redirect *string, resource *url.URL) (*auth.OAuthTokens, error) { + expires := int64(3600) + return &auth.OAuthTokens{AccessToken: "mock_access_token", TokenType: "bearer", ExpiresIn: &expires}, nil +} + +// ExchangeRefreshToken returns a basic mock access token for refresh exchange +func (p *minimalProvider) ExchangeRefreshToken(client auth.OAuthClientInformationFull, rt string, scopes []string, resource *url.URL) (*auth.OAuthTokens, error) { + expires := int64(3600) + return &auth.OAuthTokens{AccessToken: "new_mock_access_token", TokenType: "bearer", ExpiresIn: &expires}, nil +} + +// VerifyAccessToken always returns a valid mock auth info in minimal mode +func (p *minimalProvider) VerifyAccessToken(token string) (*server.AuthInfo, error) { + exp := time.Now().Add(time.Hour).Unix() + return &server.AuthInfo{Token: token, ClientID: "valid-client", Scopes: []string{"read"}, ExpiresAt: &exp}, nil +} + +// RevokeToken implements a no-op revocation to satisfy the embedded interface +func (p *minimalProvider) RevokeToken(client auth.OAuthClientInformationFull, req auth.OAuthTokenRevocationRequest) error { + return nil +} + +func Test_McpAuthRouter_RouterCreation_Validation(t *testing.T) { + mux := http.NewServeMux() + issuerHTTP, _ := url.Parse("http://auth.example.com") + err := McpAuthRouter(mux, AuthRouterOptions{ + Provider: &fullProvider{}, + IssuerUrl: issuerHTTP, + }) + if err == nil { + t.Fatalf("expected error for non-HTTPS issuer") + } + + muxOK := http.NewServeMux() + issuerHTTPS, _ := url.Parse("https://auth.example.com") + if err := McpAuthRouter(muxOK, AuthRouterOptions{ + Provider: &fullProvider{}, + IssuerUrl: issuerHTTPS, + }); err != nil { + t.Fatalf("unexpected error for valid https issuer: %v", err) + } +} + +func Test_Metadata_AuthorizationServer_Full(t *testing.T) { + mux := http.NewServeMux() + + issuer, _ := url.Parse("https://auth.example.com/") + if err := McpAuthRouter(mux, AuthRouterOptions{ + Provider: &fullProvider{}, + IssuerUrl: issuer, + ServiceDocumentationUrl: mustParseURL("https://docs.example.com"), + ScopesSupported: []string{"read", "write"}, + }); err != nil { + t.Fatalf("router init failed: %v", err) + } + + ts := httptest.NewServer(mux) + defer ts.Close() + + res, err := http.Get(ts.URL + "/.well-known/oauth-authorization-server") + if err != nil { + t.Fatalf("GET metadata failed: %v", err) + } + defer res.Body.Close() + if res.StatusCode != 200 { + t.Fatalf("expected 200, got %d", res.StatusCode) + } + + var body map[string]any + _ = json.NewDecoder(res.Body).Decode(&body) + + expectStr(t, body, "issuer", "https://auth.example.com/") + expectStr(t, body, "authorization_endpoint", "https://auth.example.com/authorize") + expectStr(t, body, "token_endpoint", "https://auth.example.com/token") + + expectArr(t, body, "response_types_supported", []string{"code"}) + expectArr(t, body, "grant_types_supported", []string{"authorization_code", "refresh_token"}) + expectArr(t, body, "code_challenge_methods_supported", []string{"S256"}) + expectArr(t, body, "token_endpoint_auth_methods_supported", containsAny("client_secret_post", "client_secret_basic")) + + // fullProvider: typically has revoke/register endpoints (depends on router implementation) + // Not asserting presence to avoid coupling with specific implementation; add expectStr(...) if needed +} + +func Test_Metadata_ProtectedResource_Full(t *testing.T) { + mux := http.NewServeMux() + + issuer, _ := url.Parse("https://auth.example.com/") + if err := McpAuthRouter(mux, AuthRouterOptions{ + Provider: &fullProvider{}, + IssuerUrl: issuer, + ServiceDocumentationUrl: mustParseURL("https://docs.example.com/"), + ScopesSupported: []string{"read", "write"}, + ResourceName: strPtr("Test API"), + }); err != nil { + t.Fatalf("router init failed: %v", err) + } + + ts := httptest.NewServer(mux) + defer ts.Close() + + res, err := http.Get(ts.URL + "/.well-known/oauth-protected-resource") + if err != nil { + t.Fatalf("GET resource metadata failed: %v", err) + } + defer res.Body.Close() + if res.StatusCode != 200 { + t.Fatalf("expected 200, got %d", res.StatusCode) + } + + var body map[string]any + _ = json.NewDecoder(res.Body).Decode(&body) + + expectStr(t, body, "resource", "https://auth.example.com/") // depends on router implementation + expectArr(t, body, "authorization_servers", []string{"https://auth.example.com/"}) + expectArr(t, body, "scopes_supported", []string{"read", "write"}) + expectStr(t, body, "resource_name", "Test API") + expectStr(t, body, "resource_documentation", "https://docs.example.com/") +} + +func Test_Metadata_Minimal_NoOptionalFields(t *testing.T) { + mux := http.NewServeMux() + + issuer, _ := url.Parse("https://auth.example.com") + if err := McpAuthRouter(mux, AuthRouterOptions{ + Provider: &minimalProvider{}, + IssuerUrl: issuer, + }); err != nil { + t.Fatalf("router init failed: %v", err) + } + + ts := httptest.NewServer(mux) + defer ts.Close() + + // Authorization server metadata: optional fields omitted + as, _ := http.Get(ts.URL + "/.well-known/oauth-authorization-server") + defer as.Body.Close() + var a map[string]any + _ = json.NewDecoder(as.Body).Decode(&a) + if _, ok := a["service_documentation"]; ok { + t.Fatalf("service_documentation should be omitted") + } + if _, ok := a["scopes_supported"]; ok { + t.Fatalf("scopes_supported should be omitted") + } + + // Protected resource metadata: optional fields omitted + pr, _ := http.Get(ts.URL + "/.well-known/oauth-protected-resource") + defer pr.Body.Close() + var p map[string]any + _ = json.NewDecoder(pr.Body).Decode(&p) + if _, ok := p["scopes_supported"]; ok { + t.Fatalf("scopes_supported should be omitted") + } + if _, ok := p["resource_name"]; ok { + t.Fatalf("resource_name should be omitted") + } + if _, ok := p["resource_documentation"]; ok { + t.Fatalf("resource_documentation should be omitted") + } +} + +func Test_Routes_Register_And_Revoke_Presence_MinimalVsFull(t *testing.T) { + issuer, _ := url.Parse("https://auth.example.com") + + // full provider: registers /register endpoint + muxFull := http.NewServeMux() + if err := McpAuthRouter(muxFull, AuthRouterOptions{ + Provider: &fullProvider{}, + IssuerUrl: issuer, + }); err != nil { + t.Fatalf("init full failed: %v", err) + } + tsFull := httptest.NewServer(muxFull) + defer tsFull.Close() + + // Test /register exists + r1, _ := http.PostForm(tsFull.URL+"/register", url.Values{ + "redirect_uris": {"https://example.com/callback"}, + }) + _ = r1.Body.Close() + if r1.StatusCode == http.StatusNotFound { + t.Fatalf("full: /register should exist (got 404)") + } + + // minimal provider: should return 404 + muxMin := http.NewServeMux() + if err := McpAuthRouter(muxMin, AuthRouterOptions{ + Provider: &minimalProvider{}, + IssuerUrl: issuer, + }); err != nil { + t.Fatalf("init minimal failed: %v", err) + } + tsMin := httptest.NewServer(muxMin) + defer tsMin.Close() + + // Test /register does not exist + mr, _ := http.PostForm(tsMin.URL+"/register", url.Values{ + "redirect_uris": {"https://example.com/callback"}, + }) + _ = mr.Body.Close() + if mr.StatusCode != http.StatusNotFound { + t.Fatalf("minimal: /register should be 404 when ClientsStore==nil, got %d", mr.StatusCode) + } +} + +// ErrInvalid is a sentinel error used by tests to simulate token verification failures +var ErrInvalid = &struct{ error }{} + +// mustParseURL parses a URL and panics on error for test setup convenience +func mustParseURL(s string) *url.URL { u, _ := url.Parse(s); return u } + +// strPtr returns a pointer to the provided string +func strPtr(s string) *string { return &s } + +// expectStr asserts a string field in a metadata map +func expectStr(t *testing.T, m map[string]any, key, want string) { + t.Helper() + v, ok := m[key] + if !ok { + t.Fatalf("missing key %q", key) + } + vs, _ := v.(string) + if vs != want { + t.Fatalf("%s mismatch: got %q want %q", key, vs, want) + } +} + +// expectArr asserts a string slice field in a metadata map +func expectArr(t *testing.T, m map[string]any, key string, want []string) { + t.Helper() + v, ok := m[key] + if !ok { + t.Fatalf("missing key %q", key) + } + arr, ok := v.([]any) + if !ok { + t.Fatalf("%s is not array", key) + } + got := make([]string, 0, len(arr)) + for _, x := range arr { + if s, ok := x.(string); ok { + got = append(got, s) + } + } + if len(want) == 1 && strings.HasPrefix(want[0], "__any__:") { + needle := strings.TrimPrefix(want[0], "__any__:") + found := false + for _, g := range got { + if g == needle { + found = true + break + } + } + if !found { + t.Fatalf("%s expected to contain %q, got %v", key, needle, got) + } + return + } + if len(got) != len(want) { + t.Fatalf("%s length mismatch: got %v want %v", key, got, want) + } + for i := range got { + if got[i] != want[i] { + t.Fatalf("%s[%d] mismatch: got %q want %q", key, i, got[i], want[i]) + } + } +} + +// containsAny encodes a loose expectation for any one of the provided values +func containsAny(values ...string) []string { + if len(values) == 0 { + return nil + } + return []string{"__any__:" + values[0]} +} diff --git a/internal/auth/server/token_verifier.go b/internal/auth/server/token_verifier.go new file mode 100644 index 0000000..11f4553 --- /dev/null +++ b/internal/auth/server/token_verifier.go @@ -0,0 +1,845 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package server + +import ( + "context" + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "strings" + "sync" + "time" + + oauthErrors "trpc.group/trpc-go/trpc-mcp-go/internal/errors" + + "github.com/lestrrat-go/jwx/v2/jwk" + "github.com/lestrrat-go/jwx/v2/jws" + "github.com/lestrrat-go/jwx/v2/jwt" +) + +// standardClaims are filtered out of Extra because they are either mapped to +// AuthInfo first-class fields or are not useful as extra metadata. +// Mapped fields: client_id -> ClientID, sub -> Subject, scope -> Scopes, +// exp -> ExpiresAt, aud -> Resource (via extractResource) +var standardClaims = map[string]bool{ + "client_id": true, + "sub": true, + "scope": true, + "exp": true, + "aud": true, + // Other common standard/housekeeping claims not needed in Extra + "iss": true, + "iat": true, + "jti": true, + "kid": true, +} + +type TokenVerifierInterface interface { + VerifyAccessToken(ctx context.Context, token string) (AuthInfo, error) +} + +// LocalJWKSConfig local JWKS configuration +type LocalJWKSConfig struct { + JWKS string // Local JWKS JSON string + File string // Local JWKS file path +} + +// RemoteJWKSConfig remote JWKS configuration +type RemoteJWKSConfig struct { + URLs []string // Remote JWKS URLs + IssuerToURL map[string]string // Mapping from issuer (iss) to remote URL + RefreshInterval time.Duration // Refresh interval +} + +// TokenVerifierConfig configuration for TokenVerifier +type TokenVerifierConfig struct { + Local *LocalJWKSConfig // Local JWKS configuration + Remote *RemoteJWKSConfig // Remote JWKS configuration + Introspection *IntrospectionConfig // Remote introspection configuration (RFC7662) +} + +// IntrospectionCredentials client credentials for introspection +type IntrospectionCredentials struct { + ClientID string + ClientSecret string +} + +// IntrospectionConfig remote introspection configuration +type IntrospectionConfig struct { + // Default introspection endpoint (optional). Used when no issuer-bound endpoint is found. + Endpoint string + // Endpoint selection by issuer (multi-tenant). + IssuerToEndpoint map[string]string + + // Default credentials and per-issuer credentials (optional). + DefaultCredentials *IntrospectionCredentials + IssuerCredentials map[string]IntrospectionCredentials + + // HTTP timeout + Timeout time.Duration + + // Cache TTL (positive) and negative cache TTL (for inactive tokens or 4xx/401, etc.). + CacheTTL time.Duration + NegativeCacheTTL time.Duration + + // Whether to fall back to introspection on JWT verification failure. + UseOnJWTFail bool +} + +// TokenVerifier holds verification configuration and helpers +type TokenVerifier struct { + localKeySet jwk.Set // Local JWKS key set + cache *jwk.Cache // Remote JWKS cache (jwx v2) + issuerToURL map[string]string // Mapping from issuer (iss) to remote URL + isRemote bool // Whether remote JWKS mode is enabled + + // RFC7662 introspection + introspectionEnabled bool + httpClient *http.Client + defaultIntrospectEP string + issuerToIntrospectEP map[string]string + defaultCreds *IntrospectionCredentials + issuerCreds map[string]IntrospectionCredentials + useIntrospectionOnFail bool + + // Simple in-memory cache + introspectCache map[string]introspectionCacheEntry + introspectCacheMu sync.RWMutex + cacheTTL time.Duration + negativeCacheTTL time.Duration +} + +type TokenVerifierFunc func(ctx context.Context, token string) (AuthInfo, error) + +func (f TokenVerifierFunc) VerifyAccessToken(ctx context.Context, token string) (AuthInfo, error) { + return f(ctx, token) +} + +// NewLocalTokenVerifier creates a TokenVerifier that uses only local JWKS +func newLocalTokenVerifier(ctx context.Context, cfg LocalJWKSConfig) (*TokenVerifier, error) { + verifier := &TokenVerifier{} + + defaultSet := jwk.NewSet() + + // Load JWKS from string + if cfg.JWKS != "" { + set, err := jwk.Parse([]byte(cfg.JWKS)) + if err != nil { + return nil, fmt.Errorf("failed to parse local JWKS: %w", err) + } + for i := 0; i < set.Len(); i++ { + key, _ := set.Key(i) + _ = defaultSet.AddKey(key) + } + } + + // Load JWKS from file + if cfg.File != "" { + set, err := jwk.ReadFile(cfg.File) + if err != nil { + return nil, fmt.Errorf("failed to parse local JWKS file: %w", err) + } + for i := 0; i < set.Len(); i++ { + key, _ := set.Key(i) + _ = defaultSet.AddKey(key) + } + } + + if defaultSet.Len() == 0 { + return nil, fmt.Errorf("must provide JWKS or File") + } + + verifier.localKeySet = defaultSet + return verifier, nil +} + +// NewRemoteTokenVerifier creates a TokenVerifier that uses only remote JWKS +func newRemoteTokenVerifier(ctx context.Context, cfg RemoteJWKSConfig) (*TokenVerifier, error) { + if len(cfg.URLs) == 0 { + return nil, fmt.Errorf("must provide at least one RemoteURL") + } + + // jwx v2 cache + cache := jwk.NewCache(ctx) + for _, url_ := range cfg.URLs { + _ = cache.Register(url_) + } + + return &TokenVerifier{ + cache: cache, + issuerToURL: func() map[string]string { + if cfg.IssuerToURL == nil { + return nil + } + m := make(map[string]string, len(cfg.IssuerToURL)) + for k, v := range cfg.IssuerToURL { + m[k] = v + } + return m + }(), + isRemote: true, + }, nil +} + +// NewIntrospectionTokenVerifier creates a TokenVerifier that uses only RFC7662 introspection +func newIntrospectionTokenVerifier(ctx context.Context, cfg IntrospectionConfig) (*TokenVerifier, error) { + verifier := &TokenVerifier{} + + to := cfg.Timeout + if to <= 0 { + to = 5 * time.Second + } + verifier.httpClient = &http.Client{Timeout: to} + verifier.defaultIntrospectEP = cfg.Endpoint + // Copy IssuerToEndpoint + if cfg.IssuerToEndpoint != nil { + verifier.issuerToIntrospectEP = make(map[string]string, len(cfg.IssuerToEndpoint)) + for k, v := range cfg.IssuerToEndpoint { + verifier.issuerToIntrospectEP[k] = v + } + } + // Copy DefaultCredentials + if cfg.DefaultCredentials != nil { + dc := *cfg.DefaultCredentials + verifier.defaultCreds = &dc + } + // Copy IssuerCredentials + if cfg.IssuerCredentials != nil { + verifier.issuerCreds = make(map[string]IntrospectionCredentials, len(cfg.IssuerCredentials)) + for k, v := range cfg.IssuerCredentials { + verifier.issuerCreds[k] = v + } + } + verifier.useIntrospectionOnFail = cfg.UseOnJWTFail + verifier.cacheTTL = cfg.CacheTTL + if verifier.cacheTTL <= 0 { + verifier.cacheTTL = 60 * time.Second + } + verifier.negativeCacheTTL = cfg.NegativeCacheTTL + if verifier.negativeCacheTTL <= 0 { + verifier.negativeCacheTTL = 15 * time.Second + } + verifier.introspectionEnabled = true + verifier.introspectCache = make(map[string]introspectionCacheEntry) + return verifier, nil +} + +// NewTokenVerifier creates a comprehensive TokenVerifier. +// Provide any one or more configurations. The SDK will automatically prefer Local → Remote → Introspection (if enabled). +func NewTokenVerifier(ctx context.Context, cfg TokenVerifierConfig) (*TokenVerifier, error) { + var verifier *TokenVerifier + var err error + + if cfg.Remote != nil && len(cfg.Remote.URLs) > 0 { + verifier, err = newRemoteTokenVerifier(ctx, *cfg.Remote) + if err != nil { + return nil, err + } + } + + if cfg.Local != nil && (cfg.Local.JWKS != "" || cfg.Local.File != "") { + localVerifier, err := newLocalTokenVerifier(ctx, *cfg.Local) + if err != nil { + return nil, err + } + + if verifier != nil { + verifier.localKeySet = localVerifier.localKeySet + } else { + verifier = localVerifier + } + } + + if verifier == nil { + // If no JWKS is provided, allow constructing an introspection-only mode + if cfg.Introspection != nil { + return newIntrospectionTokenVerifier(ctx, *cfg.Introspection) + } + return nil, errors.New("no verification method configured: configure Local JWKS (Local), or Remote JWKS (Remote), or Introspection") + } + + // Initialize introspection (optional) + if cfg.Introspection != nil { + to := cfg.Introspection.Timeout + if to <= 0 { + to = 5 * time.Second + } + verifier.httpClient = &http.Client{Timeout: to} + verifier.defaultIntrospectEP = cfg.Introspection.Endpoint + // Copy IssuerToEndpoint to avoid external mutations + if cfg.Introspection.IssuerToEndpoint != nil { + verifier.issuerToIntrospectEP = make(map[string]string, len(cfg.Introspection.IssuerToEndpoint)) + for k, v := range cfg.Introspection.IssuerToEndpoint { + verifier.issuerToIntrospectEP[k] = v + } + } + // Copy DefaultCredentials + if cfg.Introspection.DefaultCredentials != nil { + dc := *cfg.Introspection.DefaultCredentials + verifier.defaultCreds = &dc + } + // Copy IssuerCredentials + if cfg.Introspection.IssuerCredentials != nil { + verifier.issuerCreds = make(map[string]IntrospectionCredentials, len(cfg.Introspection.IssuerCredentials)) + for k, v := range cfg.Introspection.IssuerCredentials { + verifier.issuerCreds[k] = v + } + } + verifier.useIntrospectionOnFail = cfg.Introspection.UseOnJWTFail + verifier.cacheTTL = cfg.Introspection.CacheTTL + if verifier.cacheTTL <= 0 { + verifier.cacheTTL = 60 * time.Second + } + verifier.negativeCacheTTL = cfg.Introspection.NegativeCacheTTL + if verifier.negativeCacheTTL <= 0 { + verifier.negativeCacheTTL = 15 * time.Second + } + verifier.introspectionEnabled = true + verifier.introspectCache = make(map[string]introspectionCacheEntry) + } + + // Do not set an explicit "mode"; choose dynamically during Verify based on configuration + + return verifier, nil +} + +// VerifyAccessToken verifies an access token and returns AuthInfo or error +func (v *TokenVerifier) VerifyAccessToken(ctx context.Context, tokenStr string) (AuthInfo, error) { + // If no JWKS configured and introspection is enabled: directly use introspection (works for opaque/JWT) + if v.localKeySet == nil && !v.isRemote && v.introspectionEnabled { + ai, err := v.introspectAccessToken(ctx, tokenStr, "") + if err != nil { + return AuthInfo{}, oauthErrors.NewOAuthError(oauthErrors.ErrInvalidToken, "failed to verify token", "") + } + return ai, nil + } + + // Parse token first (without verifying signature) to obtain iss; if parsing fails and introspection is enabled, try introspection directly (supports opaque tokens). + unverifiedToken, err := jwt.ParseInsecure([]byte(tokenStr)) + if err != nil { + if v.introspectionEnabled { + if ai, ierr := v.introspectAccessToken(ctx, tokenStr, ""); ierr == nil { + return ai, nil + } + } + return AuthInfo{}, oauthErrors.NewOAuthError(oauthErrors.ErrInvalidToken, "malformed token: cannot parse header/payload; if you are using opaque tokens, enable Introspection", "") + } + + // Extract issuer (iss) + iss := unverifiedToken.Issuer() + if iss == "" { + return AuthInfo{}, oauthErrors.NewOAuthError(oauthErrors.ErrInvalidToken, "missing issuer (iss) in token", "") + } + + // Extract kid from JWS header + kid, err := extractKIDFromHeader(tokenStr) + if err != nil || kid == "" { + // Try introspection fallback + if v.introspectionEnabled && v.useIntrospectionOnFail { + if ai, ierr := v.introspectAccessToken(ctx, tokenStr, iss); ierr == nil { + return ai, nil + } + } + return AuthInfo{}, oauthErrors.NewOAuthError(oauthErrors.ErrInvalidToken, "missing key id (kid) in JWS header; if your tokens omit kid, ensure the JWKS only contains one key or use Introspection fallback", "") + } + + // Try to obtain target keySet + keySet, err := v.getTargetKeySet(ctx, iss, kid) + if err != nil { + return AuthInfo{}, err + } + + // Validate token with key set and basic claims, allowing time skew + token, err := jwt.Parse([]byte(tokenStr), + jwt.WithKeySet(keySet), + jwt.WithValidate(true), + jwt.WithAcceptableSkew(30*time.Second), + // RFC 9068: exp and iat are validated automatically; here we only require presence for other claims + jwt.WithRequiredClaim("exp"), + jwt.WithRequiredClaim("aud"), + jwt.WithRequiredClaim("sub"), + jwt.WithRequiredClaim("iat"), + ) + if err != nil || token == nil { + // On JWT signature/claims validation failure, optionally fall back to introspection + if v.introspectionEnabled && v.useIntrospectionOnFail { + if ai, ierr := v.introspectAccessToken(ctx, tokenStr, iss); ierr == nil { + return ai, nil + } + } + return AuthInfo{}, oauthErrors.NewOAuthError(oauthErrors.ErrInvalidToken, "signature validation failed or claims invalid; ensure JWKS is configured for issuer or enable Introspection fallback", "") + } + + // Ensure non-empty subject + if sub := token.Subject(); sub == "" { + return AuthInfo{}, oauthErrors.NewOAuthError(oauthErrors.ErrInvalidToken, "missing required 'sub' claim", "") + } + + authInfo, err := v.convertJWTToAuthInfo(token, tokenStr) + if err != nil { + return AuthInfo{}, err + } + return authInfo, nil +} + +func (v *TokenVerifier) getTargetKeySet(ctx context.Context, iss, kid string) (jwk.Set, error) { + // Prefer local JWKS + if v.localKeySet != nil { + if _, ok := v.localKeySet.LookupKeyID(kid); ok { + return v.localKeySet, nil + } + } + + // If remote mode enabled, try remote JWKS + if v.isRemote { + if url_, ok := v.issuerToURL[iss]; ok { + if v.cache != nil { + if keySet, err := v.cache.Get(ctx, url_); err == nil { + if _, ok := keySet.LookupKeyID(kid); !ok { + if refreshed, ferr := jwk.Fetch(ctx, url_); ferr == nil { + if _, ok2 := refreshed.LookupKeyID(kid); ok2 { + return refreshed, nil + } + } + } + return keySet, nil + } + } + keySet, err := jwk.Fetch(ctx, url_) + if err != nil { + return nil, fmt.Errorf("failed to fetch remote JWKS for issuer %s (url=%s): %w", iss, url_, err) + } + return keySet, nil + } + return nil, fmt.Errorf("no remote JWKS URL found for issuer %s: provide Remote.IssuerToURL mapping", iss) + } + + return nil, fmt.Errorf("no JWKS found for issuer %s: neither Local nor Remote key set available", iss) +} + +// ---- RFC7662 introspection implementation ---- + +type introspectionCacheEntry struct { + authInfo AuthInfo + inactive bool + expiresAt time.Time +} + +func (v *TokenVerifier) resolveIntrospectionEndpoint(issuer string) (string, *IntrospectionCredentials) { + ep := "" + if issuer != "" && v.issuerToIntrospectEP != nil { + if e, ok := v.issuerToIntrospectEP[issuer]; ok { + ep = e + } + } + if ep == "" { + ep = v.defaultIntrospectEP + } + var creds *IntrospectionCredentials + if issuer != "" && v.issuerCreds != nil { + if c, ok := v.issuerCreds[issuer]; ok { + cc := c + creds = &cc + } + } + if creds == nil { + creds = v.defaultCreds + } + return ep, creds +} + +func (v *TokenVerifier) introspectionCacheKey(endpoint, token string) string { + return endpoint + "|" + token +} + +func (v *TokenVerifier) loadFromIntrospectionCache(key string) (introspectionCacheEntry, bool) { + v.introspectCacheMu.RLock() + defer v.introspectCacheMu.RUnlock() + entry, ok := v.introspectCache[key] + if !ok { + return introspectionCacheEntry{}, false + } + if time.Now().After(entry.expiresAt) { + return introspectionCacheEntry{}, false + } + return entry, true +} + +func (v *TokenVerifier) storeToIntrospectionCache(key string, entry introspectionCacheEntry) { + v.introspectCacheMu.Lock() + v.introspectCache[key] = entry + v.introspectCacheMu.Unlock() +} + +func (v *TokenVerifier) introspectAccessToken(ctx context.Context, tokenStr, issuer string) (AuthInfo, error) { + if !v.introspectionEnabled { + return AuthInfo{}, errors.New("introspection not enabled") + } + endpoint, creds := v.resolveIntrospectionEndpoint(issuer) + if endpoint == "" { + return AuthInfo{}, errors.New("no introspection endpoint configured: set Introspection.Endpoint or IssuerToEndpoint for the issuer") + } + + key := v.introspectionCacheKey(endpoint, tokenStr) + if entry, ok := v.loadFromIntrospectionCache(key); ok { + if entry.inactive { + return AuthInfo{}, oauthErrors.NewOAuthError(oauthErrors.ErrInvalidToken, "inactive token", "") + } + return entry.authInfo, nil + } + + form := url.Values{} + form.Set("token", tokenStr) + form.Set("token_type_hint", "access_token") + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, strings.NewReader(form.Encode())) + if err != nil { + return AuthInfo{}, err + } + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + if creds != nil && creds.ClientID != "" { + basic := creds.ClientID + ":" + creds.ClientSecret + req.Header.Set("Authorization", "Basic "+base64.StdEncoding.EncodeToString([]byte(basic))) + } + + resp, err := v.httpClient.Do(req) + if err != nil { + return AuthInfo{}, err + } + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + if resp.StatusCode != http.StatusOK { + v.storeToIntrospectionCache(key, introspectionCacheEntry{inactive: true, expiresAt: time.Now().Add(v.negativeCacheTTL)}) + return AuthInfo{}, oauthErrors.NewOAuthError(oauthErrors.ErrInvalidToken, "introspection request failed", "") + } + + var payload map[string]interface{} + if err := json.Unmarshal(body, &payload); err != nil { + return AuthInfo{}, err + } + active, _ := payload["active"].(bool) + if !active { + v.storeToIntrospectionCache(key, introspectionCacheEntry{inactive: true, expiresAt: time.Now().Add(v.negativeCacheTTL)}) + return AuthInfo{}, oauthErrors.NewOAuthError(oauthErrors.ErrInvalidToken, "inactive token", "") + } + + ai, err := v.convertIntrospectionToAuthInfo(payload, tokenStr) + if err != nil { + return AuthInfo{}, err + } + + ttl := v.cacheTTL + if expV, ok := payload["exp"]; ok { + switch t := expV.(type) { + case float64: + expTs := time.Unix(int64(t), 0) + if expTs.After(time.Now()) { + rem := time.Until(expTs) + if rem < ttl { + ttl = rem + } + } + case json.Number: + if v, err2 := t.Int64(); err2 == nil { + expTs := time.Unix(v, 0) + if expTs.After(time.Now()) { + rem := time.Until(expTs) + if rem < ttl { + ttl = rem + } + } + } + } + } + v.storeToIntrospectionCache(key, introspectionCacheEntry{authInfo: ai, expiresAt: time.Now().Add(ttl)}) + return ai, nil +} + +func (v *TokenVerifier) convertIntrospectionToAuthInfo(payload map[string]interface{}, tokenStr string) (AuthInfo, error) { + var ai AuthInfo + ai.Token = tokenStr + if cid, _ := payload["client_id"].(string); cid != "" { + ai.ClientID = cid + } + if sc, ok := payload["scope"]; ok { + ai.Scopes = parseScopesFromRaw(sc) + } + switch exp := payload["exp"].(type) { + case float64: + ts := int64(exp) + ai.ExpiresAt = &ts + case json.Number: + if v, err := exp.Int64(); err == nil { + ai.ExpiresAt = &v + } + } + if r, err := extractResourceFromIntrospection(payload["aud"]); err == nil { + ai.Resource = r + } + + extra := make(map[string]interface{}) + for k, v := range payload { + if standardClaims[k] { + continue + } + switch k { + case "active", "username", "token_type", "token_type_hint": + continue + case "client_id", "scope", "exp", "aud", "iss", "sub", "iat", "jti": + continue + default: + extra[k] = v + } + } + if len(extra) > 0 { + ai.Extra = extra + } + return ai, nil +} + +func parseScopesFromRaw(raw interface{}) []string { + switch s := raw.(type) { + case string: + if s == "" { + return nil + } + return strings.Split(s, " ") + case []interface{}: + var scopes []string + for _, v := range s { + if str, ok := v.(string); ok && str != "" { + scopes = append(scopes, str) + } + } + if len(scopes) == 0 { + return nil + } + return scopes + default: + return nil + } +} + +func extractResourceFromIntrospection(audRaw interface{}) (*url.URL, error) { + if audRaw == nil { + return nil, nil + } + var candidates []string + switch v := audRaw.(type) { + case string: + if v != "" { + candidates = []string{v} + } + case []interface{}: + for _, it := range v { + if s, ok := it.(string); ok && s != "" { + candidates = append(candidates, s) + } + } + case []string: + candidates = v + } + for _, c := range candidates { + looksLikeURL := strings.HasPrefix(c, "http://") || strings.HasPrefix(c, "https://") || strings.Contains(c, "://") + if !looksLikeURL { + continue + } + u, err := url.Parse(c) + if err != nil || u == nil || u.Scheme == "" || u.Host == "" { + continue + } + u.Fragment = "" // Remove fragment (per RFC 8707) + return u, nil + } + return nil, nil +} + +// extractKIDFromHeader 从 JWS Header 提取 kid +func extractKIDFromHeader(tokenStr string) (string, error) { + msg, err := jws.Parse([]byte(tokenStr)) + if err != nil { + return "", fmt.Errorf("failed to parse JWS: %w", err) + } + sigs := msg.Signatures() + if len(sigs) == 0 { + return "", errors.New("no signatures found in JWS") + } + + // Prefer protected headers + if ph := sigs[0].ProtectedHeaders(); ph != nil { + if v, ok := ph.Get(jws.KeyIDKey); ok { + if kid, ok2 := v.(string); ok2 && kid != "" { + return kid, nil + } + } + } + return "", errors.New("missing kid in JWS header") +} + +// convertJWTToAuthInfo converts jwt.Token to AuthInfo structure. +func (v *TokenVerifier) convertJWTToAuthInfo(token jwt.Token, tokenStr string) (AuthInfo, error) { + authInfo := AuthInfo{Token: tokenStr} + + // Write exp -> ExpiresAt (must be done first) + if exp := token.Expiration(); !exp.IsZero() { + ts := exp.Unix() + authInfo.ExpiresAt = &ts + } else { + // This case should not be reached because WithRequiredClaim("exp") was used in Parse + // But for robustness, return invalid_token for clarity + return AuthInfo{}, oauthErrors.NewOAuthError(oauthErrors.ErrInvalidToken, "missing exp claim", "") + } + + // Extract OAuth-related fields + var err error + if authInfo.ClientID, err = extractClientID(token); err != nil { + return AuthInfo{}, err + } + if authInfo.Resource, err = extractResource(token); err != nil { + return AuthInfo{}, err + } + if authInfo.Scopes, err = extractScopes(token); err != nil { + return AuthInfo{}, err + } + + // Write subject -> AuthInfo.Subject + if s := token.Subject(); s != "" { + authInfo.Subject = s + } + + // Other custom claims + authInfo.Extra = extractExtra(token) + return authInfo, nil +} + +// extractClientID extracts client ID (optional) +func extractClientID(token jwt.Token) (string, error) { + // client_id is not mandatory for access tokens (RFC9068) + if v, ok := token.Get("client_id"); ok { + if s, ok2 := v.(string); ok2 && s != "" { + return s, nil + } + } + // Fallback to azp (often used in OIDC) + if v, ok := token.Get("azp"); ok { + if s, ok2 := v.(string); ok2 && s != "" { + return s, nil + } + } + // Missing client identifier is acceptable + return "", nil +} + +// extractScopes extracts scopes from various claim formats (scope/scp). Missing is acceptable. +func extractScopes(token jwt.Token) ([]string, error) { + var raw interface{} + // Prefer RFC6749 style "scope" (space-delimited string or array) + if v, ok := token.Get("scope"); ok { + raw = v + } else { + // Fallback to "scp" (array of strings used by some providers) + if v2, ok2 := token.Get("scp"); ok2 { + raw = v2 + } else { + // No scopes present → treat as empty without error + return nil, nil + } + } + + switch s := raw.(type) { + case string: + if s == "" { + return nil, nil + } + return strings.Split(s, " "), nil + case []string: + if len(s) == 0 { + return nil, nil + } + return s, nil + case []interface{}: + if len(s) == 0 { + return nil, nil + } + var scopes []string + for _, v := range s { + if str, ok := v.(string); ok { + scopes = append(scopes, str) + } + } + if len(scopes) == 0 { + return nil, nil + } + return scopes, nil + default: + // Unknown format → ignore rather than failing hard for compatibility + return nil, nil + } +} + +// extractResource extracts resource information +func extractResource(token jwt.Token) (*url.URL, error) { + aud := token.Audience() + if len(aud) == 0 { + return nil, fmt.Errorf("missing required 'aud' claim") + } + + // Iterate to find the first value that looks like a URL and is parseable as HTTP(S); + // if none are URLs, return nil to indicate no resource indicator provided. + for _, candidate := range aud { + if candidate == "" { + continue + } + looksLikeURL := strings.HasPrefix(candidate, "http://") || strings.HasPrefix(candidate, "https://") || strings.Contains(candidate, "://") + if !looksLikeURL { + continue + } + resourceURL, err := url.Parse(candidate) + if err != nil || resourceURL == nil { + continue + } + if resourceURL.Scheme == "" || resourceURL.Host == "" { + continue + } + resourceURL.Fragment = "" // Remove fragment (per RFC 8707) + return resourceURL, nil + } + return nil, nil +} + +// extractExtra extracts custom claims to Extra map +func extractExtra(token jwt.Token) map[string]interface{} { + all, _ := token.AsMap(context.Background()) + if len(all) == 0 { + return nil + } + + extra := make(map[string]interface{}) + for key, value := range all { + if standardClaims[key] { + continue + } + switch key { + case "active", "username", "token_type", "token_type_hint": + // Known noise in introspection/JWT contexts; exclude from Extra + continue + default: + extra[key] = value + } + } + if len(extra) == 0 { + return nil + } + return extra +} + +// Note: TokenVerifier is statically configured. After initialization, it does not support dynamically adding issuer mappings or clearing the local KeySet. diff --git a/internal/auth/server/token_verifier_test.go b/internal/auth/server/token_verifier_test.go new file mode 100644 index 0000000..c415697 --- /dev/null +++ b/internal/auth/server/token_verifier_test.go @@ -0,0 +1,840 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package server + +import ( + "context" + "crypto/rand" + "crypto/rsa" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "os" + "testing" + "time" + + "github.com/lestrrat-go/jwx/v2/jwa" + "github.com/lestrrat-go/jwx/v2/jwk" + "github.com/lestrrat-go/jwx/v2/jwt" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// Test helper functions + +// generateRSAKey generates a new RSA key pair for testing +func generateRSAKey() (*rsa.PrivateKey, error) { + return rsa.GenerateKey(rand.Reader, 2048) +} + +// createTestJWK creates a test JWK from RSA key +func createTestJWK(privateKey *rsa.PrivateKey, keyID string) (jwk.Key, error) { + // Import the public key part only + publicKey := &privateKey.PublicKey + key, err := jwk.FromRaw(publicKey) + if err != nil { + return nil, err + } + + if err := key.Set(jwk.KeyIDKey, keyID); err != nil { + return nil, err + } + + if err := key.Set(jwk.AlgorithmKey, "RS256"); err != nil { + return nil, err + } + + if err := key.Set(jwk.KeyUsageKey, "sig"); err != nil { + return nil, err + } + + return key, nil +} + +// createTestToken creates a test JWT token +func createTestToken(privateKey *rsa.PrivateKey, keyID string, claims map[string]interface{}) (string, error) { + key, err := jwk.FromRaw(privateKey) + if err != nil { + return "", err + } + + if err := key.Set(jwk.KeyIDKey, keyID); err != nil { + return "", err + } + + now := time.Now() + token := jwt.New() + + // Set standard claims + _ = token.Set(jwt.IssuerKey, "https://example.com") + _ = token.Set(jwt.SubjectKey, "user123") + _ = token.Set(jwt.AudienceKey, []string{"https://api.example.com"}) + _ = token.Set(jwt.ExpirationKey, now.Add(time.Hour)) + _ = token.Set(jwt.IssuedAtKey, now) + _ = token.Set(jwt.JwtIDKey, "jti-123") + _ = token.Set("client_id", "test-client") + _ = token.Set("scope", "read write") + _ = token.Set("kid", keyID) + + // Set custom claims + for k, v := range claims { + _ = token.Set(k, v) + } + + signed, err := jwt.Sign(token, jwt.WithKey(jwa.RS256, key)) + if err != nil { + return "", err + } + + return string(signed), nil +} + +// createTestJWKS creates a test JWKS JSON string +func createTestJWKS(keys ...jwk.Key) string { + set := jwk.NewSet() + for _, key := range keys { + _ = set.AddKey(key) + } + + buf, _ := json.Marshal(set) + return string(buf) +} + +// Test fixtures +func setupTestKeys(t *testing.T) (*rsa.PrivateKey, jwk.Key, string) { + privateKey, err := generateRSAKey() + require.NoError(t, err) + + publicKey, err := createTestJWK(privateKey, "test-key-1") + require.NoError(t, err) + + jwksJSON := createTestJWKS(publicKey) + + return privateKey, publicKey, jwksJSON +} + +func TestTokenVerifierFunc_VerifyAccessToken(t *testing.T) { + ctx := context.Background() + + // Define a fake verifier function + fn := TokenVerifierFunc(func(ctx context.Context, token string) (AuthInfo, error) { + if token == "valid" { + return AuthInfo{Token: token, ClientID: "test-client"}, nil + } + return AuthInfo{}, errors.New("invalid token") + }) + + // Success path + authInfo, err := fn.VerifyAccessToken(ctx, "valid") + assert.NoError(t, err) + assert.Equal(t, "valid", authInfo.Token) + assert.Equal(t, "test-client", authInfo.ClientID) + + // Failure path + authInfo, err = fn.VerifyAccessToken(ctx, "invalid") + assert.Error(t, err) + assert.Empty(t, authInfo.Token) +} + +// Tests for NewLocalTokenVerifier + +func TestNewLocalTokenVerifier_WithJWKSString(t *testing.T) { + ctx := context.Background() + _, _, jwksJSON := setupTestKeys(t) + + cfg := LocalJWKSConfig{ + JWKS: jwksJSON, + } + + verifier, err := newLocalTokenVerifier(ctx, cfg) + assert.NoError(t, err) + assert.NotNil(t, verifier) + assert.NotNil(t, verifier.localKeySet) + assert.False(t, verifier.isRemote) + assert.Equal(t, 1, verifier.localKeySet.Len()) +} + +func TestNewLocalTokenVerifier_WithFile(t *testing.T) { + ctx := context.Background() + _, _, jwksJSON := setupTestKeys(t) + + // Create temporary file + tmpFile, err := os.CreateTemp("", "jwks-*.json") + require.NoError(t, err) + defer os.Remove(tmpFile.Name()) + + _, err = tmpFile.WriteString(jwksJSON) + require.NoError(t, err) + _ = tmpFile.Close() + + cfg := LocalJWKSConfig{ + File: tmpFile.Name(), + } + + verifier, err := newLocalTokenVerifier(ctx, cfg) + assert.NoError(t, err) + assert.NotNil(t, verifier) + assert.Equal(t, 1, verifier.localKeySet.Len()) +} + +func TestNewLocalTokenVerifier_WithBothJWKSAndFile(t *testing.T) { + ctx := context.Background() + + // Create two different keys + _, _, jwksJSON1 := setupTestKeys(t) + + privateKey2, err := generateRSAKey() + require.NoError(t, err) + publicKey2, err := createTestJWK(privateKey2, "test-key-2") + require.NoError(t, err) + jwksJSON2 := createTestJWKS(publicKey2) + + // Create temporary file with second key + tmpFile, err := os.CreateTemp("", "jwks-*.json") + require.NoError(t, err) + defer os.Remove(tmpFile.Name()) + + _, err = tmpFile.WriteString(jwksJSON2) + require.NoError(t, err) + _ = tmpFile.Close() + + cfg := LocalJWKSConfig{ + JWKS: jwksJSON1, + File: tmpFile.Name(), + } + + verifier, err := newLocalTokenVerifier(ctx, cfg) + assert.NoError(t, err) + assert.NotNil(t, verifier) + assert.Equal(t, 2, verifier.localKeySet.Len()) // Should have both keys +} + +func TestNewLocalTokenVerifier_EmptyConfig(t *testing.T) { + ctx := context.Background() + cfg := LocalJWKSConfig{} + + verifier, err := newLocalTokenVerifier(ctx, cfg) + assert.Error(t, err) + assert.Nil(t, verifier) + assert.Contains(t, err.Error(), "must provide JWKS or File") +} + +func TestNewLocalTokenVerifier_InvalidJWKS(t *testing.T) { + ctx := context.Background() + cfg := LocalJWKSConfig{ + JWKS: "invalid-json", + } + + verifier, err := newLocalTokenVerifier(ctx, cfg) + assert.Error(t, err) + assert.Nil(t, verifier) + assert.Contains(t, err.Error(), "failed to parse local JWKS") +} + +// Tests for NewRemoteTokenVerifier + +func TestNewRemoteTokenVerifier_Success(t *testing.T) { + ctx := context.Background() + _, _, jwksJSON := setupTestKeys(t) + + // Create test server + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(jwksJSON)) + })) + defer server.Close() + + cfg := RemoteJWKSConfig{ + URLs: []string{server.URL}, + IssuerToURL: map[string]string{ + "https://example.com": server.URL, + }, + RefreshInterval: time.Minute, + } + + verifier, err := newRemoteTokenVerifier(ctx, cfg) + assert.NoError(t, err) + assert.NotNil(t, verifier) + assert.True(t, verifier.isRemote) + assert.NotNil(t, verifier.cache) + assert.Equal(t, server.URL, verifier.issuerToURL["https://example.com"]) +} + +func TestNewRemoteTokenVerifier_EmptyURLs(t *testing.T) { + ctx := context.Background() + cfg := RemoteJWKSConfig{} + + verifier, err := newRemoteTokenVerifier(ctx, cfg) + assert.Error(t, err) + assert.Nil(t, verifier) + assert.Contains(t, err.Error(), "must provide at least one RemoteURL") +} + +func TestNewRemoteTokenVerifier_DefaultRefreshInterval(t *testing.T) { + ctx := context.Background() + _, _, jwksJSON := setupTestKeys(t) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(jwksJSON)) + })) + defer server.Close() + + cfg := RemoteJWKSConfig{ + URLs: []string{server.URL}, + // RefreshInterval is 0, should use default (60 minutes) + } + + verifier, err := newRemoteTokenVerifier(ctx, cfg) + assert.NoError(t, err) + assert.NotNil(t, verifier) +} + +// Tests for NewTokenVerifier + +func TestNewTokenVerifier_LocalOnly(t *testing.T) { + ctx := context.Background() + _, _, jwksJSON := setupTestKeys(t) + + cfg := TokenVerifierConfig{ + Local: &LocalJWKSConfig{ + JWKS: jwksJSON, + }, + } + + verifier, err := NewTokenVerifier(ctx, cfg) + assert.NoError(t, err) + assert.NotNil(t, verifier) + assert.NotNil(t, verifier.localKeySet) + assert.False(t, verifier.isRemote) +} + +func TestNewTokenVerifier_RemoteOnly(t *testing.T) { + ctx := context.Background() + _, _, jwksJSON := setupTestKeys(t) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(jwksJSON)) + })) + defer server.Close() + + cfg := TokenVerifierConfig{ + Remote: &RemoteJWKSConfig{ + URLs: []string{server.URL}, + }, + } + + verifier, err := NewTokenVerifier(ctx, cfg) + assert.NoError(t, err) + assert.NotNil(t, verifier) + assert.True(t, verifier.isRemote) + assert.NotNil(t, verifier.cache) +} + +func TestNewTokenVerifier_Combined(t *testing.T) { + ctx := context.Background() + _, _, jwksJSON := setupTestKeys(t) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(jwksJSON)) + })) + defer server.Close() + + cfg := TokenVerifierConfig{ + Local: &LocalJWKSConfig{ + JWKS: jwksJSON, + }, + Remote: &RemoteJWKSConfig{ + URLs: []string{server.URL}, + }, + } + + verifier, err := NewTokenVerifier(ctx, cfg) + assert.NoError(t, err) + assert.NotNil(t, verifier) + assert.True(t, verifier.isRemote) + assert.NotNil(t, verifier.cache) + assert.NotNil(t, verifier.localKeySet) +} + +// Tests for newIntrospectionTokenVerifier + +func TestNewIntrospectionTokenVerifier_Success(t *testing.T) { + ctx := context.Background() + + cfg := IntrospectionConfig{ + Endpoint: "http://example.test/introspect", + Timeout: 2 * time.Second, + CacheTTL: 3 * time.Second, + NegativeCacheTTL: 1 * time.Second, + UseOnJWTFail: true, + IssuerToEndpoint: map[string]string{ + "https://issuer.example": "http://example.test/iss-introspect", + }, + DefaultCredentials: &IntrospectionCredentials{ClientID: "cid", ClientSecret: "sec"}, + IssuerCredentials: map[string]IntrospectionCredentials{ + "https://issuer.example": {ClientID: "icid", ClientSecret: "isec"}, + }, + } + + v, err := newIntrospectionTokenVerifier(ctx, cfg) + require.NoError(t, err) + require.NotNil(t, v) + + assert.True(t, v.introspectionEnabled) + assert.NotNil(t, v.httpClient) + assert.Equal(t, cfg.Endpoint, v.defaultIntrospectEP) + assert.Equal(t, cfg.CacheTTL, v.cacheTTL) + assert.Equal(t, cfg.NegativeCacheTTL, v.negativeCacheTTL) + assert.Equal(t, cfg.UseOnJWTFail, v.useIntrospectionOnFail) + assert.Equal(t, cfg.IssuerToEndpoint["https://issuer.example"], v.issuerToIntrospectEP["https://issuer.example"]) + require.NotNil(t, v.defaultCreds) + assert.Equal(t, "cid", v.defaultCreds.ClientID) + assert.Equal(t, "sec", v.defaultCreds.ClientSecret) + assert.Equal(t, "icid", v.issuerCreds["https://issuer.example"].ClientID) + assert.Equal(t, "isec", v.issuerCreds["https://issuer.example"].ClientSecret) +} + +func TestNewIntrospectionTokenVerifier_Defaults(t *testing.T) { + ctx := context.Background() + + cfg := IntrospectionConfig{ + Endpoint: "http://example.test/introspect", + // Timeout, CacheTTL, NegativeCacheTTL left as zero to trigger defaults + } + + v, err := newIntrospectionTokenVerifier(ctx, cfg) + require.NoError(t, err) + require.NotNil(t, v) + + // Default timeouts + require.NotNil(t, v.httpClient) + assert.Equal(t, 5*time.Second, v.httpClient.Timeout) + assert.Equal(t, 60*time.Second, v.cacheTTL) + assert.Equal(t, 15*time.Second, v.negativeCacheTTL) +} + +// Ensure NewTokenVerifier (introspection-only) constructs an introspection-enabled verifier +func TestNewTokenVerifier_IntrospectionOnly_Constructed(t *testing.T) { + ctx := context.Background() + + v, err := NewTokenVerifier(ctx, TokenVerifierConfig{ + Introspection: &IntrospectionConfig{Endpoint: "http://example.test/introspect"}, + }) + require.NoError(t, err) + require.NotNil(t, v) + + assert.True(t, v.introspectionEnabled) + assert.Nil(t, v.localKeySet) + assert.False(t, v.isRemote) +} + +// Tests for VerifyAccessToken + +func TestVerifyAccessToken_LocalSuccess(t *testing.T) { + ctx := context.Background() + privateKey, _, jwksJSON := setupTestKeys(t) + + cfg := LocalJWKSConfig{ + JWKS: jwksJSON, + } + + verifier, err := newLocalTokenVerifier(ctx, cfg) + require.NoError(t, err) + + // Create valid token + tokenStr, err := createTestToken(privateKey, "test-key-1", map[string]interface{}{ + "custom_claim": "custom_value", + }) + require.NoError(t, err) + + authInfo, err := verifier.VerifyAccessToken(ctx, tokenStr) + assert.NoError(t, err) + assert.Equal(t, tokenStr, authInfo.Token) + assert.Equal(t, "test-client", authInfo.ClientID) + assert.Equal(t, []string{"read", "write"}, authInfo.Scopes) + assert.NotNil(t, authInfo.Resource) + assert.Equal(t, "https://api.example.com", authInfo.Resource.String()) + assert.Equal(t, "custom_value", authInfo.Extra["custom_claim"]) +} + +func TestVerifyAccessToken_InvalidToken(t *testing.T) { + ctx := context.Background() + _, _, jwksJSON := setupTestKeys(t) + + cfg := LocalJWKSConfig{ + JWKS: jwksJSON, + } + + verifier, err := newLocalTokenVerifier(ctx, cfg) + require.NoError(t, err) + + authInfo, err := verifier.VerifyAccessToken(ctx, "invalid-token") + assert.Error(t, err) + assert.Empty(t, authInfo) +} + +func TestVerifyAccessToken_ExpiredToken(t *testing.T) { + ctx := context.Background() + privateKey, _, jwksJSON := setupTestKeys(t) + + cfg := LocalJWKSConfig{ + JWKS: jwksJSON, + } + + verifier, err := newLocalTokenVerifier(ctx, cfg) + require.NoError(t, err) + + // Create expired token + key, err := jwk.FromRaw(privateKey) + require.NoError(t, err) + + err = key.Set(jwk.KeyIDKey, "test-key-1") + require.NoError(t, err) + + now := time.Now() + token := jwt.New() + + _ = token.Set(jwt.IssuerKey, "https://example.com") + _ = token.Set(jwt.SubjectKey, "user123") + _ = token.Set(jwt.AudienceKey, []string{"https://api.example.com"}) + _ = token.Set(jwt.ExpirationKey, now.Add(-time.Hour)) // Expired 1 hour ago + _ = token.Set(jwt.IssuedAtKey, now.Add(-2*time.Hour)) + _ = token.Set(jwt.JwtIDKey, "jti-123") + _ = token.Set("client_id", "test-client") + _ = token.Set("scope", "read write") + _ = token.Set("kid", "test-key-1") + + signed, err := jwt.Sign(token, jwt.WithKey(jwa.RS256, key)) + require.NoError(t, err) + + authInfo, err := verifier.VerifyAccessToken(ctx, string(signed)) + assert.Error(t, err) + assert.Empty(t, authInfo) +} + +func TestVerifyAccessToken_MissingRequiredClaims(t *testing.T) { + ctx := context.Background() + privateKey, _, jwksJSON := setupTestKeys(t) + + cfg := LocalJWKSConfig{ + JWKS: jwksJSON, + } + + verifier, err := newLocalTokenVerifier(ctx, cfg) + require.NoError(t, err) + + // Create token missing required claims + key, err := jwk.FromRaw(privateKey) + require.NoError(t, err) + + err = key.Set(jwk.KeyIDKey, "test-key-1") + require.NoError(t, err) + + token := jwt.New() + _ = token.Set(jwt.IssuerKey, "https://example.com") + // Missing other required claims + _ = token.Set("kid", "test-key-1") + + signed, err := jwt.Sign(token, jwt.WithKey(jwa.RS256, key)) + require.NoError(t, err) + + authInfo, err := verifier.VerifyAccessToken(ctx, string(signed)) + assert.Error(t, err) + assert.Empty(t, authInfo) +} + +func TestVerifyAccessToken_NoMatchingKey(t *testing.T) { + ctx := context.Background() + privateKey, _, jwksJSON := setupTestKeys(t) + + cfg := LocalJWKSConfig{ + JWKS: jwksJSON, + } + + verifier, err := newLocalTokenVerifier(ctx, cfg) + require.NoError(t, err) + + // Create token with different key ID + tokenStr, err := createTestToken(privateKey, "different-key-id", nil) + require.NoError(t, err) + + authInfo, err := verifier.VerifyAccessToken(ctx, tokenStr) + assert.Error(t, err) + assert.Empty(t, authInfo) +} + +// Tests for extractScopes + +func TestExtractScopes_StringFormat(t *testing.T) { + token := jwt.New() + _ = token.Set("scope", "read write admin") + + scopes, err := extractScopes(token) + assert.NoError(t, err) + assert.Equal(t, []string{"read", "write", "admin"}, scopes) +} + +func TestExtractScopes_ArrayFormat(t *testing.T) { + token := jwt.New() + _ = token.Set("scope", []string{"read", "write", "admin"}) + + scopes, err := extractScopes(token) + assert.NoError(t, err) + assert.Equal(t, []string{"read", "write", "admin"}, scopes) +} + +func TestExtractScopes_EmptyString(t *testing.T) { + token := jwt.New() + _ = token.Set("scope", "") + + scopes, err := extractScopes(token) + assert.NoError(t, err) + assert.Empty(t, scopes) +} + +func TestExtractScopes_EmptyArray(t *testing.T) { + token := jwt.New() + _ = token.Set("scope", []string{}) + + scopes, err := extractScopes(token) + assert.NoError(t, err) + assert.Empty(t, scopes) +} + +// Tests for extractResource + +func TestExtractResource_ValidURL(t *testing.T) { + token := jwt.New() + _ = token.Set(jwt.AudienceKey, []string{"https://api.example.com/resource"}) + + resource, err := extractResource(token) + assert.NoError(t, err) + assert.NotNil(t, resource) + assert.Equal(t, "https://api.example.com/resource", resource.String()) +} + +func TestExtractResource_URLWithFragment(t *testing.T) { + token := jwt.New() + _ = token.Set(jwt.AudienceKey, []string{"https://api.example.com/resource#fragment"}) + + resource, err := extractResource(token) + assert.NoError(t, err) + assert.NotNil(t, resource) + assert.Equal(t, "https://api.example.com/resource", resource.String()) // Fragment should be removed +} + +func TestExtractResource_MissingAudience(t *testing.T) { + token := jwt.New() + + resource, err := extractResource(token) + assert.Error(t, err) + assert.Nil(t, resource) +} + +// Tests for extractExtra + +func TestExtractExtra_WithCustomClaims(t *testing.T) { + token := jwt.New() + _ = token.Set(jwt.IssuerKey, "https://example.com") // Standard claim + _ = token.Set("custom_claim1", "value1") // Custom claim + _ = token.Set("custom_claim2", 123) // Custom claim + _ = token.Set("client_id", "test-client") // Standard claim + + extra := extractExtra(token) + assert.NotNil(t, extra) + assert.Equal(t, "value1", extra["custom_claim1"]) + assert.Equal(t, 123, extra["custom_claim2"]) + assert.NotContains(t, extra, "iss") // Standard claims should be excluded + assert.NotContains(t, extra, "client_id") // Standard claims should be excluded +} + +func TestExtractExtra_NoCustomClaims(t *testing.T) { + token := jwt.New() + _ = token.Set(jwt.IssuerKey, "https://example.com") + _ = token.Set("client_id", "test-client") + + extra := extractExtra(token) + assert.Nil(t, extra) // Should return nil for omitempty +} + +// Integration tests + +func TestVerifyAccessToken_RemoteJWKS(t *testing.T) { + ctx := context.Background() + privateKey, _, jwksJSON := setupTestKeys(t) + + // Create test server for JWKS + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(jwksJSON)) + })) + defer server.Close() + + cfg := RemoteJWKSConfig{ + URLs: []string{server.URL}, + IssuerToURL: map[string]string{ + "https://example.com": server.URL, + }, + } + + verifier, err := newRemoteTokenVerifier(ctx, cfg) + require.NoError(t, err) + + // Create valid token + tokenStr, err := createTestToken(privateKey, "test-key-1", nil) + require.NoError(t, err) + + authInfo, err := verifier.VerifyAccessToken(ctx, tokenStr) + assert.NoError(t, err) + assert.Equal(t, tokenStr, authInfo.Token) + assert.Equal(t, "test-client", authInfo.ClientID) +} + +func TestVerifyAccessToken_MixedMode_LocalKeyFound(t *testing.T) { + ctx := context.Background() + privateKey, _, jwksJSON := setupTestKeys(t) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(jwksJSON)) + })) + defer server.Close() + + cfg := TokenVerifierConfig{ + Local: &LocalJWKSConfig{ + JWKS: jwksJSON, + }, + Remote: &RemoteJWKSConfig{ + URLs: []string{server.URL}, + IssuerToURL: map[string]string{ + "https://example.com": server.URL, + }, + }, + } + + verifier, err := NewTokenVerifier(ctx, cfg) + require.NoError(t, err) + + // Create valid token + tokenStr, err := createTestToken(privateKey, "test-key-1", nil) + require.NoError(t, err) + + authInfo, err := verifier.VerifyAccessToken(ctx, tokenStr) + assert.NoError(t, err) + assert.Equal(t, tokenStr, authInfo.Token) +} + +// Introspection-only mode: no JWKS configured, only introspection is used. +func TestVerifyAccessToken_IntrospectionOnly_Mode(t *testing.T) { + ctx := context.Background() + + // Fake introspection endpoint which returns active token and minimal payload + introspectCalls := 0 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + introspectCalls++ + _ = r.ParseForm() + token := r.FormValue("token") + // return an active response regardless of token content + w.Header().Set("Content-Type", "application/json") + resp := map[string]interface{}{ + "active": true, + "scope": "read write", + "client_id": "cli-123", + "exp": float64(time.Now().Add(5 * time.Minute).Unix()), + "aud": "https://api.example.com", + } + // echo part to ensure parser tolerates arbitrary fields + if token != "" { + resp["token_hash"] = len(token) + } + _ = json.NewEncoder(w).Encode(resp) + })) + defer server.Close() + + v, err := NewTokenVerifier(ctx, TokenVerifierConfig{ + Introspection: &IntrospectionConfig{ + Endpoint: server.URL, + Timeout: 2 * time.Second, + CacheTTL: 2 * time.Second, + NegativeCacheTTL: 1 * time.Second, + UseOnJWTFail: true, + }, + }) + require.NoError(t, err) + + // Opaque token scenario + ai, err := v.VerifyAccessToken(ctx, "opaque-token-abc") + assert.NoError(t, err) + assert.Equal(t, "cli-123", ai.ClientID) + assert.ElementsMatch(t, []string{"read", "write"}, ai.Scopes) + assert.NotNil(t, ai.ExpiresAt) + + // Cache hit path + ai2, err := v.VerifyAccessToken(ctx, "opaque-token-abc") + assert.NoError(t, err) + assert.Equal(t, ai.ClientID, ai2.ClientID) + assert.LessOrEqual(t, introspectCalls, 2) // first call + maybe cache check +} + +// Key rotation: first JWKS does not contain target kid, second fetch returns rotated JWKS. +func TestVerifyAccessToken_RemoteJWKS_KeyRotation_RefreshOnKidMiss(t *testing.T) { + ctx := context.Background() + + // old key (won't match token) + oldPriv, err := generateRSAKey() + require.NoError(t, err) + oldPub, err := createTestJWK(oldPriv, "old-key") + require.NoError(t, err) + + // new key (used to sign token) + newPriv, err := generateRSAKey() + require.NoError(t, err) + newPub, err := createTestJWK(newPriv, "new-key") + require.NoError(t, err) + + jwksOld := createTestJWKS(oldPub) + jwksNew := createTestJWKS(newPub) + + // JWKS server: first call -> old, subsequent -> new + callCount := 0 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if callCount == 0 { + _, _ = w.Write([]byte(jwksOld)) + } else { + _, _ = w.Write([]byte(jwksNew)) + } + callCount++ + })) + defer server.Close() + + cfg := RemoteJWKSConfig{ + URLs: []string{server.URL}, + IssuerToURL: map[string]string{ + "https://example.com": server.URL, + }, + RefreshInterval: time.Minute, + } + + verifier, err := newRemoteTokenVerifier(ctx, cfg) + require.NoError(t, err) + + // Token signed by new key (kid=new-key). First cache lookup sees old JWKS + tokenStr, err := createTestToken(newPriv, "new-key", nil) + require.NoError(t, err) + + authInfo, err := verifier.VerifyAccessToken(ctx, tokenStr) + assert.NoError(t, err) + assert.Equal(t, tokenStr, authInfo.Token) + + // Expect at least two server calls: initial cache fetch + forced refresh + assert.GreaterOrEqual(t, callCount, 2) +} diff --git a/internal/auth/server/types.go b/internal/auth/server/types.go new file mode 100644 index 0000000..78015d2 --- /dev/null +++ b/internal/auth/server/types.go @@ -0,0 +1,40 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package server + +import "net/url" + +// AuthInfo holds information about a validated access token +// and is provided to request handlers +type AuthInfo struct { + // Token is the original access token string + Token string `json:"token"` + + // ClientID is the client identifier associated with this token + ClientID string `json:"clientId"` + + // Subject is the principal (end-user or client) the token represents + // Typically comes from the JWT 'sub' claim or introspection response + Subject string `json:"subject,omitempty"` + + // Scopes are the permission scopes granted with this token + Scopes []string `json:"scopes"` + + // ExpiresAt is the token expiration time in seconds since Unix epoch + // If nil it means no expiration was provided + ExpiresAt *int64 `json:"expiresAt,omitempty"` + + // Resource is the RFC 8707 resource server identifier for which this token is valid + // If set it must match the resource identifier of the MCP server excluding any fragment + // If nil it means no resource was provided + Resource *url.URL `json:"resource,omitempty"` + + // Extra contains any additional data attached to this token + // Used for passing custom claims or metadata alongside authentication info + // If nil it means no extra data was provided + Extra map[string]interface{} `json:"extra,omitempty"` +} diff --git a/internal/auth/types.go b/internal/auth/types.go new file mode 100644 index 0000000..b7b6d3d --- /dev/null +++ b/internal/auth/types.go @@ -0,0 +1,271 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package auth + +import "net/http" + +// OAuthClientMetadata defines RFC 7591 OAuth 2.0 Dynamic Client Registration metadata +type OAuthClientMetadata struct { + RedirectURIs []string `json:"redirect_uris"` // Allowed redirect URIs for the client + TokenEndpointAuthMethod string `json:"token_endpoint_auth_method,omitempty"` // Client auth method at token endpoint + GrantTypes []string `json:"grant_types,omitempty"` // Supported grant types + ResponseTypes []string `json:"response_types,omitempty"` // Supported response types + ClientName *string `json:"client_name,omitempty"` // Human readable client name + ClientURI *string `json:"client_uri,omitempty"` // Client homepage URL + LogoURI *string `json:"logo_uri,omitempty"` // Client logo URL + Scope *string `json:"scope,omitempty"` // Default requested scopes as space separated string + Contacts []string `json:"contacts,omitempty"` // Admin contact emails + TosURI *string `json:"tos_uri,omitempty"` // Terms of service URL + PolicyURI *string `json:"policy_uri,omitempty"` // Privacy policy URL + JwksURI *string `json:"jwks_uri,omitempty"` // URL to client JWKS + Jwks interface{} `json:"jwks,omitempty"` // Inline JWKS object + SoftwareID *string `json:"software_id,omitempty"` // Software identifier + SoftwareVersion *string `json:"software_version,omitempty"` // Software version + SoftwareStatement *string `json:"software_statement,omitempty"` // Software statement assertion +} + +// OAuthClientInformation defines RFC 7591 OAuth 2.0 Dynamic Client Registration client information +type OAuthClientInformation struct { + ClientID string `json:"client_id"` // Issued client identifier + ClientSecret string `json:"client_secret,omitempty"` // Issued client secret if applicable + ClientIDIssuedAt *int64 `json:"client_id_issued_at,omitempty"` // Issue time in seconds since epoch + ClientSecretExpiresAt *int64 `json:"client_secret_expires_at,omitempty"` // Secret expiry time in seconds since epoch +} + +// OAuthClientInformationFull defines RFC 7591 OAuth 2.0 Dynamic Client Registration full response +type OAuthClientInformationFull struct { + OAuthClientMetadata + OAuthClientInformation +} + +// OAuthProtectedResourceMetadata defines RFC 9728 OAuth Protected Resource metadata +type OAuthProtectedResourceMetadata struct { + Resource string `json:"resource"` // Resource identifier URI + AuthorizationServers []string `json:"authorization_servers,omitempty"` // Authorization server issuers supporting this resource + JWKSURI *string `json:"jwks_uri,omitempty"` // JWKS URI used by the resource + ScopesSupported []string `json:"scopes_supported,omitempty"` // Supported scopes + BearerMethodsSupported []string `json:"bearer_methods_supported,omitempty"` // Supported bearer presentation methods + ResourceSigningAlgs []string `json:"resource_signing_alg_values_supported,omitempty"` // Supported signing algorithms + ResourceName *string `json:"resource_name,omitempty"` // Human friendly resource name + ResourceDocumentation *string `json:"resource_documentation,omitempty"` // Documentation URL + ResourcePolicyURI *string `json:"resource_policy_uri,omitempty"` // Policy URL + ResourceTOSURI *string `json:"resource_tos_uri,omitempty"` // Terms of service URL + TLSCertBoundAT *bool `json:"tls_client_certificate_bound_access_tokens,omitempty"` // Whether MTLS bound AT are required or supported + AuthzDetailsTypes []string `json:"authorization_details_types_supported,omitempty"` // Supported authorization details types + DPoPSigningAlgs []string `json:"dpop_signing_alg_values_supported,omitempty"` // Supported DPoP signing algorithms + DPoPBoundATRequired *bool `json:"dpop_bound_access_tokens_required,omitempty"` // Whether DPoP bound access tokens are required +} + +// OAuthTokens defines the OAuth 2.1 token response +type OAuthTokens struct { + AccessToken string `json:"access_token"` // Access token value required non empty + IDToken *string `json:"id_token,omitempty"` // OIDC ID token optional non empty if present + TokenType string `json:"token_type"` // Token type for example Bearer required non empty + ExpiresIn *int64 `json:"expires_in,omitempty"` // Access token lifetime in seconds optional positive if present + Scope *string `json:"scope,omitempty"` // Granted scope as space separated string optional non empty if present + RefreshToken *string `json:"refresh_token,omitempty"` // Refresh token optional non empty if present +} + +// OAuthTokenRevocationRequest represents a token revocation request payload +type OAuthTokenRevocationRequest struct { + Token string `json:"token"` // Token to revoke + TokenTypeHint string `json:"token_type_hint,omitempty"` // Optional token type hint +} + +// AuthorizationServerMetadata represents OAuth 2.0 or OpenID Connect server metadata +type AuthorizationServerMetadata interface { + // GetIssuer returns the issuer identifier + GetIssuer() string + // GetAuthorizationEndpoint returns the authorization endpoint URL + GetAuthorizationEndpoint() string + // GetTokenEndpoint returns the token endpoint URL + GetTokenEndpoint() string + // GetResponseTypesSupported returns supported response types + GetResponseTypesSupported() []string + // GetGrantTypesSupported returns supported grant types + GetGrantTypesSupported() []string + // GetTokenEndpointAuthMethodsSupported returns supported client auth methods for the token endpoint + GetTokenEndpointAuthMethodsSupported() []string +} + +// OAuthMetadata defines OAuth 2.0 Authorization Server Metadata per RFC 8414 +type OAuthMetadata struct { + Issuer string `json:"issuer"` // Issuer identifier + AuthorizationEndpoint string `json:"authorization_endpoint"` // Authorization endpoint URL + TokenEndpoint string `json:"token_endpoint"` // Token endpoint URL + RegistrationEndpoint *string `json:"registration_endpoint,omitempty"` // Dynamic client registration endpoint + ScopesSupported []string `json:"scopes_supported,omitempty"` // Supported scopes + ResponseTypesSupported []string `json:"response_types_supported"` // Supported response types + ResponseModesSupported []string `json:"response_modes_supported,omitempty"` // Supported response modes + GrantTypesSupported []string `json:"grant_types_supported,omitempty"` // Supported grant types + TokenEndpointAuthMethodsSupported []string `json:"token_endpoint_auth_methods_supported,omitempty"` // Supported token endpoint auth methods + TokenEndpointAuthSigningAlgValuesSupported []string `json:"token_endpoint_auth_signing_alg_values_supported,omitempty"` // Supported signing algs for client auth + ServiceDocumentation *string `json:"service_documentation,omitempty"` // Service documentation URL + RevocationEndpoint *string `json:"revocation_endpoint,omitempty"` // Token revocation endpoint + RevocationEndpointAuthMethodsSupported []string `json:"revocation_endpoint_auth_methods_supported,omitempty"` // Supported auth methods for revocation + RevocationEndpointAuthSigningAlgValuesSupported []string `json:"revocation_endpoint_auth_signing_alg_values_supported,omitempty"` // Supported signing algs for revocation + IntrospectionEndpoint *string `json:"introspection_endpoint,omitempty"` // Token introspection endpoint + IntrospectionEndpointAuthMethodsSupported []string `json:"introspection_endpoint_auth_methods_supported,omitempty"` // Supported auth methods for introspection + IntrospectionEndpointAuthSigningAlgValuesSupported []string `json:"introspection_endpoint_auth_signing_alg_values_supported,omitempty"` // Supported signing algs for introspection + CodeChallengeMethodsSupported []string `json:"code_challenge_methods_supported,omitempty"` // Supported PKCE methods +} + +// GetIssuer returns the issuer identifier +func (m OAuthMetadata) GetIssuer() string { + return m.Issuer +} + +// GetAuthorizationEndpoint returns the authorization endpoint URL +func (m OAuthMetadata) GetAuthorizationEndpoint() string { + return m.AuthorizationEndpoint +} + +// GetTokenEndpoint returns the token endpoint URL +func (m OAuthMetadata) GetTokenEndpoint() string { + return m.TokenEndpoint +} + +// GetResponseTypesSupported returns supported response types +func (m OAuthMetadata) GetResponseTypesSupported() []string { + return m.ResponseTypesSupported +} + +// GetGrantTypesSupported returns supported grant types +func (m OAuthMetadata) GetGrantTypesSupported() []string { + return m.GrantTypesSupported +} + +// GetTokenEndpointAuthMethodsSupported returns supported client auth methods for the token endpoint +func (m OAuthMetadata) GetTokenEndpointAuthMethodsSupported() []string { + return m.TokenEndpointAuthMethodsSupported +} + +// OpenIdProviderMetadata defines OpenID Connect Discovery 1.0 provider metadata +type OpenIdProviderMetadata struct { + Issuer string `json:"issuer"` // Issuer identifier + AuthorizationEndpoint string `json:"authorization_endpoint"` // Authorization endpoint URL + TokenEndpoint string `json:"token_endpoint"` // Token endpoint URL + UserinfoEndpoint *string `json:"userinfo_endpoint,omitempty"` // Userinfo endpoint URL + JwksURI string `json:"jwks_uri"` // JWKS URI + RegistrationEndpoint *string `json:"registration_endpoint,omitempty"` // Dynamic client registration endpoint + ScopesSupported []string `json:"scopes_supported,omitempty"` // Supported scopes + ResponseTypesSupported []string `json:"response_types_supported"` // Supported response types + ResponseModesSupported []string `json:"response_modes_supported,omitempty"` // Supported response modes + GrantTypesSupported []string `json:"grant_types_supported,omitempty"` // Supported grant types + AcrValuesSupported []string `json:"acr_values_supported,omitempty"` // Supported ACR values + SubjectTypesSupported []string `json:"subject_types_supported"` // Supported subject types + IdTokenSigningAlgValuesSupported []string `json:"id_token_signing_alg_values_supported"` // Supported ID token signing algs + IdTokenEncryptionAlgValuesSupported []string `json:"id_token_encryption_alg_values_supported,omitempty"` // Supported ID token encryption algs + IdTokenEncryptionEncValuesSupported []string `json:"id_token_encryption_enc_values_supported,omitempty"` // Supported ID token encryption enc values + UserinfoSigningAlgValuesSupported []string `json:"userinfo_signing_alg_values_supported,omitempty"` // Supported userinfo signing algs + UserinfoEncryptionAlgValuesSupported []string `json:"userinfo_encryption_alg_values_supported,omitempty"` // Supported userinfo encryption algs + UserinfoEncryptionEncValuesSupported []string `json:"userinfo_encryption_enc_values_supported,omitempty"` // Supported userinfo encryption enc values + RequestObjectSigningAlgValuesSupported []string `json:"request_object_signing_alg_values_supported,omitempty"` // Supported request object signing algs + RequestObjectEncryptionAlgValuesSupported []string `json:"request_object_encryption_alg_values_supported,omitempty"` // Supported request object encryption algs + RequestObjectEncryptionEncValuesSupported []string `json:"request_object_encryption_enc_values_supported,omitempty"` // Supported request object encryption enc values + TokenEndpointAuthMethodsSupported []string `json:"token_endpoint_auth_methods_supported,omitempty"` // Supported token endpoint auth methods + TokenEndpointAuthSigningAlgValuesSupported []string `json:"token_endpoint_auth_signing_alg_values_supported,omitempty"` // Supported signing algs for token endpoint auth + DisplayValuesSupported []string `json:"display_values_supported,omitempty"` // Supported display values + ClaimTypesSupported []string `json:"claim_types_supported,omitempty"` // Supported claim types + ClaimsSupported []string `json:"claims_supported,omitempty"` // Supported claims + ServiceDocumentation *string `json:"service_documentation,omitempty"` // Service documentation URL + ClaimsLocalesSupported []string `json:"claims_locales_supported,omitempty"` // Supported claims locales + UiLocalesSupported []string `json:"ui_locales_supported,omitempty"` // Supported UI locales + ClaimsParameterSupported *bool `json:"claims_parameter_supported,omitempty"` // Whether claims parameter is supported + RequestParameterSupported *bool `json:"request_parameter_supported,omitempty"` // Whether request parameter is supported + RequestUriParameterSupported *bool `json:"request_uri_parameter_supported,omitempty"` // Whether request_uri is supported + RequireRequestUriRegistration *bool `json:"require_request_uri_registration,omitempty"` // Whether request_uri registration is required + OpPolicyUri *string `json:"op_policy_uri,omitempty"` // OP policy URL + OpTosUri *string `json:"op_tos_uri,omitempty"` // OP terms of service URL +} + +// AuthOptions contains configuration options for the OAuth authorization process +type AuthOptions struct { + ServerUrl string // OAuth server URL + ResourceMetadataUrl *string // Resource metadata URL + AuthorizationCode *string // Authorization code to exchange + Scope *string // Requested scopes as space separated string + ProtocolVersion *string // OAuth protocol version string + FetchFn FetchFunc // Custom HTTP request function +} + +// DiscoveryOptions contains options for discovering OAuth server metadata +type DiscoveryOptions struct { + ServerUrl string // Base server URL for discovery + ResourceMetadataUrl *string // Resource metadata URL for RFC 9728 + FetchFn FetchFunc // Custom HTTP request function + ProtocolVersion *string // Protocol version hint +} + +// GetIssuer returns the issuer identifier +func (m OpenIdProviderMetadata) GetIssuer() string { + return m.Issuer +} + +// GetAuthorizationEndpoint returns the authorization endpoint URL +func (m OpenIdProviderMetadata) GetAuthorizationEndpoint() string { + return m.AuthorizationEndpoint +} + +// GetTokenEndpoint returns the token endpoint URL +func (m OpenIdProviderMetadata) GetTokenEndpoint() string { + return m.TokenEndpoint +} + +// GetResponseTypesSupported returns supported response types +func (m OpenIdProviderMetadata) GetResponseTypesSupported() []string { + return m.ResponseTypesSupported +} + +// GetGrantTypesSupported returns supported grant types +func (m OpenIdProviderMetadata) GetGrantTypesSupported() []string { + return m.GrantTypesSupported +} + +// GetTokenEndpointAuthMethodsSupported returns supported client auth methods for the token endpoint +func (m OpenIdProviderMetadata) GetTokenEndpointAuthMethodsSupported() []string { + return m.TokenEndpointAuthMethodsSupported +} + +// OpenIdProviderDiscoveryMetadata merges OpenID Provider metadata with OAuth 2.0 fields for discovery +type OpenIdProviderDiscoveryMetadata struct { + OpenIdProviderMetadata // Embedded OIDC provider metadata + CodeChallengeMethodsSupported []string `json:"code_challenge_methods_supported,omitempty"` // Supported PKCE methods +} + +// GetIssuer returns the issuer identifier +func (m OpenIdProviderDiscoveryMetadata) GetIssuer() string { + return m.OpenIdProviderMetadata.Issuer +} + +// GetAuthorizationEndpoint returns the authorization endpoint URL +func (m OpenIdProviderDiscoveryMetadata) GetAuthorizationEndpoint() string { + return m.OpenIdProviderMetadata.AuthorizationEndpoint +} + +// GetTokenEndpoint returns the token endpoint URL +func (m OpenIdProviderDiscoveryMetadata) GetTokenEndpoint() string { + return m.OpenIdProviderMetadata.TokenEndpoint +} + +// GetResponseTypesSupported returns supported response types +func (m OpenIdProviderDiscoveryMetadata) GetResponseTypesSupported() []string { + return m.OpenIdProviderMetadata.ResponseTypesSupported +} + +// GetGrantTypesSupported returns supported grant types +func (m OpenIdProviderDiscoveryMetadata) GetGrantTypesSupported() []string { + return m.OpenIdProviderMetadata.GrantTypesSupported +} + +// GetTokenEndpointAuthMethodsSupported returns supported client auth methods for the token endpoint +func (m OpenIdProviderDiscoveryMetadata) GetTokenEndpointAuthMethodsSupported() []string { + return m.OpenIdProviderMetadata.TokenEndpointAuthMethodsSupported +} + +// FetchFunc is a customizable HTTP fetch function used by discovery and auth flows +type FetchFunc func(url string, req *http.Request) (*http.Response, error) diff --git a/internal/auth/utils.go b/internal/auth/utils.go new file mode 100644 index 0000000..01e5ac7 --- /dev/null +++ b/internal/auth/utils.go @@ -0,0 +1,94 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package auth + +import ( + "fmt" + "net/url" + "strings" +) + +// Utilities for handling OAuth resource URIs. + +// ResourceURLFromServerURL converts a server URL to a resource URL by removing the fragment. +func ResourceURLFromServerURL(u interface{}) (*url.URL, error) { + var resourceURL *url.URL + var err error + + switch v := u.(type) { + case string: + resourceURL, err = url.Parse(v) + if err != nil { + return nil, err + } + case *url.URL: + resourceURL, err = url.Parse(v.String()) + if err != nil { + return nil, err + } + default: + return nil, fmt.Errorf("unsupported URL type") + } + + // Remove fragment + resourceURL.Fragment = "" + return resourceURL, nil +} + +// CheckResourceAllowedParams represents the parameters for CheckResourceAllowed function +type CheckResourceAllowedParams struct { + RequestedResource interface{} // URL string or *url.URL + ConfiguredResource interface{} // URL string or *url.URL +} + +// CheckResourceAllowed checks if a requested resource URL matches a configured resource URL. +func CheckResourceAllowed(params CheckResourceAllowedParams) (bool, error) { + requested, err := parseURL(params.RequestedResource) + if err != nil { + return false, err + } + + configured, err := parseURL(params.ConfiguredResource) + if err != nil { + return false, err + } + + // Compare the origin (scheme, domain, and port) + if requested.Scheme != configured.Scheme || + requested.Host != configured.Host { + return false, nil + } + + // Handle cases like requested=/foo and configured=/foo/ + if len(requested.Path) < len(configured.Path) { + return false, nil + } + + requestedPath := requested.Path + if !strings.HasSuffix(requestedPath, "/") { + requestedPath = requestedPath + "/" + } + + configuredPath := configured.Path + if !strings.HasSuffix(configuredPath, "/") { + configuredPath = configuredPath + "/" + } + + return strings.HasPrefix(requestedPath, configuredPath), nil +} + +// parseURL is a helper function to parse URL from string or *url.URL +func parseURL(u interface{}) (*url.URL, error) { + switch v := u.(type) { + case string: + return url.Parse(v) + case *url.URL: + return url.Parse(v.String()) + default: + return nil, fmt.Errorf("unsupported URL type") + } +} diff --git a/internal/auth/utils_test.go b/internal/auth/utils_test.go new file mode 100644 index 0000000..1dac578 --- /dev/null +++ b/internal/auth/utils_test.go @@ -0,0 +1,228 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package auth + +import ( + "net/url" + "testing" +) + +func TestResourceURLFromServerURL_JSParity(t *testing.T) { + t.Run("remove fragments", func(t *testing.T) { + got, err := ResourceURLFromServerURL("https://example.com/path#fragment") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if got.String() != "https://example.com/path" { + t.Fatalf("got %q, want %q", got.String(), "https://example.com/path") + } + + got, err = ResourceURLFromServerURL("https://example.com#fragment") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if got.String() != "https://example.com" && got.String() != "https://example.com/" { + t.Fatalf("got %q, want https://example.com(/)", got.String()) + } + + got, err = ResourceURLFromServerURL("https://example.com/path?query=1#fragment") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if got.String() != "https://example.com/path?query=1" { + t.Fatalf("got %q, want %q", got.String(), "https://example.com/path?query=1") + } + }) + + t.Run("no fragment -> unchanged", func(t *testing.T) { + cases := []string{ + "https://example.com", + "https://example.com/path", + "https://example.com/path?query=1", + } + for _, in := range cases { + got, err := ResourceURLFromServerURL(in) + if err != nil { + t.Fatalf("unexpected error for %q: %v", in, err) + } + // url.String() may print with or without a trailing slash when the path is rooted; both are accepted. + want := in + if got.String() != want && !(want == "https://example.com" && got.String() == "https://example.com/") { + t.Fatalf("got %q, want %q", got.String(), want) + } + } + }) + + t.Run("keep everything else unchanged (except fragment)", func(t *testing.T) { + got, err := ResourceURLFromServerURL("https://example.com:443/path") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if got.String() != "https://example.com:443/path" { + t.Fatalf("got %q, want %q", got.String(), "https://example.com:443/path") + } + + got, err = ResourceURLFromServerURL("https://example.com:8080/path") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if got.String() != "https://example.com:8080/path" { + t.Fatalf("got %q, want %q", got.String(), "https://example.com:8080/path") + } + + got, err = ResourceURLFromServerURL("https://example.com/?foo=bar&baz=qux") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if got.String() != "https://example.com/?foo=bar&baz=qux" && got.String() != "https://example.com?foo=bar&baz=qux" { + t.Fatalf("got %q, want https://example.com/?foo=bar&baz=qux", got.String()) + } + + got, err = ResourceURLFromServerURL("https://example.com/") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if got.String() != "https://example.com/" { + t.Fatalf("got %q, want %q", got.String(), "https://example.com/") + } + + got, err = ResourceURLFromServerURL("https://example.com/path/") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if got.String() != "https://example.com/path/" { + t.Fatalf("got %q, want %q", got.String(), "https://example.com/path/") + } + }) + + t.Run("unsupported type -> error", func(t *testing.T) { + _, err := ResourceURLFromServerURL(123) + if err == nil { + t.Fatal("expected error for unsupported type") + } + }) +} + +func TestCheckResourceAllowed_JSParity(t *testing.T) { + type args struct { + req interface{} + cfg interface{} + } + tests := []struct { + name string + a args + want bool + wantErr bool + }{ + { + name: "identical URLs", + a: args{"https://example.com/path", "https://example.com/path"}, + want: true, + }, + { + name: "identical origins at root", + a: args{"https://example.com/", "https://example.com/"}, + want: true, + }, + { + name: "different paths -> false", + a: args{"https://example.com/path1", "https://example.com/path2"}, + want: false, + }, + { + name: "requested root vs configured path -> false", + a: args{"https://example.com/", "https://example.com/path"}, + want: false, + }, + { + name: "different domain -> false", + a: args{"https://example.com/path", "https://example.org/path"}, + want: false, + }, + { + name: "different port -> false", + a: args{"https://example.com:8080/path", "https://example.com/path"}, + want: false, + }, + { + name: "path prefix but not by segment (mcpxxxx vs mcp) -> false", + a: args{"https://example.com/mcpxxxx", "https://example.com/mcp"}, + want: false, + }, + { + name: "requested shorter than configured -> false", + a: args{"https://example.com/folder", "https://example.com/folder/subfolder"}, + want: false, + }, + { + name: "requested is subpath of configured -> true", + a: args{"https://example.com/api/v1", "https://example.com/api"}, + want: true, + }, + { + name: "trailing slash handling: requested has slash, configured no slash -> true", + a: args{"https://example.com/mcp/", "https://example.com/mcp"}, + want: true, + }, + { + name: "trailing slash handling: requested no slash, configured has slash -> false (requested shorter)", + a: args{"https://example.com/folder", "https://example.com/folder/"}, + want: false, + }, + { + name: "invalid requested URL -> error", + a: args{"https://%zz", "https://example.com/path"}, + wantErr: true, + }, + { + name: "invalid configured URL -> error", + a: args{"https://example.com/path", "://bad_url"}, + wantErr: true, + }, + { + name: "unsupported requested type -> error", + a: args{123, "https://example.com/path"}, + wantErr: true, + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + got, err := CheckResourceAllowed(CheckResourceAllowedParams{ + RequestedResource: tt.a.req, + ConfiguredResource: tt.a.cfg, + }) + if tt.wantErr { + if err == nil { + t.Fatalf("expected error, got nil (got=%v)", got) + } + return + } + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if got != tt.want { + t.Fatalf("got %v, want %v", got, tt.want) + } + }) + } +} + +func TestResourceURLFromServerURL_URLInput(t *testing.T) { + u, _ := url.Parse("https://example.com/A/B#frag") + got, err := ResourceURLFromServerURL(u) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if got.Fragment != "" { + t.Fatalf("fragment not removed: %q", got.Fragment) + } + if got.Path != "/A/B" { + t.Fatalf("path changed: %q", got.Path) + } +} diff --git a/internal/errors/auth.go b/internal/errors/auth.go new file mode 100644 index 0000000..8bf7eaa --- /dev/null +++ b/internal/errors/auth.go @@ -0,0 +1,97 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package errors + +import ( + "errors" +) + +// OAuthErrorCode represents an OAuth 2.1 error code +type OAuthErrorCode error + +// OAuthError represents a structured OAuth 2.1 error +type OAuthError struct { + ErrorCode string + Message string + ErrorURI string +} + +// OAuthErrorResponse represents the JSON response for OAuth errors +type OAuthErrorResponse struct { + Error string `json:"error"` + ErrorDescription string `json:"error_description,omitempty"` + ErrorURI string `json:"error_uri,omitempty"` +} + +// Standard OAuth error codes +var ( + ErrInvalidRequest OAuthErrorCode = errors.New("invalid_request") + ErrInvalidClient OAuthErrorCode = errors.New("invalid_client") + ErrInvalidGrant OAuthErrorCode = errors.New("invalid_grant") + ErrUnauthorizedClient OAuthErrorCode = errors.New("unauthorized_client") + ErrUnsupportedGrantType OAuthErrorCode = errors.New("unsupported_grant_type") + ErrInvalidScope OAuthErrorCode = errors.New("invalid_scope") + ErrAccessDenied OAuthErrorCode = errors.New("access_denied") + ErrServerError OAuthErrorCode = errors.New("server_error") + ErrTemporarilyUnavailable OAuthErrorCode = errors.New("temporarily_unavailable") + ErrUnsupportedResponseType OAuthErrorCode = errors.New("unsupported_response_type") + ErrUnsupportedTokenType OAuthErrorCode = errors.New("unsupported_token_type") + ErrInvalidToken OAuthErrorCode = errors.New("invalid_token") + ErrMethodNotAllowed OAuthErrorCode = errors.New("method_not_allowed") + ErrTooManyRequests OAuthErrorCode = errors.New("too_many_requests") + ErrInvalidClientMetadata OAuthErrorCode = errors.New("invalid_client_metadata") + ErrInsufficientScope OAuthErrorCode = errors.New("insufficient_scope") +) + +// OAuthErrorMapping maps error strings to their corresponding OAuthErrorCode +// This replaces the need for large switch statements when parsing error responses +var OAuthErrorMapping = map[string]OAuthErrorCode{ + "invalid_request": ErrInvalidRequest, + "invalid_client": ErrInvalidClient, + "invalid_grant": ErrInvalidGrant, + "unauthorized_client": ErrUnauthorizedClient, + "unsupported_grant_type": ErrUnsupportedGrantType, + "invalid_scope": ErrInvalidScope, + "access_denied": ErrAccessDenied, + "server_error": ErrServerError, + "temporarily_unavailable": ErrTemporarilyUnavailable, + "unsupported_response_type": ErrUnsupportedResponseType, + "unsupported_token_type": ErrUnsupportedTokenType, + "invalid_token": ErrInvalidToken, + "method_not_allowed": ErrMethodNotAllowed, + "too_many_requests": ErrTooManyRequests, + "invalid_client_metadata": ErrInvalidClientMetadata, + "insufficient_scope": ErrInsufficientScope, +} + +// NewOAuthError creates a new OAuthError +func NewOAuthError(errCode OAuthErrorCode, message string, uri string) OAuthError { + err := OAuthError{ + ErrorCode: errCode.Error(), + } + if uri != "" { + err.ErrorURI = uri + } + if message != "" { + err.Message = message + } + return err +} + +// ToResponseStruct converts OAuthError into OAuthErrorResponse for JSON encoding +func (o OAuthError) ToResponseStruct() *OAuthErrorResponse { + return &OAuthErrorResponse{ + Error: o.ErrorCode, + ErrorDescription: o.Message, + ErrorURI: o.ErrorURI, + } +} + +// Error implements the error interface +func (o OAuthError) Error() string { + return o.ErrorCode +} diff --git a/internal/errors/auth_test.go b/internal/errors/auth_test.go new file mode 100644 index 0000000..c771627 --- /dev/null +++ b/internal/errors/auth_test.go @@ -0,0 +1,48 @@ +// Tencent is pleased to support the open source community by making trpc-mcp-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-mcp-go is licensed under the Apache License Version 2.0. + +package errors_test + +import ( + "testing" + "trpc.group/trpc-go/trpc-mcp-go/internal/errors" +) + +func TestNewOAuthError(t *testing.T) { + err := errors.NewOAuthError(errors.ErrInvalidRequest, "missing parameter", "https://example.com/docs") + + if err.ErrorCode != "invalid request" { + t.Errorf("expected error code 'invalid request', got %s", err.ErrorCode) + } + if err.Message != "missing parameter" { + t.Errorf("expected message 'missing parameter', got %s", err.Message) + } + if err.ErrorURI != "https://example.com/docs" { + t.Errorf("expected URI 'https://example.com/docs', got %s", err.ErrorURI) + } +} + +func TestToResponseStruct(t *testing.T) { + err := errors.NewOAuthError(errors.ErrInvalidClient, "bad client id", "") + resp := err.ToResponseStruct() + + if resp.Error != "invalid client" { + t.Errorf("expected 'invalid client', got %s", resp.Error) + } + if resp.ErrorDescription != "bad client id" { + t.Errorf("expected description 'bad client id', got %s", resp.ErrorDescription) + } + if resp.ErrorURI != "" { + t.Errorf("expected empty URI, got %s", resp.ErrorURI) + } +} + +func TestErrorMethod(t *testing.T) { + err := errors.NewOAuthError(errors.ErrServerError, "internal failure", "") + if err.Error() != "server error" { + t.Errorf("expected 'server error', got %s", err.Error()) + } +} diff --git a/internal/errors/errors.go b/internal/errors/errors.go index b45aa1a..ef975c7 100644 --- a/internal/errors/errors.go +++ b/internal/errors/errors.go @@ -7,7 +7,9 @@ // Package mcperrors defines common error types and constants package errors -import "errors" +import ( + "errors" +) // Common errors var ( diff --git a/server.go b/server.go index 04c3115..5b761f4 100644 --- a/server.go +++ b/server.go @@ -12,8 +12,16 @@ import ( "errors" "fmt" "net/http" + "net/url" "sync" "sync/atomic" + + "golang.org/x/time/rate" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth/server" + sh "trpc.group/trpc-go/trpc-mcp-go/internal/auth/server/handler" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth/server/middleware" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth/server/router" ) // Common errors @@ -55,6 +63,110 @@ const ( // Multiple HTTPContextFunc will be executed in the order they are registered. type HTTPContextFunc func(ctx context.Context, r *http.Request) context.Context +// AuditConfig defines configuration for audit middleware +type AuditConfig struct { + // Whether to enable audit logging + Enabled bool + + // Audit level: none, basic, detailed, full + Level string + + // Whether to hash sensitive data + HashSensitiveData bool + + // Whether to include request/response body + IncludeRequestBody bool + IncludeResponseBody bool + + // Custom endpoint patterns to audit + EndpointPatterns []string + + // Patterns to exclude from auditing + ExcludePatterns []string + + // Custom metadata extractor function + MetadataExtractor func(*http.Request) map[string]interface{} + + // Custom risk assessor function + RiskAssessor func(map[string]interface{}) (string, []string) +} + +// BearerAuthConfig defines configuration for Bearer token authentication +type BearerAuthConfig struct { + Enabled bool + + // Required: Token validator + Verifier server.TokenVerifierInterface + + // Optional: List of required scopes + RequiredScopes []string + + // Optional: Write resource_metadata for WWW-Authenticate + ResourceMetadataURL *string + + // Optional: Restrict accepted issuer (extra authorization check) + Issuer string + + // Optional: Restrict accepted audiences/resources (extra authorization check) + Audience []string +} + +// OAuthRoutesConfig defines configuration for OAuth 2.1 server routes. +type OAuthRoutesConfig struct { + // OAuth server implementation + Provider server.OAuthServerProvider + + // Canonical issuer identifier (iss claim) + IssuerURL *url.URL + + // Root URL for OAuth endpoints + BaseURL *url.URL + + // Optional link to service documentation + ServiceDocumentationURL *url.URL + + // Supported OAuth scopes + ScopesSupported []string + + // Optional human-readable resource name + ResourceName *string + + // Rate limit for /authorize endpoint + AuthorizationRateLimit *rate.Limiter + + // Rate limit for /token endpoint + TokenRateLimit *rate.Limiter + + // Resolve client_id from refresh token + ResolveClientIDFromRT func(rt string) (string, bool) + + // Rate limit for dynamic client registration + RegistrationRateLimit *sh.RegisterRateLimitConfig + + // Rate limit for token revocation + RevocationRateLimit *sh.RevocationRateLimitConfig +} + +// OAuthMetadataConfig defines configuration for exposing OAuth server metadata. +type OAuthMetadataConfig struct { + // Core OAuth server metadata + OAuthMetadata OAuthMetadata + + // Optional resource server URL + ResourceServerURL *url.URL + + // Optional service documentation URL + ServiceDocumentationURL *url.URL + + // Scopes advertised in metadata + ScopesSupported []string + + // Optional human-readable resource name + ResourceName *string +} + +type OAuthMetadata = auth.OAuthMetadata + // serverConfig stores all server configuration options type serverConfig struct { // Basic configuration @@ -74,11 +186,20 @@ type serverConfig struct { // HTTP context functions for extracting information from HTTP requests httpContextFuncs []HTTPContextFunc + // Audit middleware configuration + auditConfig *AuditConfig + + // Bearer authenticate configuration + bearerAuth *BearerAuthConfig + // Tool list filter function toolListFilter ToolListFilter // Method name modifier for external customization. methodNameModifier MethodNameModifier + + // Route installers for adding extra endpoints (e.g. OAuth, metadata). + routerInstallers []func(*http.ServeMux) error } // ServerNotificationHandler defines a function that handles notifications on the server side. @@ -100,6 +221,7 @@ type Server struct { requestID atomic.Int64 // Request ID counter for generating unique request IDs. notificationHandlers map[string]ServerNotificationHandler // Map of notification handlers by method name. notificationMu sync.RWMutex // Mutex for notification handlers map. + rootHandler http.Handler // Server's top-level HTTP handler including the core MCP endpoint and any extra routes. } // NewServer creates a new MCP server @@ -113,6 +235,8 @@ func NewServer(name, version string, options ...ServerOption) *Server { postSSEEnabled: true, getSSEEnabled: true, notificationBufferSize: defaultNotificationBufferSize, + auditConfig: nil, + bearerAuth: nil, } // Create server with provided serverInfo @@ -155,7 +279,6 @@ func (s *Server) initComponents() { if s.config.methodNameModifier != nil { toolManager.withMethodNameModifier(s.config.methodNameModifier) } - // Only set tool list filter if not nil. if s.config.toolListFilter != nil { toolManager.withToolListFilter(s.config.toolListFilter) } @@ -165,7 +288,6 @@ func (s *Server) initComponents() { resourceManager := newResourceManager() s.resourceManager = resourceManager - // Create prompt manager. promptManager := newPromptManager() s.promptManager = promptManager @@ -207,12 +329,39 @@ func (s *Server) initComponents() { // Inject logger into httpServerHandler if provided. if s.logger != nil { - // This is the httpServerHandler option version. httpOptions = append(httpOptions, withServerTransportLogger(s.logger)) } + // Enable Bearer token auth middleware if configured + if s.config.bearerAuth != nil && s.config.bearerAuth.Enabled { + authWrap := convertToAuthMiddleware(s.config.bearerAuth) + httpOptions = append(httpOptions, withTransportAuthEnabled(authWrap)) + } + + // Enable audit logging if configured + if s.config.auditConfig != nil && s.config.auditConfig.Enabled { + auditOpts := convertToMiddlewareOptions(s.config.auditConfig) + httpOptions = append(httpOptions, withTransportAuditEnabled( + middleware.AuditMiddleware(auditOpts), + )) + } + // Create HTTP handler. s.httpHandler = newHTTPServerHandler(s.mcpHandler, s.config.path, httpOptions...) + + mux := http.NewServeMux() + mux.Handle(s.config.path+"/", s.httpHandler) + + // Install additional routes (OAuth, resource metadata, .well-known, etc.) + for _, install := range s.config.routerInstallers { + _ = install(mux) + } + + // Exposing mux externally + s.customServer = &http.Server{Addr: s.config.addr, Handler: mux} + + // Expose mux as the server's root handler + s.rootHandler = mux } // ServerOption server option function. @@ -271,6 +420,34 @@ func WithHTTPContextFunc(fn HTTPContextFunc) ServerOption { } } +// WithAudit enables audit logging for the server with the specified configuration. +// The audit middleware will log all HTTP requests and responses based on the configuration. +// +// Example: +// +// server := mcp.NewServer("my-server", "1.0", +// mcp.WithAudit(&mcp.AuditConfig{ +// Enabled: true, +// Level: "detailed", +// EndpointPatterns: []string{"/mcp/", "/oauth2/"}, +// HashSensitiveData: true, +// }), +// ) +func WithAudit(config *AuditConfig) ServerOption { + return func(s *Server) { + s.config.auditConfig = config + } +} + +// WithBearerAuth configures the server to use Bearer token authentication. +// The provided BearerAuthConfig specifies how tokens are verified, +// what scopes are required, and optionally, metadata for WWW-Authenticate responses. +func WithBearerAuth(config *BearerAuthConfig) ServerOption { + return func(s *Server) { + s.config.bearerAuth = config + } +} + // WithStatelessMode sets whether the server uses stateless mode // In stateless mode, the server won't generate session IDs and won't validate session IDs in client requests // Each request will use a temporary session, which is only valid during request processing @@ -313,10 +490,68 @@ func WithServerAddress(addr string) ServerOption { } } +// WithOAuthRoutes installs standard OAuth 2.1 endpoints into the server, +// such as /authorize, /token, /revoke, and /register, depending on the +// provided AuthRouterOptions and the provider's capabilities. +func WithOAuthRoutes(cfg OAuthRoutesConfig) ServerOption { + return withHTTPRoutes(func(mux *http.ServeMux) error { + base := cfg.BaseURL + if base == nil { + base = cfg.IssuerURL + } + + opts := router.AuthRouterOptions{ + Provider: cfg.Provider, + IssuerUrl: cfg.IssuerURL, + BaseUrl: base, + ServiceDocumentationUrl: cfg.ServiceDocumentationURL, + ScopesSupported: cfg.ScopesSupported, + ResourceName: cfg.ResourceName, + + AuthorizationOptions: &sh.AuthorizationHandlerOptions{ + Provider: cfg.Provider, + RateLimit: cfg.AuthorizationRateLimit, + }, + TokenOptions: &sh.TokenHandlerOptions{ + Provider: cfg.Provider, + RateLimit: cfg.TokenRateLimit, + }, + ClientRegistrationOptions: &sh.ClientRegistrationHandlerOptions{ + ClientsStore: cfg.Provider.ClientsStore(), + RateLimit: cfg.RegistrationRateLimit, + }, + RevocationOptions: &sh.RevocationHandlerOptions{ + Provider: cfg.Provider, + RateLimit: cfg.RevocationRateLimit, + }, + } + return router.McpAuthRouter(mux, opts) + }) +} + +// WithOAuthMetadata installs the .well-known OAuth metadata endpoints +// (e.g. /.well-known/oauth-authorization-server and +// /.well-known/oauth-protected-resource) into the server. +// The returned metadata is constructed from the given AuthMetadataOptions. +func WithOAuthMetadata(cfg OAuthMetadataConfig) ServerOption { + return withHTTPRoutes(func(mux *http.ServeMux) error { + opts := router.AuthMetadataOptions{ + OAuthMetadata: cfg.OAuthMetadata, + ResourceServerUrl: cfg.ResourceServerURL, + ServiceDocumentationUrl: cfg.ServiceDocumentationURL, + ScopesSupported: cfg.ScopesSupported, + ResourceName: cfg.ResourceName, + } + return router.McpAuthMetadataRouter(mux, opts) + }) +} + // Start starts the server func (s *Server) Start() error { if s.customServer != nil { - s.customServer.Handler = s.Handler() + if s.customServer.Handler == nil { + s.customServer.Handler = s.Handler() + } return s.customServer.ListenAndServe() } return http.ListenAndServe(s.config.addr, s.Handler()) @@ -570,9 +805,18 @@ func (s *Server) GetActiveSessions() ([]string, error) { return s.getActiveSessions() } -// Handler returns the http.Handler for the server. -// This can be used to integrate the MCP server into existing HTTP servers. +// Handler returns the top-level http.Handler exposed by the server. +// This handler always includes the core MCP endpoint (e.g., /mcp). +// Depending on the configured ServerOptions, it may also include +// additional routes such as OAuth endpoints and .well-known metadata. +// +// You can pass this directly to an http.Server, or mount it into +// an existing HTTP mux as the unified entry point for MCP and +// any configured auxiliary endpoints. func (s *Server) Handler() http.Handler { + if s.rootHandler == nil { + return s.rootHandler + } return s.httpHandler } @@ -657,3 +901,83 @@ func (s *Server) handleServerNotification(ctx context.Context, notification *JSO } return nil } + +// convertToMiddlewareOptions converts AuditConfig to AuditMiddlewareOptions +func convertToMiddlewareOptions(config *AuditConfig) *middleware.AuditMiddlewareOptions { + level := middleware.AuditLevelBasic + switch config.Level { + case "detailed": + level = middleware.AuditLevelDetailed + case "full": + level = middleware.AuditLevelFull + } + + return &middleware.AuditMiddlewareOptions{ + Level: level, + HashSensitiveData: config.HashSensitiveData, + IncludeRequestBody: config.IncludeRequestBody, + IncludeResponseBody: config.IncludeResponseBody, + EndpointPatterns: config.EndpointPatterns, + ExcludePatterns: config.ExcludePatterns, + MetadataExtractor: config.MetadataExtractor, + } +} + +// convertToAuthMiddleware converts BearerAuthConfig to an HTTP middleware wrapper using RequireBearerAuth +// It also adapts the auth info stored by middleware into the server-level context so downstream code can use server.GetAuthInfo +func convertToAuthMiddleware(config *BearerAuthConfig) func(http.Handler) http.Handler { + if config == nil || !config.Enabled || config.Verifier == nil { + return nil + } + + opts := middleware.BearerAuthMiddlewareOptions{ + Verifier: config.Verifier, + RequiredScopes: config.RequiredScopes, + ResourceMetadataURL: config.ResourceMetadataURL, + Issuer: config.Issuer, + Audience: config.Audience, + } + + bearer := middleware.RequireBearerAuth(opts) + + // Compose an adapter that maps middleware.AuthInfoKey -> server.WithAuthInfo + return func(next http.Handler) http.Handler { + // Wrap the downstream handler to translate context + adapter := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if v := r.Context().Value(middleware.AuthInfoKey); v != nil { + if ai, ok := v.(server.AuthInfo); ok { + // Avoid token transparent transmission + aiCopy := ai + aiCopy.Token = "" + r = r.WithContext(server.WithAuthInfo(r.Context(), &aiCopy)) + } + } + next.ServeHTTP(w, r) + }) + return bearer(adapter) + } +} + +// withTransportAuditEnabled enables audit logging by wrapping the handler with given middleware +func withTransportAuditEnabled(wrap func(http.Handler) http.Handler) func(*httpServerHandler) { + return func(h *httpServerHandler) { + h.auditEnabled = (wrap != nil) + h.auditWrap = wrap + } +} + +// withTransportAuthEnabled enables bearer authentication by wrapping the handler with given middleware +func withTransportAuthEnabled(wrap func(http.Handler) http.Handler) func(*httpServerHandler) { + return func(h *httpServerHandler) { + h.authEnabled = (wrap != nil) + h.authWrap = wrap + } +} + +// withHTTPRoutes registers a custom installer function that can +// attach additional HTTP routes to the server's root mux. +func withHTTPRoutes(install func(*http.ServeMux) error) ServerOption { + return func(s *Server) { + s.config.routerInstallers = append(s.config.routerInstallers, install) + } +} diff --git a/stdio_server.go b/stdio_server.go index 3891ffa..afba947 100644 --- a/stdio_server.go +++ b/stdio_server.go @@ -582,7 +582,7 @@ type stdioServerInternal struct { func (s *stdioServerInternal) HandleRequest(ctx context.Context, rawMessage json.RawMessage) (interface{}, error) { var request JSONRPCRequest if err := json.Unmarshal(rawMessage, &request); err != nil { - return newJSONRPCErrorResponse(nil, -32700, "Parse error", nil), nil + return newJSONRPCErrorResponse(nil, ErrCodeParse, "Parse error", nil), nil } s.parent.logger.Debugf("Handling request: %s (ID: %v)", request.Method, request.ID) @@ -611,11 +611,11 @@ func (s *stdioServerInternal) HandleRequest(ctx context.Context, rawMessage json case MethodPing: return s.handlePing(ctx, request) default: - return newJSONRPCErrorResponse(request.ID, -32601, "Method not found", nil), nil + return newJSONRPCErrorResponse(request.ID, ErrCodeMethodNotFound, "Method not found", nil), nil } if err != nil { - return newJSONRPCErrorResponse(request.ID, -32603, "Internal error", err.Error()), nil + return newJSONRPCErrorResponse(request.ID, ErrCodeInternal, "Internal error", err.Error()), nil } // Check if result is already a JSON-RPC response or error (has jsonrpc field). diff --git a/streamable_client.go b/streamable_client.go index 9aba314..ba58ea6 100644 --- a/streamable_client.go +++ b/streamable_client.go @@ -19,6 +19,8 @@ import ( "sync" "time" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth/client" "trpc.group/trpc-go/trpc-mcp-go/internal/httputil" "trpc.group/trpc-go/trpc-mcp-go/internal/retry" ) @@ -85,6 +87,9 @@ type streamableHTTPClientTransport struct { // Client reference for accessing rootsProvider. client *Client + + // OAuth client provider + oauthProvider client.OAuthClientProvider } // NotificationHandler is a handler for notifications. @@ -191,6 +196,13 @@ func withTransportHTTPReqHandlerOption(option HTTPReqHandlerOption) transportOpt } } +// withTransportOAuthProvider adds an option for OAuth client provider +func withTransportOAuthProvider(p client.OAuthClientProvider) transportOption { + return func(t *streamableHTTPClientTransport) { + t.oauthProvider = p + } +} + // start is a no-op for streamableHTTPClientTransport. func (t *streamableHTTPClientTransport) start(ctx context.Context) error { return nil @@ -241,6 +253,11 @@ func (t *streamableHTTPClientTransport) send( return nil, fmt.Errorf("%w: %v", ErrRequestSerialization, err) } + ctx, err = t.ensureAuth(ctx) + if err != nil { + return nil, fmt.Errorf("authentication failed: %w", err) + } + // Create HTTP request httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, t.serverURL.String(), bytes.NewReader(reqBytes)) if err != nil { @@ -250,6 +267,7 @@ func (t *streamableHTTPClientTransport) send( httpReq.URL.Path = t.path } + t.setBasicHeaders(httpReq) // Set request headers - accept both SSE and JSON responses httpReq.Header.Set(httputil.ContentTypeHeader, httputil.ContentTypeJSON) httpReq.Header.Set(httputil.AcceptHeader, httputil.ContentTypeJSON+", "+httputil.ContentTypeSSE) @@ -264,6 +282,13 @@ func (t *streamableHTTPClientTransport) send( httpReq.Header.Set(httputil.LastEventIDHeader, t.lastEventID) } + // Authorization from ctx + t.setAuthorizationHeader(ctx, httpReq) + + if authInfo, ok := client.GetAuthInfo(ctx); ok && authInfo != nil && authInfo.AccessToken != "" { + httpReq.Header.Set("Authorization", "Bearer "+authInfo.AccessToken) + } + // Add custom headers for key, values := range t.httpHeaders { for _, value := range values { @@ -277,6 +302,51 @@ func (t *streamableHTTPClientTransport) send( return nil, fmt.Errorf("%w: %v", ErrHTTPRequestFailed, err) } + // 204/400: try once with fresh auth (rebuild request inside) + if httpResp.StatusCode == http.StatusNoContent || httpResp.StatusCode == http.StatusBadRequest { + httpResp.Body.Close() + return t.retryWithFreshAuth(ctx, reqBytes, options) + } + + // 401/403: refresh/ensure auth, then REBUILD a new request and resend ONCE + if httpResp.StatusCode == http.StatusUnauthorized || httpResp.StatusCode == http.StatusForbidden { + httpResp.Body.Close() + + if _, err := t.ensureAuth(ctx); err == nil { + httpReq2, err := http.NewRequestWithContext(ctx, http.MethodPost, t.serverURL.String(), bytes.NewReader(reqBytes)) + if err != nil { + return nil, fmt.Errorf("%w: %v", ErrHTTPRequestCreation, err) + } + if len(t.path) != 0 { + httpReq2.URL.Path = t.path + } + + // headers again (same as first send) + t.setBasicHeaders(httpReq2) + httpReq2.Header.Set(httputil.ContentTypeHeader, httputil.ContentTypeJSON) + httpReq2.Header.Set(httputil.AcceptHeader, httputil.ContentTypeJSON+", "+httputil.ContentTypeSSE) + if t.sessionID != "" && !t.isStateless { + httpReq2.Header.Set(httputil.SessionIDHeader, t.sessionID) + } + if options != nil && options.lastEventID != "" { + httpReq2.Header.Set(httputil.LastEventIDHeader, options.lastEventID) + } else if t.lastEventID != "" { + httpReq2.Header.Set(httputil.LastEventIDHeader, t.lastEventID) + } + t.setAuthorizationHeader(ctx, httpReq2) // new token if refreshed + for k, values := range t.httpHeaders { + for _, value := range values { + httpReq2.Header.Add(k, value) + } + } + + httpResp, err = t.httpReqHandler.Handle(ctx, t.httpClient, httpReq2) + if err != nil { + return nil, fmt.Errorf("%w: %v", ErrHTTPRequestFailed, err) + } + } + } + // Handle session ID if sessionID := httpResp.Header.Get(httputil.SessionIDHeader); sessionID != "" { t.setSessionID(sessionID) @@ -634,6 +704,11 @@ func (t *streamableHTTPClientTransport) connectGetSSE(ctx context.Context) error return fmt.Errorf("cannot establish GET SSE connection: session ID is empty") } + ctx, err := t.ensureAuth(ctx) + if err != nil { + return fmt.Errorf("authentication failed: %w", err) + } + // Build GET request req, err := http.NewRequestWithContext(ctx, http.MethodGet, t.serverURL.String(), nil) if err != nil { @@ -650,6 +725,10 @@ func (t *streamableHTTPClientTransport) connectGetSSE(ctx context.Context) error req.Header.Set(httputil.LastEventIDHeader, t.lastEventID) } + if authInfo, ok := client.GetAuthInfo(ctx); ok && authInfo != nil && authInfo.AccessToken != "" { + req.Header.Set("Authorization", "Bearer "+authInfo.AccessToken) + } + // Add custom headers for key, values := range t.httpHeaders { for _, value := range values { @@ -867,6 +946,10 @@ func (t *streamableHTTPClientTransport) sendResponseToServer(response interface{ } } + if authInfo, ok := client.GetAuthInfo(ctx); ok && authInfo != nil && authInfo.AccessToken != "" { + httpReq.Header.Set("Authorization", "Bearer "+authInfo.AccessToken) + } + // Add session ID if available if t.sessionID != "" { httpReq.Header.Set(httputil.SessionIDHeader, t.sessionID) // Use correct MCP protocol header: Mcp-Session-Id. @@ -972,3 +1055,150 @@ func (t *streamableHTTPClientTransport) establishGetSSEConnection() { t.establishGetSSE() } + +// ensureAuth ensures that the current request context carries valid authentication information +func (t *streamableHTTPClientTransport) ensureAuth(ctx context.Context) (context.Context, error) { + if t.oauthProvider == nil { + return ctx, nil + } + + if info, ok := client.GetAuthInfo(ctx); ok && info != nil && !client.IsTokenExpired(info) { + return ctx, nil + } + + _, err := client.Auth(t.oauthProvider, auth.AuthOptions{ + ServerUrl: t.serverURL.String(), + }) + if err != nil { + return client.WithAuthErr(ctx, err), err + } + + tokens, terr := t.oauthProvider.Tokens() + if terr != nil { + return client.WithAuthErr(ctx, err), terr + } + if tokens == nil { + return client.WithAuthErr(ctx, fmt.Errorf("no tokens after auth")), fmt.Errorf("no tokens") + } + info := client.ConvertTokensToAuthInfo(tokens) + return client.WithAuthInfo(ctx, info), nil +} + +// setBasicHeaders sets the basic headers that are common to all HTTP requests +func (t *streamableHTTPClientTransport) setBasicHeaders(req *http.Request) { + // Set content type and accept headers + req.Header.Set(httputil.ContentTypeHeader, httputil.ContentTypeJSON) + req.Header.Set(httputil.AcceptHeader, httputil.ContentTypeJSON+", "+httputil.ContentTypeSSE) + + // Set session ID if available and not in stateless mode + if t.sessionID != "" && !t.isStateless { + req.Header.Set(httputil.SessionIDHeader, t.sessionID) + } + + // Set last event ID if available + if t.lastEventID != "" { + req.Header.Set(httputil.LastEventIDHeader, t.lastEventID) + } + + // Set path if specified + if len(t.path) != 0 { + req.URL.Path = t.path + } + + // Add custom headers + for key, values := range t.httpHeaders { + for _, value := range values { + req.Header.Add(key, value) + } + } +} + +// setAuthorizationHeader sets the Authorization header using context auth info +func (t *streamableHTTPClientTransport) setAuthorizationHeader(ctx context.Context, req *http.Request) { + if authInfo, ok := client.GetAuthInfo(ctx); ok && authInfo != nil && authInfo.AccessToken != "" { + req.Header.Set("Authorization", "Bearer "+authInfo.AccessToken) + } +} + +// retryWithFreshAuth handles retry logic with fresh authentication +func (t *streamableHTTPClientTransport) retryWithFreshAuth(ctx context.Context, reqBytes []byte, options *streamOptions) (*json.RawMessage, error) { + ctx, err := t.ensureAuth(ctx) + if err != nil { + return nil, fmt.Errorf("re-authentication failed: %w", err) + } + + // Create a new HTTP request + httpReq2, err := http.NewRequestWithContext(ctx, http.MethodPost, t.serverURL.String(), bytes.NewReader(reqBytes)) + if err != nil { + return nil, fmt.Errorf("%w: %v", ErrHTTPRequestCreation, err) + } + + // Reset all headers + t.setBasicHeaders(httpReq2) + + // Set up Authorization using the new context + t.setAuthorizationHeader(ctx, httpReq2) + + // Process options specific to this request + if options != nil && options.lastEventID != "" { + httpReq2.Header.Set(httputil.LastEventIDHeader, options.lastEventID) + } + + // Send a retry request + httpResp, err := t.httpReqHandler.Handle(ctx, t.httpClient, httpReq2) + if err != nil { + return nil, fmt.Errorf("%w: %v", ErrHTTPRequestFailed, err) + } + defer httpResp.Body.Close() + + // Handle session IDs + if sessionID := httpResp.Header.Get(httputil.SessionIDHeader); sessionID != "" { + t.setSessionID(sessionID) + t.isStateless = false + } + + // Check the content type + contentType := httpResp.Header.Get(httputil.ContentTypeHeader) + if strings.Contains(contentType, httputil.ContentTypeSSE) { + // Handle SSE Responses, reqID is set to nil because this is a retry + return t.handleSSEResponse(ctx, httpResp, nil, options) + } + + // Check status code + if httpResp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("%w: status code %d", ErrHTTPRequestFailed, httpResp.StatusCode) + } + + // Read the response body + respBytes, err := io.ReadAll(httpResp.Body) + if err != nil { + return nil, fmt.Errorf("failed to read response body: %w", err) + } + + // Parse JSON response + var jsonResp map[string]interface{} + if err := json.Unmarshal(respBytes, &jsonResp); err != nil { + return nil, fmt.Errorf("%w: %v", ErrResponseParsing, err) + } + + // Check if it is an error response + if _, hasError := jsonResp["error"]; hasError { + rawMessage := json.RawMessage(respBytes) + return &rawMessage, nil + } + + // Extraction results section + resultData, ok := jsonResp["result"] + if !ok { + return nil, ErrMissingResultField + } + + // Serialized result is JSON + resultBytes, err := json.Marshal(resultData) + if err != nil { + return nil, fmt.Errorf("%w: %v", ErrResponseSerialization, err) + } + + rawMessage := json.RawMessage(resultBytes) + return &rawMessage, nil +} diff --git a/streamable_server.go b/streamable_server.go index 450041a..48637b3 100644 --- a/streamable_server.go +++ b/streamable_server.go @@ -11,10 +11,12 @@ import ( "encoding/json" "fmt" "net/http" + "strings" "sync" "sync/atomic" "time" + "trpc.group/trpc-go/trpc-mcp-go/internal/auth/server" "trpc.group/trpc-go/trpc-mcp-go/internal/httputil" "trpc.group/trpc-go/trpc-mcp-go/internal/sseutil" ) @@ -73,6 +75,18 @@ type httpServerHandler struct { // Response manager for server-to-client requests. responseManager *responseManager + + // Enable bearer auth middleware if true + authEnabled bool + + // Auth middleware wrapper applied if authEnabled + authWrap func(http.Handler) http.Handler + + // Enable audit logging if true + auditEnabled bool + + // Audit middleware wrapper applied if auditEnabled + auditWrap func(http.Handler) http.Handler } // getSSEConnection represents a GET SSE connection @@ -90,6 +104,8 @@ type getSSEConnection struct { sseResponder *sseResponder } +// ServerAuthConfig and NewAuthHTTPContextFunc were removed (replaced by RequireBearerAuth middleware) + // newHTTPServerHandler creates an HTTP server handler func newHTTPServerHandler(handler requestHandler, serverPath string, options ...func(*httpServerHandler)) *httpServerHandler { h := &httpServerHandler{ @@ -195,29 +211,49 @@ func withTransportHTTPContextFuncs(funcs []HTTPContextFunc) func(*httpServerHand // ServeHTTP implements the http.Handler interface func (h *httpServerHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { - if !h.isValidPath(r.URL.Path) { - if h.serverPath == "" { - http.Error(w, fmt.Sprintf("Path not found: %s (expected: %s)", r.URL.Path, h.serverPath), http.StatusNotFound) + if len(h.httpContextFuncs) > 0 { + enriched := r.Context() + for _, fn := range h.httpContextFuncs { + enriched = fn(enriched, r) } - return + r = r.WithContext(enriched) } - switch r.Method { - case http.MethodPost: - h.handlePost(r.Context(), w, r) - case http.MethodGet: - if !h.enableGetSSE { - w.Header().Set("Allow", "POST, DELETE") - http.Error(w, "GET method not enabled", http.StatusMethodNotAllowed) + var core http.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !h.isValidPath(r.URL.Path) { + if h.serverPath == "" { + http.Error(w, fmt.Sprintf("Path not found: %s (expected: %s)", r.URL.Path, h.serverPath), http.StatusNotFound) + } return } - h.handleGet(r.Context(), w, r) - case http.MethodDelete: - h.handleDelete(r.Context(), w, r) - default: - w.Header().Set("Allow", "POST, GET, DELETE") - http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) + + switch r.Method { + case http.MethodPost: + h.handlePost(r.Context(), w, r) + case http.MethodGet: + if !h.enableGetSSE { + w.Header().Set("Allow", "POST, DELETE") + http.Error(w, "GET method not enabled", http.StatusMethodNotAllowed) + return + } + h.handleGet(r.Context(), w, r) + case http.MethodDelete: + h.handleDelete(r.Context(), w, r) + default: + w.Header().Set("Allow", "POST, GET, DELETE") + http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) + } + }) + + // Apply auth middleware first, then audit middleware + if h.auditEnabled && h.auditWrap != nil { + core = h.auditWrap(core) + } + if h.authEnabled && h.authWrap != nil { + core = h.authWrap(core) } + + core.ServeHTTP(w, r) } type baseMessage struct { @@ -338,6 +374,9 @@ func (h *httpServerHandler) handlePostRequest(ctx context.Context, w http.Respon if session != nil { reqCtx = setSessionToContext(reqCtx, session) } + if authInfo, ok := server.GetAuthInfo(ctx); ok { + reqCtx = server.WithAuthInfo(reqCtx, authInfo) + } resp, err := h.requestHandler.handleRequest(reqCtx, &req, session) if err != nil { h.logger.Infof("Request processing failed: %v", err) @@ -365,6 +404,9 @@ func (h *httpServerHandler) handlePostRequest(ctx context.Context, w http.Respon if session != nil { reqCtx = setSessionToContext(reqCtx, session) } + if authInfo, ok := server.GetAuthInfo(ctx); ok { + reqCtx = server.WithAuthInfo(reqCtx, authInfo) + } resp, err := h.requestHandler.handleRequest(reqCtx, &req, session) if err != nil { h.logger.Infof("Request processing failed: %v", err) @@ -556,6 +598,12 @@ func (h *httpServerHandler) handleGet(ctx context.Context, w http.ResponseWriter return } + // Perform authentication context checks, + enrichedCtx := ctx + for _, fn := range h.httpContextFuncs { + enrichedCtx = fn(enrichedCtx, r) + } + // Check if streaming is supported flusher, ok := w.(http.Flusher) if !ok { @@ -570,7 +618,7 @@ func (h *httpServerHandler) handleGet(ctx context.Context, w http.ResponseWriter flusher.Flush() // Create context, for canceling connection - connCtx, cancelConn := context.WithCancel(ctx) + connCtx, cancelConn := context.WithCancel(enrichedCtx) localCancelFunc = cancelConn // Assign to the variable captured by defer // Check if there's already a GET SSE connection @@ -775,7 +823,10 @@ func (h *httpServerHandler) isValidPath(requestPath string) bool { if h.serverPath == "" { return true } - return requestPath == h.serverPath + sp := strings.TrimSuffix(h.serverPath, "/") + rp := strings.TrimSuffix(requestPath, "/") + + return rp == sp || strings.HasPrefix(requestPath, sp+"/") } // responseManager manages pending requests and their response channels.