Skip to content

Commit 27790de

Browse files
technicalpicklestechnicalpickles
committed
feat(test): add rate limit testing support to OAuth test server
Add rate limit injection capabilities to the OAuth test server for testing autoDetectResource() retry behavior with 429 responses. ## Changes - Add MCPRateLimitCount, MCPRateLimitRetryAfter, MCPRateLimitUseResetAt to ErrorMode - Add rate limit check to oauthMiddleware (before auth check) - Add /.well-known/oauth-protected-resource endpoint (RFC 9728) - Add ProtectedResourceMetadata type - Add MCPURL and ProtectedResourceMetadataURL to ServerResult - Add ResetRateLimitCounter() helper for tests - Add tests for rate limiting and protected resource metadata ## Usage ```go server := oauthserver.Start(t, oauthserver.Options{ ErrorMode: oauthserver.ErrorMode{ MCPRateLimitCount: 2, // Return 429 twice MCPRateLimitRetryAfter: 5, // With Retry-After: 5 }, }) ```
1 parent 6b13013 commit 27790de

6 files changed

Lines changed: 225 additions & 27 deletions

File tree

‎tests/oauthserver/discovery.go‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,3 +66,31 @@ func (s *OAuthTestServer) handleDiscovery(w http.ResponseWriter, r *http.Request
6666
return
6767
}
6868
}
69+
70+
// buildProtectedResourceMetadata constructs RFC 9728 Protected Resource Metadata.
71+
func (s *OAuthTestServer) buildProtectedResourceMetadata() *ProtectedResourceMetadata {
72+
return &ProtectedResourceMetadata{
73+
Resource: s.issuerURL + "/mcp",
74+
AuthorizationServers: []string{s.issuerURL},
75+
ScopesSupported: s.options.SupportedScopes,
76+
BearerMethodsSupported: []string{"header"},
77+
}
78+
}
79+
80+
// handleProtectedResourceMetadata handles GET /.well-known/oauth-protected-resource requests.
81+
// This implements RFC 9728 Protected Resource Metadata for resource auto-detection.
82+
func (s *OAuthTestServer) handleProtectedResourceMetadata(w http.ResponseWriter, r *http.Request) {
83+
if r.Method != http.MethodGet {
84+
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
85+
return
86+
}
87+
88+
metadata := s.buildProtectedResourceMetadata()
89+
90+
w.Header().Set("Content-Type", "application/json")
91+
w.Header().Set("Cache-Control", "public, max-age=3600")
92+
if err := json.NewEncoder(w).Encode(metadata); err != nil {
93+
http.Error(w, "Failed to encode metadata", http.StatusInternalServerError)
94+
return
95+
}
96+
}

