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:
parent
f0cce8559f
commit
76365a61e3
2 changed files with 159 additions and 4 deletions
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue