fix: use io.LimitReader in key fetcher (#558)

## Motivation and Context

Replace server-only `http.MaxBytesReader(nil, ...)` with
`io.LimitReader` to cap the fetched body to 4KB in key fetcher. This
avoids passing a `nil` ResponseWriter and aligns with Go idioms. If more
than 4KB is received, return a clear error. Previously `MaxBytesReader`
returned a misleading `http: request body too large` error, which is
server-oriented. With this change the error becomes `HTTP auth key
response too large`.

## How Has This Been Tested?
<!-- Have you tested this in a real application? Which scenarios were
tested? -->

Add tests covering success, oversized bodies, connection failure,
non-200 status, and read failures. This covers almost all lines of
`FetchKey` method of `DefaultHTTPKeyFetcher`. For tests, I also added a
new initialiser `NewDefaultHTTPKeyFetcherWithClient` so that the HTTP
client could be modified for test use.

Tests use `httptest.NewTLSServer` plus a helper to construct a client
pinned to the test server.

`auth` package test coverage increased from 56.8% to 59.3%.

## Breaking Changes
<!-- Will users need to update their code or configurations? -->

None.

## Types of changes
<!-- What types of changes does your code introduce? Put an `x` in all
the boxes that apply: -->
- [x] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update

## Checklist
<!-- Go over all the following points, and put an `x` in all the boxes
that apply. -->
- [x] I have read the [MCP
Documentation](https://modelcontextprotocol.io)
- [x] My code follows the repository's style guidelines
- [x] New and existing tests pass locally
- [x] I have added appropriate error handling
- [ ] I have added or updated documentation as needed

## Additional context
<!-- Add any other context, implementation notes, or design decisions
-->

Signed-off-by: Ville Vesilehto <ville@vesilehto.fi>
This commit is contained in:
Ville Vesilehto 2025-09-26 13:24:07 +03:00 • committed by GitHub
commit 76365a61e3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 159 additions and 4 deletions

View file

@ -18,6 +18,9 @@ import (
"github.com/modelcontextprotocol/registry/internal/config"
)
// MaxKeyResponseSize is the maximum size of the response body from the HTTP endpoint.
const MaxKeyResponseSize = 4096
// HTTPTokenExchangeInput represents the input for HTTP-based authentication
type HTTPTokenExchangeInput struct {
Body struct {
@ -51,6 +54,12 @@ func NewDefaultHTTPKeyFetcher() *DefaultHTTPKeyFetcher {
}
}
// NewDefaultHTTPKeyFetcherWithClient creates a new HTTP key fetcher with a custom HTTP client.
// This is primarily useful in tests to inject transports or TLS settings.
func NewDefaultHTTPKeyFetcherWithClient(client *http.Client) *DefaultHTTPKeyFetcher {
return &DefaultHTTPKeyFetcher{client: client}
}
// FetchKey fetches the public key from the well-known HTTP endpoint
func (f *DefaultHTTPKeyFetcher) FetchKey(ctx context.Context, domain string) (string, error) {
url := fmt.Sprintf("https://%s/.well-known/mcp-registry-auth", domain)
@ -73,13 +82,16 @@ func (f *DefaultHTTPKeyFetcher) FetchKey(ctx context.Context, domain string) (st
return "", fmt.Errorf("HTTP %d: failed to fetch key from %s", resp.StatusCode, url)
}
// Limit response size to prevent DoS attacks
resp.Body = http.MaxBytesReader(nil, resp.Body, 4096)
body, err := io.ReadAll(resp.Body)
// Limit response size to prevent DoS attacks.
// Read up to MaxKeyResponseSize+1 and error if exceeded.
limited := io.LimitReader(resp.Body, MaxKeyResponseSize+1)
body, err := io.ReadAll(limited)
if err != nil {
return "", fmt.Errorf("failed to read response body: %w", err)
}
if len(body) > MaxKeyResponseSize {
return "", fmt.Errorf("HTTP auth key response too large")
}
return strings.TrimSpace(string(body)), nil
}

View file

@ -1,11 +1,16 @@
package auth_test
import (
"bytes"
"context"
"crypto/ed25519"
"crypto/tls"
"encoding/base64"
"encoding/hex"
"fmt"
"net"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
@ -18,6 +23,26 @@ import (
"github.com/modelcontextprotocol/registry/internal/config"
)
const wellKnownPath = "/.well-known/mcp-registry-auth"
func newClientForTLSServer(t *testing.T, srv *httptest.Server) *http.Client {
t.Helper()
dialAddr := srv.Listener.Addr().String()
transport := &http.Transport{
TLSClientConfig: &tls.Config{InsecureSkipVerify: true}, //nolint:gosec // testing only
DialContext: func(ctx context.Context, network, _ string) (net.Conn, error) {
d := &net.Dialer{}
return d.DialContext(ctx, network, dialAddr)
},
ForceAttemptHTTP2: false,
MaxIdleConns: 10,
IdleConnTimeout: 30 * time.Second,
TLSHandshakeTimeout: 5 * time.Second,
}
return &http.Client{Transport: transport, Timeout: 10 * time.Second}
}
// MockHTTPKeyFetcher for testing
type MockHTTPKeyFetcher struct {
keyResponses map[string]string
@ -241,3 +266,121 @@ func TestDefaultHTTPKeyFetcher_FetchKey(t *testing.T) {
_, err := fetcher.FetchKey(context.Background(), "nonexistent-test-domain-12345.com")
assert.Error(t, err)
}
func TestDefaultHTTPKeyFetcher(t *testing.T) {
tests := []struct {
name string
handler http.HandlerFunc
wantErrSub string
customClient *http.Client
expectOK bool
wantBody string
}{
{
name: "oversized body",
handler: func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != wellKnownPath {
w.WriteHeader(http.StatusNotFound)
return
}
w.WriteHeader(http.StatusOK)
w.Header().Set("Content-Type", "text/plain")
_, _ = w.Write(bytes.Repeat([]byte("A"), 6000))
},
wantErrSub: "too large",
},
{
name: "non-OK status",
handler: func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path == wellKnownPath {
w.WriteHeader(http.StatusInternalServerError)
return
}
w.WriteHeader(http.StatusNotFound)
},
wantErrSub: "HTTP 500",
},
{
name: "connection failure",
handler: nil,
customClient: &http.Client{
Transport: &http.Transport{
TLSClientConfig: &tls.Config{InsecureSkipVerify: true}, //nolint:gosec // testing only
DialContext: func(context.Context, string, string) (net.Conn, error) {
return nil, fmt.Errorf("dial blocked")
},
},
Timeout: 5 * time.Second,
},
wantErrSub: "failed to fetch key",
},
{
name: "response body failure",
handler: func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != wellKnownPath {
w.WriteHeader(http.StatusNotFound)
return
}
w.Header().Set("Content-Type", "text/plain")
w.Header().Set("Content-Length", "100")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte("PARTIAL"))
},
wantErrSub: "failed to read response body",
},
{
name: "success",
handler: func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != wellKnownPath {
w.WriteHeader(http.StatusNotFound)
return
}
w.Header().Set("Content-Type", "text/plain")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte("response"))
},
expectOK: true,
wantBody: "response",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if tt.handler != nil {
srv := httptest.NewTLSServer(tt.handler)
defer srv.Close()
c := newClientForTLSServer(t, srv)
f := auth.NewDefaultHTTPKeyFetcherWithClient(c)
got, err := f.FetchKey(context.Background(), "example.com")
if tt.expectOK {
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if got != tt.wantBody {
t.Fatalf("unexpected body: got %q want %q", got, tt.wantBody)
}
return
}
if err == nil || !strings.Contains(err.Error(), tt.wantErrSub) {
t.Fatalf("got err=%v, want substring %q", err, tt.wantErrSub)
}
return
}
f := auth.NewDefaultHTTPKeyFetcherWithClient(tt.customClient)
got, err := f.FetchKey(context.Background(), "example.com")
if tt.expectOK {
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if got != tt.wantBody {
t.Fatalf("unexpected body: got %q want %q", got, tt.wantBody)
}
return
}
if err == nil || !strings.Contains(err.Error(), tt.wantErrSub) {
t.Fatalf("got err=%v, want substring %q", err, tt.wantErrSub)
}
})
}
}