‎tests/oauthserver/mcp.go‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,19 @@ func (s *OAuthTestServer) createMCPServer() *mcpserver.MCPServer {
5555
// oauthMiddleware wraps the MCP handler with OAuth authentication
5656
func (s *OAuthTestServer) oauthMiddleware(next http.Handler) http.Handler {
5757
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
58+
// Check for rate limit injection BEFORE auth check
59+
if s.options.ErrorMode.MCPRateLimitCount > 0 {
60+
s.mu.Lock()
61+
s.mcpRateLimitHits++
62+
hits := s.mcpRateLimitHits
63+
s.mu.Unlock()
64+
65+
if hits <= s.options.ErrorMode.MCPRateLimitCount {
66+
s.sendMCPRateLimited(w)
67+
return
68+
}
69+
}
70+
5871
// Check for Bearer token authentication
5972
authHeader := r.Header.Get("Authorization")
6073
if authHeader == "" || !strings.HasPrefix(authHeader, "Bearer ") {
@@ -109,3 +122,21 @@ func (s *OAuthTestServer) sendMCPUnauthorized(w http.ResponseWriter) {
109122
w.WriteHeader(http.StatusUnauthorized)
110123
w.Write([]byte(`{"error": "unauthorized", "error_description": "Bearer token required"}`))
111124
}
125+
126+
// sendMCPRateLimited sends a 429 response for rate limit testing
127+
func (s *OAuthTestServer) sendMCPRateLimited(w http.ResponseWriter) {
128+
w.Header().Set("Content-Type", "application/json")
129+
130+
if s.options.ErrorMode.MCPRateLimitRetryAfter > 0 && !s.options.ErrorMode.MCPRateLimitUseResetAt {
131+
w.Header().Set("Retry-After", fmt.Sprintf("%d", s.options.ErrorMode.MCPRateLimitRetryAfter))
132+
}
133+
134+
w.WriteHeader(http.StatusTooManyRequests)
135+
136+
if s.options.ErrorMode.MCPRateLimitUseResetAt {
137+
resetAt := time.Now().Add(time.Duration(s.options.ErrorMode.MCPRateLimitRetryAfter) * time.Second).Unix()
138+
w.Write([]byte(fmt.Sprintf(`{"error": "rate_limited", "reset_at": %d}`, resetAt)))
139+
} else {
140+
w.Write([]byte(`{"error": "rate_limited", "error_description": "Too many requests"}`))
141+
}
142+
}

‎tests/oauthserver/options.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,11 @@ type ErrorMode struct {
4242
// Device code errors
4343
DeviceSlowPoll bool // Return `slow_down` on device polling
4444
DeviceExpired bool // Return `expired_token` on device polling
45+
46+
// MCP endpoint rate limiting (for resource auto-detection testing)
47+
MCPRateLimitCount int // Return 429 this many times before real response
48+
MCPRateLimitRetryAfter int // Retry-After header value (seconds, 0 = omit header)
49+
MCPRateLimitUseResetAt bool // Use JSON body with reset_at instead of Retry-After header
4550
}
4651

4752
// Options configures the OAuth test server behavior.

‎tests/oauthserver/server.go‎

Lines changed: 47 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,9 @@ type OAuthTestServer struct {
3030
refreshTokens map[string]*RefreshTokenData
3131
issuedTokens []TokenInfo
3232

33+
// Rate limit tracking
34+
mcpRateLimitHits int
35+
3336
mu sync.RWMutex
3437

3538
// Test reference
@@ -39,13 +42,15 @@ type OAuthTestServer struct {
3942
// ServerResult contains everything needed to configure a test client.
4043
type ServerResult struct {
4144
// Server URLs
42-
IssuerURL string
43-
AuthorizationEndpoint string
44-
TokenEndpoint string
45-
JWKSURL string
46-
RegistrationEndpoint string // Empty if DCR disabled
47-
DeviceAuthorizationEndpoint string // Empty if device code disabled
48-
ProtectedResourceURL string // For WWW-Authenticate detection tests
45+
IssuerURL string
46+
AuthorizationEndpoint string
47+
TokenEndpoint string
48+
JWKSURL string
49+
RegistrationEndpoint string // Empty if DCR disabled
50+
DeviceAuthorizationEndpoint string // Empty if device code disabled
51+
ProtectedResourceURL string // For WWW-Authenticate detection tests
52+
MCPURL string // MCP endpoint URL (for resource auto-detection testing)
53+
ProtectedResourceMetadataURL string // RFC 9728 metadata URL
4954

5055
// Pre-registered test client (confidential)
5156
ClientID string
@@ -135,16 +140,18 @@ func StartOnPort(t *testing.T, port int, opts Options) *ServerResult {
135140

136141
// Build result
137142
result := &ServerResult{
138-
IssuerURL: issuerURL,
139-
AuthorizationEndpoint: issuerURL + "/authorize",
140-
TokenEndpoint: issuerURL + "/token",
141-
JWKSURL: issuerURL + "/jwks.json",
142-
ProtectedResourceURL: issuerURL + "/protected",
143-
ClientID: confidentialClient.ClientID,
144-
ClientSecret: confidentialClient.ClientSecret,
145-
PublicClientID: publicClient.ClientID,
146-
Shutdown: s.Shutdown,
147-
Server: s,
143+
IssuerURL: issuerURL,
144+
AuthorizationEndpoint: issuerURL + "/authorize",
145+
TokenEndpoint: issuerURL + "/token",
146+
JWKSURL: issuerURL + "/jwks.json",
147+
ProtectedResourceURL: issuerURL + "/protected",
148+
MCPURL: issuerURL + "/mcp",
149+
ProtectedResourceMetadataURL: issuerURL + "/.well-known/oauth-protected-resource",
150+
ClientID: confidentialClient.ClientID,
151+
ClientSecret: confidentialClient.ClientSecret,
152+
PublicClientID: publicClient.ClientID,
153+
Shutdown: s.Shutdown,
154+
Server: s,
148155
}
149156

150157
if opts.EnableDCR {
@@ -226,16 +233,18 @@ func Start(t *testing.T, opts Options) *ServerResult {
226233

227234
// Build result
228235
result := &ServerResult{
229-
IssuerURL: issuerURL,
230-
AuthorizationEndpoint: issuerURL + "/authorize",
231-
TokenEndpoint: issuerURL + "/token",
232-
JWKSURL: issuerURL + "/jwks.json",
233-
ProtectedResourceURL: issuerURL + "/protected",
234-
ClientID: confidentialClient.ClientID,
235-
ClientSecret: confidentialClient.ClientSecret,
236-
PublicClientID: publicClient.ClientID,
237-
Shutdown: s.Shutdown,
238-
Server: s,
236+
IssuerURL: issuerURL,
237+
AuthorizationEndpoint: issuerURL + "/authorize",
238+
TokenEndpoint: issuerURL + "/token",
239+
JWKSURL: issuerURL + "/jwks.json",
240+
ProtectedResourceURL: issuerURL + "/protected",
241+
MCPURL: issuerURL + "/mcp",
242+
ProtectedResourceMetadataURL: issuerURL + "/.well-known/oauth-protected-resource",
243+
ClientID: confidentialClient.ClientID,
244+
ClientSecret: confidentialClient.ClientSecret,
245+
PublicClientID: publicClient.ClientID,
246+
Shutdown: s.Shutdown,
247+
Server: s,
239248
}
240249

241250
if opts.EnableDCR {
@@ -256,6 +265,9 @@ func (s *OAuthTestServer) setupRoutes(mux *http.ServeMux) {
256265
mux.HandleFunc("/.well-known/openid-configuration", s.handleDiscovery)
257266
}
258267

268+
// Protected Resource Metadata (RFC 9728) - always available for resource auto-detection
269+
mux.HandleFunc("/.well-known/oauth-protected-resource", s.handleProtectedResourceMetadata)
270+
259271
// JWKS endpoint
260272
mux.HandleFunc("/jwks.json", s.handleJWKS)
261273

@@ -446,6 +458,14 @@ func (s *OAuthTestServer) SetErrorMode(mode ErrorMode) {
446458
s.options.ErrorMode = mode
447459
}
448460

461+
// ResetRateLimitCounter resets the MCP rate limit hit counter.
462+
// Call this between tests to start fresh.
463+
func (s *OAuthTestServer) ResetRateLimitCounter() {
464+
s.mu.Lock()
465+
defer s.mu.Unlock()
466+
s.mcpRateLimitHits = 0
467+
}
468+
449469
// GetAuthorizationCodes returns pending authorization codes (for debugging).
450470
func (s *OAuthTestServer) GetAuthorizationCodes() []AuthCodeInfo {
451471
s.mu.RLock()

‎tests/oauthserver/server_test.go‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -960,3 +960,109 @@ func TestDeviceCode_VerificationPage(t *testing.T) {
960960
assert.Contains(t, string(body), "Device Verification")
961961
assert.Contains(t, string(body), deviceResp.UserCode) // Should pre-fill the user code
962962
}
963+
964+
// --- MCP Rate Limiting Tests ---
965+
966+
func TestMCPRateLimiting(t *testing.T) {
967+
t.Run("returns 429 with Retry-After header", func(t *testing.T) {
968+
server := Start(t, Options{
969+
ErrorMode: ErrorMode{
970+
MCPRateLimitCount: 2,
971+
MCPRateLimitRetryAfter: 5,
972+
},
973+
})
974+
defer server.Shutdown()
975+
976+
client := &http.Client{}
977+
978+
// First request should return 429
979+
resp1, err := client.Post(server.MCPURL, "application/json", strings.NewReader("{}"))
980+
require.NoError(t, err)
981+
defer resp1.Body.Close()
982+
assert.Equal(t, http.StatusTooManyRequests, resp1.StatusCode)
983+
assert.Equal(t, "5", resp1.Header.Get("Retry-After"))
984+
985+
// Second request should also return 429
986+
resp2, err := client.Post(server.MCPURL, "application/json", strings.NewReader("{}"))
987+
require.NoError(t, err)
988+
defer resp2.Body.Close()
989+
assert.Equal(t, http.StatusTooManyRequests, resp2.StatusCode)
990+
991+
// Third request should return 401 (past rate limit count)
992+
resp3, err := client.Post(server.MCPURL, "application/json", strings.NewReader("{}"))
993+
require.NoError(t, err)
994+
defer resp3.Body.Close()
995+
assert.Equal(t, http.StatusUnauthorized, resp3.StatusCode)
996+
assert.Contains(t, resp3.Header.Get("WWW-Authenticate"), "resource_metadata")
997+
})
998+
999+
t.Run("returns 429 with reset_at in body", func(t *testing.T) {
1000+
server := Start(t, Options{
1001+
ErrorMode: ErrorMode{
1002+
MCPRateLimitCount: 1,
1003+
MCPRateLimitRetryAfter: 10,
1004+
MCPRateLimitUseResetAt: true,
1005+
},
1006+
})
1007+
defer server.Shutdown()
1008+
1009+
client := &http.Client{}
1010+
1011+
resp, err := client.Post(server.MCPURL, "application/json", strings.NewReader("{}"))
1012+
require.NoError(t, err)
1013+
defer resp.Body.Close()
1014+
1015+
assert.Equal(t, http.StatusTooManyRequests, resp.StatusCode)
1016+
assert.Empty(t, resp.Header.Get("Retry-After"), "Should not have Retry-After header when using reset_at")
1017+
1018+
body, _ := io.ReadAll(resp.Body)
1019+
assert.Contains(t, string(body), "reset_at")
1020+
})
1021+
1022+
t.Run("ResetRateLimitCounter works", func(t *testing.T) {
1023+
server := Start(t, Options{
1024+
ErrorMode: ErrorMode{
1025+
MCPRateLimitCount: 1,
1026+
MCPRateLimitRetryAfter: 1,
1027+
},
1028+
})
1029+
defer server.Shutdown()
1030+
1031+
client := &http.Client{}
1032+
1033+
// First request returns 429
1034+
resp1, _ := client.Post(server.MCPURL, "application/json", strings.NewReader("{}"))
1035+
resp1.Body.Close()
1036+
assert.Equal(t, http.StatusTooManyRequests, resp1.StatusCode)
1037+
1038+
// Reset counter
1039+
server.Server.ResetRateLimitCounter()
1040+
1041+
// Next request should return 429 again (counter reset)
1042+
resp2, _ := client.Post(server.MCPURL, "application/json", strings.NewReader("{}"))
1043+
resp2.Body.Close()
1044+
assert.Equal(t, http.StatusTooManyRequests, resp2.StatusCode)
1045+
})
1046+
}
1047+
1048+
// --- Protected Resource Metadata Tests ---
1049+
1050+
func TestProtectedResourceMetadata(t *testing.T) {
1051+
server := Start(t, Options{})
1052+
defer server.Shutdown()
1053+
1054+
resp, err := http.Get(server.ProtectedResourceMetadataURL)
1055+
require.NoError(t, err)
1056+
defer resp.Body.Close()
1057+
1058+
assert.Equal(t, http.StatusOK, resp.StatusCode)
1059+
assert.Equal(t, "application/json", resp.Header.Get("Content-Type"))
1060+
1061+
var metadata map[string]interface{}
1062+
err = json.NewDecoder(resp.Body).Decode(&metadata)
1063+
require.NoError(t, err)
1064+
1065+
assert.Equal(t, server.MCPURL, metadata["resource"])
1066+
assert.Contains(t, metadata, "authorization_servers")
1067+
assert.Contains(t, metadata, "scopes_supported")
1068+
}

‎tests/oauthserver/types.go‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,14 @@ type TokenErrorResponse struct {
8686
ErrorURI string `json:"error_uri,omitempty"`
8787
}
8888

89+
// ProtectedResourceMetadata represents OAuth 2.0 Protected Resource Metadata (RFC 9728).
90+
type ProtectedResourceMetadata struct {
91+
Resource string `json:"resource"`
92+
AuthorizationServers []string `json:"authorization_servers"`
93+
ScopesSupported []string `json:"scopes_supported,omitempty"`
94+
BearerMethodsSupported []string `json:"bearer_methods_supported,omitempty"`
95+
}
96+
8997
// DiscoveryMetadata represents OAuth 2.0 Authorization Server Metadata (RFC 8414).
9098
type DiscoveryMetadata struct {
9199
Issuer string `json:"issuer"`

0 commit comments

Comments
 (0)