feat: Implement OAuth in the client #296
Conversation
WalkthroughOAuth 2.0 client support is introduced to the MCP Go client, including PKCE, token management, and dynamic client registration. New utility functions, types, and error handling mechanisms are provided for OAuth integration. The transport layer is extended to inject OAuth tokens into HTTP requests and handle authorization errors. Comprehensive tests and an example usage with documentation are added. Changes
Assessment against linked issues
Suggested labels
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 golangci-lint (1.64.8)Error: you are using a configuration file for golangci-lint v2 with golangci-lint v1: please use golangci-lint v2 📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
✨ Finishing Touches
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
ezynda3
changed the title
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 10
🧹 Nitpick comments (14)
examples/oauth_client/README.md (1)
43-50: Include warning about using example URLsThe README uses example URLs that need to be replaced. Consider adding a prominent warning to emphasize that these are placeholder values that must be replaced with actual endpoint URLs.
```go const ( // Replace with your MCP server URL serverURL = "https://api.example.com/v1/mcp" // Use a localhost redirect URI for this example redirectURI = "http://localhost:8085/oauth/callback" ) + +// IMPORTANT: The URLs above are examples only. You must replace them with your actual MCP server URL +// and configure your OAuth server to accept the redirect URI you specify.</blockquote></details> <details> <summary>client/transport/oauth_test.go (1)</summary><blockquote> `51-92`: **Consider adding edge cases to memory token store test** While the basic functionality of the memory token store is tested, consider adding tests for: 1. Overwriting an existing token 2. Handling nil token values 3. Testing token expiration integration These additional tests would provide more comprehensive coverage of potential edge cases. </blockquote></details> <details> <summary>client/transport/streamable_http_oauth_test.go (2)</summary><blockquote> `24-30`: **Check `json.Encoder.Encode` return value in test server handler** The Encode call can fail (e.g., if the connection is closed) and the error is silently ignored. Because the test server controls the response, capturing failures helps surface unexpected regressions earlier. ```diff - json.NewEncoder(w).Encode(map[string]any{ + if err := json.NewEncoder(w).Encode(map[string]any{ "jsonrpc": "2.0", "id": 1, "result": "success", - }) + }); err != nil { + t.Errorf("handler Encode failed: %v", err) + }🧰 Tools
🪛 GitHub Check: lint
[failure] 25-25:
Error return value of(*encoding/json.Encoder).Encodeis not checked (errcheck)
33-43: HandletokenStore.SaveTokenerror
SaveTokenreturns an error (e.g., disk-full for file-based stores). Swallowing it in tests hides real issues and causes the linter failure you’re seeing.- tokenStore.SaveToken(validToken) + if err := tokenStore.SaveToken(validToken); err != nil { + t.Fatalf("SaveToken failed: %v", err) + }🧰 Tools
🪛 GitHub Check: lint
[failure] 42-42:
Error return value oftokenStore.SaveTokenis not checked (errcheck)client/oauth_test.go (4)
25-40: Check Encode error in test serverSame issue as in the transport test: ignore of
json.Encoder.Encodetriggerserrchecklint failure.- json.NewEncoder(w).Encode(map[string]any{ + if err := json.NewEncoder(w).Encode(map[string]any{ ... - }) + }); err != nil { + t.Errorf("handler Encode failed: %v", err) + }🧰 Tools
🪛 GitHub Check: lint
[failure] 28-28:
Error return value of(*encoding/json.Encoder).Encodeis not checked (errcheck)🪛 GitHub Actions: golangci-lint
[error] 28-28: golangci-lint: Error return value of
(*encoding/json.Encoder).Encodeis not checked (errcheck)
43-53: HandleSaveTokenerror- tokenStore.SaveToken(validToken) + if err := tokenStore.SaveToken(validToken); err != nil { + t.Fatalf("SaveToken failed: %v", err) + }🧰 Tools
🪛 GitHub Check: lint
[failure] 52-52:
Error return value oftokenStore.SaveTokenis not checked (errcheck)
63-91: Test does not confirm Authorization header propagation
NewOAuthStreamableHttpClientis started, but no RPC/notification is sent. Unless the higher-level client automatically performs aninitializecall (which this test does not assert), the server never receives a request; thus the OAuth path is untested.
Consider invoking a dummy RPC to ensure thatAuthorizationis attached and 401 handling works.
93-123: Minor: table-driven tests improve readability for error-helper utilitiesYou can collapse the assertions for
IsOAuthAuthorizationRequiredError/GetOAuthHandlerinto a table-driven loop, reducing duplication.client/transport/streamable_http.go (2)
459-467: Public API: consider namingGetOAuthHandler/IsOAuthEnablederror semanticsThe getters are convenient, but exposing the handler directly leaks implementation details (e.g., token store internals) and allows callers to mutate internal state.
If read-only access is intended, document it explicitly or expose narrower functionality (e.g., methods to start auth flow). Otherwise, keep as is.
430-437: Remember to drain/close body on 401 branchWhen returning early on 401 you skip the deferred
resp.Body.Close()(it’s after the return). Addio.Copy(io.Discard, resp.Body)+resp.Body.Close()before returning to avoid leaking connections.- if resp.StatusCode == http.StatusUnauthorized && c.oauthHandler != nil { - return &OAuthAuthorizationRequiredError{ - Handler: c.oauthHandler, - } + if resp.StatusCode == http.StatusUnauthorized && c.oauthHandler != nil { + io.Copy(io.Discard, resp.Body) + resp.Body.Close() + return &OAuthAuthorizationRequiredError{Handler: c.oauthHandler} }examples/oauth_client/main.go (1)
205-214: Surface unsupported-OS error fromopenBrowserWhen
runtime.GOOSis not recognised we silently replace the error with our own, but never return it.
Consider returning the error so the caller can decide how to proceed (e.g. print the URL once, then exit).-func openBrowser(url string) { +func openBrowser(url string) error { var err error @@ } - if err != nil { - log.Printf("Failed to open browser: %v", err) - fmt.Printf("Please open the following URL in your browser: %s\n", url) - } + if err != nil { + // Caller can log or decide to fall back. + return fmt.Errorf("openBrowser: %w", err) + } + return nil }Call-sites can treat a non-nil error the same way you already do.
This keeps the helper composable and testable.client/oauth.go (1)
26-34: DuplicateWithOAuthoptions can lead to confusing behaviour
NewOAuthStreamableHttpClientalways appends its owntransport.WithOAuth(oauthConfig)to the supplied options.
If the caller already provided aWithOAuthoption (perhaps with different scopes or a different token store) the last option silently wins.Consider guarding against duplicates or, at minimum, documenting that callers must not pass an explicit OAuth option.
client/transport/oauth.go (2)
61-66: Add safety margin when checking token expiryNetwork latency and clock-skew can cause a token to be considered valid locally but rejected by the server a few seconds later.
Add a grace period (e.g. 30 s) when evaluatingIsExpired.-return time.Now().After(t.ExpiresAt) +const skew = 30 * time.Second +return time.Now().Add(skew).After(t.ExpiresAt)
386-447: State parameter is accepted but never verified
ProcessAuthorizationResponsereceivesstateyet never uses it.
Although the example verifies the state before calling this method, the
handler itself should still compare the server-echoedstateto the one it
generated to prevent CSRF when used directly.At a minimum document this expectation; preferably add the check and return an
error on mismatch.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
Cache: Disabled due to data retention organization setting
Knowledge Base: Disabled due to data retention organization setting
📒 Files selected for processing (10)
client/oauth.go(1 hunks)client/oauth_test.go(1 hunks)client/transport/oauth.go(1 hunks)client/transport/oauth_test.go(1 hunks)client/transport/oauth_utils.go(1 hunks)client/transport/oauth_utils_test.go(1 hunks)client/transport/streamable_http.go(8 hunks)client/transport/streamable_http_oauth_test.go(1 hunks)examples/oauth_client/README.md(1 hunks)examples/oauth_client/main.go(1 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (6)
client/transport/oauth_utils_test.go (2)
client/transport/oauth_utils.go (4)
GenerateRandomString(11-17)GenerateCodeVerifier(20-23)GenerateCodeChallenge(26-31)GenerateState(34-36)client/oauth.go (3)
GenerateCodeVerifier(38-38)GenerateCodeChallenge(41-41)GenerateState(44-44)
client/transport/streamable_http.go (2)
client/oauth.go (3)
OAuthConfig(10-10)OAuthAuthorizationRequiredError(47-47)GetOAuthHandler(56-61)client/transport/oauth.go (3)
OAuthConfig(18-34)NewOAuthHandler(120-129)OAuthHandler(111-117)
client/transport/oauth_utils.go (1)
client/oauth.go (3)
GenerateCodeVerifier(38-38)GenerateCodeChallenge(41-41)GenerateState(44-44)
client/transport/oauth_test.go (3)
client/oauth.go (2)
Token(13-13)NewMemoryTokenStore(22-22)client/transport/oauth.go (2)
Token(45-58)NewMemoryTokenStore(75-77)client/transport/oauth_utils.go (1)
ValidateRedirectURI(39-56)
client/oauth.go (4)
client/transport/oauth.go (6)
OAuthConfig(18-34)Token(45-58)TokenStore(37-42)MemoryTokenStore(69-72)NewMemoryTokenStore(75-77)OAuthHandler(111-117)client/transport/streamable_http.go (4)
StreamableHTTPCOption(21-21)WithOAuth(37-41)NewStreamableHTTP(76-95)OAuthAuthorizationRequiredError(145-147)client/client.go (2)
Client(16-25)NewClient(44-54)client/transport/oauth_utils.go (3)
GenerateCodeVerifier(20-23)GenerateCodeChallenge(26-31)GenerateState(34-36)
client/transport/oauth.go (2)
client/oauth.go (5)
TokenStore(16-16)Token(13-13)NewMemoryTokenStore(22-22)MemoryTokenStore(19-19)OAuthConfig(10-10)client/client.go (1)
Client(16-25)
🪛 GitHub Check: lint
client/transport/streamable_http_oauth_test.go
[failure] 25-25:
Error return value of (*encoding/json.Encoder).Encode is not checked (errcheck)
[failure] 42-42:
Error return value of tokenStore.SaveToken is not checked (errcheck)
client/oauth_test.go
[failure] 28-28:
Error return value of (*encoding/json.Encoder).Encode is not checked (errcheck)
[failure] 52-52:
Error return value of tokenStore.SaveToken is not checked (errcheck)
examples/oauth_client/main.go
[failure] 181-181:
Error return value of w.Write is not checked (errcheck)
🪛 GitHub Actions: golangci-lint
client/oauth_test.go
[error] 28-28: golangci-lint: Error return value of (*encoding/json.Encoder).Encode is not checked (errcheck)
🔇 Additional comments (8)
client/transport/oauth_utils_test.go (3)
8-31: Well-implemented comprehensive test for random string generationThe test thoroughly verifies the
GenerateRandomStringfunction by checking different string lengths and ensuring randomness through non-repetition of generated values.
33-69: Thorough PKCE verification testsThe test for code verifier and challenge generation is comprehensive, covering expected length validation, randomness verification, and challenge consistency for the same verifier. This ensures the PKCE implementation is robust and compliant with OAuth standards.
71-88: Appropriate state parameter testingThe test effectively verifies that state parameters have the expected length and are generated with sufficient randomness to prevent CSRF attacks.
examples/oauth_client/README.md (1)
3-38: Well-structured OAuth integration documentationThe README provides excellent documentation on the OAuth client capabilities, usage, and workflow. It gives users a clear understanding of how the OAuth flow works with the MCP client.
client/transport/oauth_test.go (2)
8-49: Well-structured token expiration testThe test cases for token expiration cover all relevant scenarios (valid, expired, and no expiration) using a table-driven testing approach.
94-139: Comprehensive redirect URI validation testsThe test cases thoroughly verify redirect URI validation rules, covering all security-critical scenarios including HTTPS requirement, localhost exception, and invalid schemes.
client/transport/oauth_utils.go (2)
11-17: Secure random string generation approachThe implementation correctly uses crypto/rand for secure random generation and base64 URL encoding for string conversion, which is appropriate for security-sensitive OAuth parameters.
20-36: Compliant PKCE and state parameter implementationThe code verifier, challenge, and state generation functions follow OAuth 2.0 and PKCE (RFC 7636) specifications with appropriate string lengths and cryptographic methods.
There was a problem hiding this comment.
Actionable comments posted: 3
♻️ Duplicate comments (1)
client/transport/oauth.go (1)
287-305:⚠️ Potential issueValidate
authBaseURLto prevent returning relative pathsThe
getDefaultEndpointsfunction still doesn't validate thatauthBaseURLhas a scheme and host after clearing the path. This could result in returning relative endpoint paths that will fail when used withhttp.NewRequest.Apply this fix to validate the URL:
func (h *OAuthHandler) getDefaultEndpoints(baseURL string) *AuthServerMetadata { // Parse the base URL to extract the authority parsedURL, err := url.Parse(baseURL) if err != nil { return nil } // Discard any path component to get the authorization base URL parsedURL.Path = "" authBaseURL := parsedURL.String() + + // Ensure we have a valid URL with scheme and host + if parsedURL.Scheme == "" || parsedURL.Host == "" { + return nil + } return &AuthServerMetadata{ Issuer: authBaseURL, AuthorizationEndpoint: authBaseURL + "/authorize", TokenEndpoint: authBaseURL + "/token", RegistrationEndpoint: authBaseURL + "/register", } }
🧹 Nitpick comments (3)
client/transport/oauth.go (3)
36-42: Consider enhancing TokenStore interface for token revocationThe
TokenStoreinterface could benefit from methods for deleting or invalidating tokens, which would be useful for implementing logout or token revocation features.// TokenStore is an interface for storing and retrieving OAuth tokens type TokenStore interface { // GetToken returns the current token GetToken() (*Token, error) // SaveToken saves a token SaveToken(token *Token) error + // DeleteToken removes the stored token + // Useful for implementing logout or token revocation + DeleteToken() error }This would require updating the
MemoryTokenStoreimplementation accordingly.
125-128: Make HTTP timeout configurableThe HTTP client timeout is hardcoded to 30 seconds, which might not be appropriate for all environments or operations.
type OAuthConfig struct { // ... existing fields + // HTTPTimeout is the timeout for HTTP requests (defaults to 30 seconds) + HTTPTimeout time.Duration } func NewOAuthHandler(config OAuthConfig) *OAuthHandler { if config.TokenStore == nil { config.TokenStore = NewMemoryTokenStore() } + + timeout := 30 * time.Second + if config.HTTPTimeout > 0 { + timeout = config.HTTPTimeout + } return &OAuthHandler{ config: config, - httpClient: &http.Client{Timeout: 30 * time.Second}, + httpClient: &http.Client{Timeout: timeout}, } }
443-466: Consider validating the redirect URI in GetAuthorizationURLWhile the function correctly adds the redirect URI to the authorization URL, it doesn't validate it. It might be worth adding a validation to ensure the redirect URI is properly URL-encoded and has a valid format.
func (h *OAuthHandler) GetAuthorizationURL(ctx context.Context, state, codeChallenge string) (string, error) { metadata, err := h.getServerMetadata(ctx) if err != nil { return "", fmt.Errorf("failed to get server metadata: %w", err) } + // Validate redirect URI + if _, err := url.Parse(h.config.RedirectURI); err != nil { + return "", fmt.Errorf("invalid redirect URI: %w", err) + } params := url.Values{} params.Set("response_type", "code") params.Set("client_id", h.config.ClientID) params.Set("redirect_uri", h.config.RedirectURI) params.Set("state", state)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
Cache: Disabled due to data retention organization setting
Knowledge Base: Disabled due to data retention organization setting
📒 Files selected for processing (5)
client/transport/oauth.go(1 hunks)client/transport/oauth_test.go(1 hunks)client/transport/oauth_utils.go(1 hunks)client/transport/streamable_http.go(9 hunks)client/transport/streamable_http_oauth_test.go(1 hunks)
✅ Files skipped from review due to trivial changes (1)
- client/transport/oauth_utils.go
🚧 Files skipped from review as they are similar to previous changes (3)
- client/transport/oauth_test.go
- client/transport/streamable_http.go
- client/transport/streamable_http_oauth_test.go
🧰 Additional context used
🧬 Code Graph Analysis (1)
client/transport/oauth.go (3)
client/oauth.go (5)
TokenStore(16-16)Token(13-13)NewMemoryTokenStore(22-22)MemoryTokenStore(19-19)OAuthConfig(10-10)client/client.go (1)
Client(16-25)client/transport/streamable_http.go (1)
ErrOAuthAuthorizationRequired(146-146)
🔇 Additional comments (5)
client/transport/oauth.go (5)
140-145: Looks good - Fixed empty token checkThe code now properly checks if the access token is empty before considering it valid, addressing the previous issue where empty access tokens would be treated as valid.
232-242: Improved metadata URL validationThe code now properly checks if
AuthServerMetadataURLis empty and returns a clear error message instead of attempting to parse an empty string.
68-95: Memory store implementation looks goodThe
MemoryTokenStoreimplementation properly uses a mutex for thread safety and has appropriate error handling.
60-66: Token expiration check is well implementedThe
IsExpiredmethod correctly handles the case whereExpiresAtis zero and properly compares with the current time.
307-378: The dynamic client registration looks goodThe
RegisterClientmethod properly implements OAuth dynamic client registration, handling both public and confidential clients correctly.
pottekkat
added
type: enhancement
|
@ezynda3 -- modulo the couple of comments above, this PR looks good. Could you please help with getting this merged? Thanks. |
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (5)
client/oauth.go (2)
56-61: Useerrors.Asto recognize wrapped errorsThe current implementation of
GetOAuthHandleruses direct type assertion which will fail when errors are wrapped. This is the same issue that was previously fixed inIsOAuthAuthorizationRequiredError.func GetOAuthHandler(err error) *transport.OAuthHandler { - if oauthErr, ok := err.(*OAuthAuthorizationRequiredError); ok { - return oauthErr.Handler - } + var oauthErr *OAuthAuthorizationRequiredError + if errors.As(err, &oauthErr) { + return oauthErr.Handler + } return nil }
3-7:⚠️ Potential issueAdd the
errorspackage to importsThe
IsOAuthAuthorizationRequiredErrorfunction on line 52 useserrors.As, but theerrorspackage is not imported. This will cause compilation errors.import ( + "errors" "fmt" "github.com/mark3labs/mcp-go/client/transport" )client/transport/oauth.go (3)
294-312: ImprovegetDefaultEndpointsto handle empty base URLsThe function may return relative paths if the base URL lacks a scheme or host. This would cause runtime failures when used in HTTP requests.
func (h *OAuthHandler) getDefaultEndpoints(baseURL string) *AuthServerMetadata { // Parse the base URL to extract the authority parsedURL, err := url.Parse(baseURL) if err != nil { return nil } // Discard any path component to get the authorization base URL parsedURL.Path = "" authBaseURL := parsedURL.String() + + // Check if the URL has scheme and host + if authBaseURL == "" { + return nil + } return &AuthServerMetadata{ Issuer: authBaseURL, AuthorizationEndpoint: authBaseURL + "/authorize", TokenEndpoint: authBaseURL + "/token", RegistrationEndpoint: authBaseURL + "/register", } }
194-197: Parse OAuth error responses properlyThe current implementation doesn't attempt to parse OAuth error responses, which typically follow a standardized JSON format. This makes error handling less informative.
if resp.StatusCode != http.StatusOK { body, _ := io.ReadAll(resp.Body) + + // Try to parse as OAuth error response + var oauthError struct { + Error string `json:"error"` + ErrorDescription string `json:"error_description,omitempty"` + } + + if err := json.Unmarshal(body, &oauthError); err == nil && oauthError.Error != "" { + return nil, fmt.Errorf("oauth error: %s - %s", oauthError.Error, oauthError.ErrorDescription) + } + return nil, fmt.Errorf("refresh token request failed with status %d: %s", resp.StatusCode, body) }This same pattern should also be applied to error handling in
ProcessAuthorizationResponse(lines 427-430) andRegisterClient(lines 364-367).
387-406: Add state parameter validationThe
ProcessAuthorizationResponsefunction doesn't validate the state parameter against an expected value, which is important for preventing CSRF attacks.func (h *OAuthHandler) ProcessAuthorizationResponse(ctx context.Context, code, state, codeVerifier string) error { + // Validate state parameter (assuming you store expected state elsewhere) + // This could be done by comparing with a stored value or using a signed state + // For example, if you have a stateStore interface: + // if !h.stateStore.ValidateState(state) { + // return errors.New("invalid state parameter, possible CSRF attack") + // } + metadata, err := h.getServerMetadata(ctx) if err != nil { return fmt.Errorf("failed to get server metadata: %w", err) }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
client/oauth.go(1 hunks)client/transport/oauth.go(1 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (1)
client/oauth.go (4)
client/transport/oauth.go (6)
OAuthConfig(18-34)Token(45-58)TokenStore(37-42)MemoryTokenStore(69-72)NewMemoryTokenStore(75-77)OAuthHandler(111-117)client/transport/streamable_http.go (4)
StreamableHTTPCOption(22-22)WithOAuth(44-48)NewStreamableHTTP(84-103)OAuthAuthorizationRequiredError(155-157)client/client.go (2)
Client(16-25)NewClient(44-54)client/transport/oauth_utils.go (3)
GenerateCodeVerifier(21-24)GenerateCodeChallenge(27-32)GenerateState(35-37)
🪛 GitHub Check: lint
client/oauth.go
[failure] 52-52:
undefined: errors (typecheck)
client/transport/oauth.go
[failure] 216-216:
invalid character U+2026 '…' in identifier
[failure] 225-225:
syntax error: unexpected name context in argument list; possibly missing comma or )) (typecheck)
[failure] 216-216:
invalid character U+2026 '…' in identifier
[failure] 225-225:
syntax error: unexpected name context in argument list; possibly missing comma or ) (typecheck)
[failure] 216-216:
illegal character U+2026 '…' (typecheck)
[failure] 225-225:
missing ',' in argument list (typecheck)
[failure] 226-226:
missing ',' in argument list (typecheck)
[failure] 227-227:
expected operand, found '}' (typecheck)
🪛 GitHub Check: test
client/transport/oauth.go
[failure] 216-216:
invalid character U+2026 '…' in identifier
[failure] 225-225:
syntax error: unexpected name context in argument list; possibly missing comma or )
🪛 GitHub Actions: go
client/transport/oauth.go
[error] 216-216: Invalid character U+2026 '…' in identifier at line 216.
🔇 Additional comments (5)
client/oauth.go (3)
26-35: The function logic looks goodThe
NewOAuthStreamableHttpClientfunction correctly creates a new client with OAuth support by adding the OAuth option to the transport options. It properly handles errors from creating the transport.
37-45: Good use of variable declaration for function aliasingUsing variable declarations to alias the transport functions is a clean approach for re-exporting functionality from the transport package.
50-53: The implementation is correctThe function uses
errors.Ascorrectly to check if the error is (or wraps) anOAuthAuthorizationRequiredError.🧰 Tools
🪛 GitHub Check: lint
[failure] 52-52:
undefined: errors (typecheck)client/transport/oauth.go (2)
142-145: Good improvement for empty access token checkI see you've added a check for empty access tokens in
getValidToken, which was a previously identified issue. This ensures that empty tokens won't be treated as valid.
209-212: Good handling of refresh tokensThe code now properly preserves old refresh tokens when the server doesn't provide a new one, which was a previously identified issue.
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (2)
client/transport/oauth.go (2)
196-199: 🛠️ Refactor suggestionParse OAuth error responses properly for better diagnostics.
OAuth 2.0 servers return errors in a standardized JSON format, but the current implementation doesn't attempt to parse these responses.
if resp.StatusCode != http.StatusOK { body, _ := io.ReadAll(resp.Body) + + // Try to parse as OAuth error response + var oauthError struct { + Error string `json:"error"` + ErrorDescription string `json:"error_description,omitempty"` + } + + if err := json.Unmarshal(body, &oauthError); err == nil && oauthError.Error != "" { + return nil, fmt.Errorf("oauth error: %s - %s", oauthError.Error, oauthError.ErrorDescription) + } + return nil, fmt.Errorf("refresh token request failed with status %d: %s", resp.StatusCode, body) }This same pattern should be applied to other HTTP response error handling sections (lines 467-470, 547-550).
397-415: 🛠️ Refactor suggestionValidate parsed URL in getDefaultEndpoints to prevent relative paths.
If
baseURLlacks scheme/host, the resulting endpoints will be relative paths like/authorize,/token, which will fail at runtime when passed to HTTP requests.// getDefaultEndpoints returns default OAuth endpoints based on the base URL func (h *OAuthHandler) getDefaultEndpoints(baseURL string) *AuthServerMetadata { // Parse the base URL to extract the authority parsedURL, err := url.Parse(baseURL) if err != nil { return nil } // Discard any path component to get the authorization base URL parsedURL.Path = "" authBaseURL := parsedURL.String() + + // Validate that we have a complete URL with scheme and host + if parsedURL.Scheme == "" || parsedURL.Host == "" { + return nil + } return &AuthServerMetadata{ Issuer: authBaseURL, AuthorizationEndpoint: authBaseURL + "/authorize", TokenEndpoint: authBaseURL + "/token", RegistrationEndpoint: authBaseURL + "/register", } }
🧹 Nitpick comments (3)
client/transport/oauth_test.go (2)
54-95: Consider verifying all token fields in the MemoryTokenStore test.The test verifies access token, token type, and refresh token, but doesn't verify that
ExpiresInandExpiresAtfields are properly preserved when retrieving tokens.if retrievedToken.RefreshToken != token.RefreshToken { t.Errorf("Expected refresh token to be %s, got %s", token.RefreshToken, retrievedToken.RefreshToken) } + if retrievedToken.ExpiresIn != token.ExpiresIn { + t.Errorf("Expected expires in to be %d, got %d", token.ExpiresIn, retrievedToken.ExpiresIn) + } + // Check that ExpiresAt is preserved (within a small tolerance to account for time precision differences) + if !retrievedToken.ExpiresAt.Equal(token.ExpiresAt) { + t.Errorf("Expected expires at to be %v, got %v", token.ExpiresAt, retrievedToken.ExpiresAt) + }
196-221: Improve error condition specificity in server metadata test.The error message check is very loose, making it hard to determine what specific error condition is being tested. Consider using a more targeted approach by mocking the HTTP client to return specific errors or test against specific error types rather than string content.
- // Verify the error message contains something about a connection error - // since we're now trying to connect to the well-known endpoint - if !strings.Contains(err.Error(), "connection refused") && - !strings.Contains(err.Error(), "failed to send protected resource request") { - t.Errorf("Expected error message to contain connection error, got %s", err.Error()) - } + // Verify that we get an error related to the server metadata discovery + if err == nil { + t.Fatalf("Expected error when getting server metadata with empty URL") + } + // The error should come from attempted discovery, not from invalid URL formatting + if strings.Contains(err.Error(), "invalid URL") { + t.Errorf("Got URL formatting error instead of discovery error: %s", err.Error()) + }client/transport/oauth.go (1)
369-390: Document the limitation of base URL extraction from redirect URI.The current approach of extracting a base URL from the redirect URI assumes that the redirect URI is on the same domain as the OAuth server, which may not always be true.
// extractBaseURL extracts the base URL from the first request func (h *OAuthHandler) extractBaseURL() (string, error) { // If we have a base URL from a previous request, use it if h.baseURL != "" { return h.baseURL, nil } + // Note: This approach assumes the redirect URI is on the same domain as the OAuth server, + // which may not always be true. In production, it's better to explicitly provide the base URL. + // Otherwise, we need to infer it from the redirect URI if h.config.RedirectURI == "" { return "", fmt.Errorf("no base URL available and no redirect URI provided") }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
client/oauth.go(1 hunks)client/transport/oauth.go(1 hunks)client/transport/oauth_test.go(1 hunks)client/transport/streamable_http_oauth_test.go(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (2)
- client/transport/streamable_http_oauth_test.go
- client/oauth.go
🧰 Additional context used
🧬 Code Graph Analysis (1)
client/transport/oauth.go (4)
client/oauth.go (5)
TokenStore(17-17)Token(14-14)NewMemoryTokenStore(23-23)MemoryTokenStore(20-20)OAuthConfig(11-11)client/client.go (1)
Client(16-25)client/transport/streamable_http.go (1)
ErrOAuthAuthorizationRequired(152-152)mcp/types.go (1)
Resource(613-628)
🔇 Additional comments (11)
client/transport/oauth_test.go (4)
11-52: Well-structured test for token expiration logic.The test comprehensively covers the
IsExpiredmethod with valid, expired, and no-expiration test cases.
97-157: Comprehensive test for redirect URI validation.The test covers a wide range of use cases including secure HTTPS URIs, localhost variants (domain, IPv4, IPv6), and various invalid cases.
159-194: Well-implemented test for empty access token handling.The test correctly verifies that attempting to get an authorization header with an empty access token returns the appropriate error type.
223-263: Good test coverage for CSRF protection.The test properly verifies state parameter validation logic, including both matching and non-matching states, as well as edge cases like empty expected state.
client/transport/oauth.go (7)
18-34: Well-designed and documented OAuth configuration structure.The
OAuthConfigstruct includes all necessary configuration options with clear documentation.
36-66: Clean token model with proper expiration handling.The
Tokenstruct includes all standard OAuth token fields, and theIsExpiredmethod correctly handles tokens with no expiration time.
68-95: Thread-safe in-memory token store implementation.The
MemoryTokenStoreprovides a simple, thread-safe implementation with proper error handling.
142-160: Robust token validation logic with proper error handling.The
getValidTokenmethod correctly addresses the previous issue with empty access tokens by adding the check fortoken.AccessToken != "".
257-335: Comprehensive OAuth metadata discovery with RFC 9728 support.The implementation correctly follows the RFC 9728 standard for OAuth Protected Resource metadata discovery and includes multiple fallback mechanisms.
490-502: Proper implementation of state validation for CSRF protection.The state parameter validation is correctly implemented to prevent CSRF attacks, addressing the previous review comment.
570-596: Well-implemented authorization URL generation with PKCE support.The method properly stores the state for later validation and includes all necessary OAuth parameters.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (3)
client/transport/oauth_test.go (3)
54-95: MemoryTokenStore implementation looks goodThe test properly verifies both error handling for empty stores and correct token persistence/retrieval functionality, checking all token fields.
Consider adding a test case for concurrent access to validate that the mutex in
MemoryTokenStoreprovides proper thread safety.
196-221: Metadata URL empty handling test could be more robustWhile the test correctly verifies that an empty metadata URL produces an error, the error message validation is somewhat brittle, relying on string matching that could break with implementation changes.
Consider refactoring this test to use a mock HTTP client that returns a specific error, making the test more resilient to implementation changes in error messages.
262-302: State validation test covers critical security casesThe test correctly verifies CSRF protection by checking:
- Rejection of non-matching state values
- Handling of empty expected states
This is essential for preventing cross-site request forgery attacks in the OAuth flow.
Consider adding a positive test case that mocks a successful token response when the state matches, to complete the coverage of the authorization flow.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
client/transport/oauth.go(1 hunks)client/transport/oauth_test.go(1 hunks)
🔇 Additional comments (12)
client/transport/oauth_test.go (4)
1-52: Test coverage for Token.IsExpired is comprehensive!The test thoroughly covers all key scenarios for token expiration:
- Valid non-expired tokens
- Expired tokens
- Tokens without expiration information
The test structure is clean with a good table-driven approach.
97-157: Redirect URI validation test is thorough and security-focusedThe test covers critical security scenarios including:
- Valid HTTPS URIs
- Localhost exceptions (both hostname and IP formats including IPv6)
- Invalid HTTP URIs
- Empty and malformed URIs
This is important for preventing redirect-based attacks.
159-194: Good handling of empty access token caseThe test verifies that authorization attempts with an empty access token properly trigger the authorization required error, which is an important security check.
223-260: OAuth error formatting test is clear and focusedThe test properly validates that OAuth errors are formatted correctly both with and without descriptions.
client/transport/oauth.go (8)
1-140: Well-structured OAuth client implementation with comprehensive type definitionsThe code provides a complete implementation of an OAuth 2.0 client with:
- Clear configuration options including PKCE support
- Thread-safe token storage
- Proper token expiration handling
- Clean interfaces for extensibility
The
OAuthConfigincludes all necessary fields for a compliant OAuth 2.0 client, and the token-related structures and methods handle the token lifecycle correctly.
142-160: Token validation logic is secure and properly implementedThe
getValidTokenmethod correctly:
- Checks for token expiration
- Verifies the access token is not empty
- Falls back to refresh token flow when needed
- Returns an authorization required error when a valid token isn't available
This ensures that only valid tokens are used for requests, which is crucial for security.
162-228: Token refresh implementation follows OAuth 2.0 best practicesThe refresh token flow properly:
- Sends the correct parameters to the token endpoint
- Uses the client secret only when available
- Sets appropriate headers for the request
- Handles error responses by parsing OAuth error formats
- Preserves the old refresh token if a new one isn't provided
- Sets proper token expiration time
This ensures reliable token refresh according to the OAuth 2.0 specification.
230-275: Good error handling and utility methodsThe implementation includes:
- Public wrappers for internal methods
- Getter methods for configuration values
- Proper OAuth error extraction and formatting
The
extractOAuthErrorfunction intelligently tries to parse standard OAuth error responses first, falling back to raw error messages when needed.
277-432: Excellent RFC 9728 metadata discovery implementationThe metadata discovery process is robust, following a cascading approach:
- Uses explicitly provided metadata URL if available
- Tries RFC 9728 OAuth Protected Resource discovery
- Falls back to OpenID Connect discovery
- Tries OAuth Authorization Server Metadata
- Constructs default endpoints as a last resort
The implementation uses
sync.Oncefor thread safety and efficient caching, and includes proper error handling at each step. Base URL extraction is also handled gracefully.
434-462: Secure endpoint generation with proper validationThe
getDefaultEndpointsmethod:
- Parses the base URL correctly
- Validates that the URL has both scheme and host components
- Returns descriptive errors for invalid URLs
- Constructs standard OAuth endpoints
This ensures that even when falling back to default endpoints, the URLs are valid and secure.
464-535: Dynamic client registration implementation is comprehensiveThe client registration flow:
- Verifies server support for registration
- Adapts request parameters based on client type (public vs. confidential)
- Includes all required registration parameters
- Updates client configuration with server-provided credentials
- Properly handles error responses
This enables the MCP client to work with OAuth servers that support dynamic registration.
537-643: Authorization flow implements critical security measuresThe authorization flow implementation:
- Properly validates state parameters to prevent CSRF attacks
- Cleans up state after validation
- Includes PKCE support for enhanced security
- Correctly handles the token exchange process
- Sets appropriate token expiration
- Constructs properly formatted authorization URLs
These security measures are essential for a safe OAuth implementation.
ezynda3
requested review from
pottekkat and
rwjblue-glean
and removed request for
giridhar-murthy-glean and
david-hamilton-glean
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
client/transport/streamable_http_oauth_test.go (2)
15-51: Improve test server logic clarity and consider edge cases.The test server logic mixes request counting with authorization header validation, which can be confusing. Consider refactoring for better readability.
func TestStreamableHTTP_WithOAuth(t *testing.T) { // Track request count to simulate 401 on first request, then success requestCount := 0 - authHeaderReceived := "" // Create a test server that requires OAuth server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - // Capture the Authorization header - authHeaderReceived = r.Header.Get("Authorization") // Check for Authorization header if requestCount == 0 { // First request - simulate 401 to test error handling requestCount++ w.WriteHeader(http.StatusUnauthorized) return } + // Capture and verify the Authorization header for subsequent requests + authHeaderReceived := r.Header.Get("Authorization") // Subsequent requests - verify the Authorization header if authHeaderReceived != "Bearer test-token" { t.Errorf("Expected Authorization header 'Bearer test-token', got '%s'", authHeaderReceived) w.WriteHeader(http.StatusUnauthorized) return }Additionally, consider adding test cases for:
- Expired tokens
- Invalid token format
- Missing Authorization header scenarios
119-143: Good authorization flow verification with minor improvement opportunity.The second request test properly verifies successful OAuth flow. However, there's a scope issue with the
authHeaderReceivedvariable that was captured in the outer scope but verified here.Consider moving the authorization header verification logic inside the test server handler for each specific request, or use a more structured approach to capture and verify headers per request to avoid potential race conditions or confusion.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
client/oauth.go(1 hunks)client/transport/streamable_http.go(10 hunks)client/transport/streamable_http_oauth_test.go(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (2)
- client/transport/streamable_http.go
- client/oauth.go
🔇 Additional comments (6)
client/transport/streamable_http_oauth_test.go (6)
1-13: LGTM! Clean and appropriate imports.The package declaration and imports are well-organized and include all necessary dependencies for OAuth transport testing.
54-64: LGTM! Proper token store setup with error handling.The token store initialization and token saving are correctly implemented with appropriate error handling.
67-89: LGTM! Comprehensive OAuth configuration and verification.The OAuth configuration setup and subsequent verification of OAuth enablement and handler presence are well-implemented and thorough.
92-116: Excellent OAuth error handling verification.The test properly verifies that the first request returns an
OAuthAuthorizationRequiredErrorwith the correct handler, demonstrating good error type checking and OAuth flow understanding.
145-192: LGTM! Comprehensive unauthorized scenario testing.This test effectively covers the scenario where no valid token is available, properly verifying that
OAuthAuthorizationRequiredErroris returned with the appropriate handler.
194-218: LGTM! Clean OAuth enablement flag testing.The test properly verifies OAuth enablement status for both enabled and disabled configurations, providing good coverage of the feature flag functionality.
Description
Fixes #295
Type of Change
Checklist
MCP Spec Compliance
Additional Information
Summary by CodeRabbit
Summary by CodeRabbit
New Features
Bug Fixes
Tests
Documentation