diff --git a/internal/cmd/proxy.go b/internal/cmd/proxy.go index 49341c66..80ea872d 100644 --- a/internal/cmd/proxy.go +++ b/internal/cmd/proxy.go @@ -19,6 +19,7 @@ import ( "github.com/github/gh-aw-mcpg/internal/httputil" "github.com/github/gh-aw-mcpg/internal/logger" "github.com/github/gh-aw-mcpg/internal/proxy" + "github.com/github/gh-aw-mcpg/internal/util" "github.com/spf13/cobra" ) @@ -304,7 +305,7 @@ func runProxy(cmd *cobra.Command, args []string) error { if tlsCfg != nil { fmt.Fprintf(stderr, " CA cert: %s\n", tlsCfg.CACertPath) fmt.Fprintf(stderr, "\nConnect with:\n") - fmt.Fprintf(stderr, " export GH_HOST=%s\n", clientAddr(actualAddr)) + fmt.Fprintf(stderr, " export GH_HOST=%s\n", util.ClientAddr(actualAddr)) fmt.Fprintf(stderr, " export NODE_EXTRA_CA_CERTS=%s\n", tlsCfg.CACertPath) fmt.Fprintf(stderr, " export SSL_CERT_FILE=%s\n", tlsCfg.CACertPath) fmt.Fprintf(stderr, " export GIT_SSL_CAINFO=%s\n", tlsCfg.CACertPath) @@ -326,25 +327,6 @@ func runProxy(cmd *cobra.Command, args []string) error { return nil } -// clientAddr returns a client-friendly address from a listener address. -// When the host is a wildcard (0.0.0.0, ::, or empty), it substitutes -// "localhost" so the printed GH_HOST value is usable from a client. -// -// Note: output.go uses "127.0.0.1" for the same wildcard substitution in -// the gateway config output, while this function uses "localhost" because -// GH_HOST must be a resolvable hostname for the gh CLI. -func clientAddr(addr string) string { - host, port, err := net.SplitHostPort(addr) - if err != nil { - return addr - } - switch host { - case "", "0.0.0.0", "::", "[::]": - return net.JoinHostPort("localhost", port) - } - return addr -} - // proxyForcePublicReposIfNeeded checks if GITHUB_REPOSITORY is public and, if so, // overrides the allow-only policy's repos field to "public". This prevents agents // in public-repo workflows from reading private repos through the proxy. diff --git a/internal/cmd/proxy_test.go b/internal/cmd/proxy_test.go index 431e5c52..84fcb7ed 100644 --- a/internal/cmd/proxy_test.go +++ b/internal/cmd/proxy_test.go @@ -13,6 +13,7 @@ import ( "github.com/stretchr/testify/require" "github.com/github/gh-aw-mcpg/internal/config" + "github.com/github/gh-aw-mcpg/internal/util" ) // TestDetectGuardWasm_FileNotFound tests that detectGuardWasm returns empty string @@ -527,7 +528,7 @@ func TestClientAddr(t *testing.T) { for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { - assert.Equal(t, tc.expected, clientAddr(tc.input)) + assert.Equal(t, tc.expected, util.ClientAddr(tc.input)) }) } } diff --git a/internal/config/error_format.go b/internal/config/error_format.go new file mode 100644 index 00000000..65b80f5e --- /dev/null +++ b/internal/config/error_format.go @@ -0,0 +1,39 @@ +package config + +import ( + "errors" + "strings" + + "github.com/BurntSushi/toml" +) + +// AppendConfigDocsFooter appends standard documentation links to an error message. +func AppendConfigDocsFooter(sb *strings.Builder) { + sb.WriteString("\n\nPlease check your configuration against the MCP Gateway specification at:") + sb.WriteString("\n" + ConfigSpecURL) + sb.WriteString("\n\nJSON Schema reference:") + sb.WriteString("\n" + SchemaURL) +} + +// FormatConfigError returns a rich diagnostic message for TOML parse errors. +// When err wraps a toml.ParseError, it returns ParseError.ErrorWithUsage() which +// includes a source-code snippet and column pointer, e.g.: +// +// toml: line 5 (field command): expected "=", got "[" instead +// +// 3 | [servers.github] +// 4 | command = "docker" +// 5 | [servers.github +// | ^ +// +// For all other error types, it falls back to err.Error(). +func FormatConfigError(err error) string { + if err == nil { + return "" + } + var perr toml.ParseError + if errors.As(err, &perr) { + return perr.ErrorWithUsage() + } + return err.Error() +} diff --git a/internal/config/validation_errors.go b/internal/config/validation_errors.go index c5837767..d1341644 100644 --- a/internal/config/validation_errors.go +++ b/internal/config/validation_errors.go @@ -1,11 +1,8 @@ package config import ( - "errors" "fmt" "strings" - - "github.com/BurntSushi/toml" ) // Documentation URL constants @@ -85,14 +82,6 @@ func UnsupportedField(fieldName, message, jsonPath, suggestion string) *Validati return InvalidValue(fieldName, message, jsonPath, suggestion) } -// AppendConfigDocsFooter appends standard documentation links to an error message -func AppendConfigDocsFooter(sb *strings.Builder) { - sb.WriteString("\n\nPlease check your configuration against the MCP Gateway specification at:") - sb.WriteString("\n" + ConfigSpecURL) - sb.WriteString("\n\nJSON Schema reference:") - sb.WriteString("\n" + SchemaURL) -} - // InvalidPattern creates a ValidationError for values that don't match a required pattern. // Used by validation_schema.go for container, mount, URL, and other pattern validations. func InvalidPattern(fieldName, value, jsonPath, suggestion string) *ValidationError { @@ -128,26 +117,3 @@ func SchemaValidationError(serverType, message, jsonPath, suggestion string) *Va suggestion, ) } - -// FormatConfigError returns a rich diagnostic message for TOML parse errors. -// When err wraps a toml.ParseError, it returns ParseError.ErrorWithUsage() which -// includes a source-code snippet and column pointer, e.g.: -// -// toml: line 5 (field command): expected "=", got "[" instead -// -// 3 | [servers.github] -// 4 | command = "docker" -// 5 | [servers.github -// | ^ -// -// For all other error types, it falls back to err.Error(). -func FormatConfigError(err error) string { - if err == nil { - return "" - } - var perr toml.ParseError - if errors.As(err, &perr) { - return perr.ErrorWithUsage() - } - return err.Error() -} diff --git a/internal/guard/guard.go b/internal/guard/guard.go index a3e8022d..a9f3a008 100644 --- a/internal/guard/guard.go +++ b/internal/guard/guard.go @@ -62,3 +62,11 @@ type AgentLabelsPayload struct { // RequestState represents any state that the guard needs to pass from request to response // This is useful when the guard needs to carry information from LabelResource to LabelResponse type RequestState interface{} + +// IsSafeOutputsServer reports whether serverID identifies a safe-outputs server. +// Safe-outputs servers use write-sink guards and are subject to special DIFC +// sink-visibility policy (see internal/server/guard_init.go). +// Matches the canonical "safe-outputs" ID and the legacy "safeoutputs" alias. +func IsSafeOutputsServer(serverID string) bool { + return serverID == "safe-outputs" || serverID == "safeoutputs" +} diff --git a/internal/httputil/httputil.go b/internal/httputil/httputil.go index 03acbbf5..a74fa57f 100644 --- a/internal/httputil/httputil.go +++ b/internal/httputil/httputil.go @@ -6,7 +6,6 @@ import ( "encoding/json" "net/http" - "github.com/github/gh-aw-mcpg/internal/difc" "github.com/github/gh-aw-mcpg/internal/logger" ) @@ -55,17 +54,3 @@ func IsTransientHTTPError(statusCode int) bool { } return transient } - -// WriteSimpleHealthResponse writes a minimal {"status":"ok"} JSON health response. -// This is the standard health response shape for endpoints that do not have access -// to detailed server status (e.g. the proxy), keeping them consistent with each other. -func WriteSimpleHealthResponse(w http.ResponseWriter) { - WriteJSONResponse(w, http.StatusOK, map[string]string{"status": "ok"}) -} - -// WriteReflectResponse writes the DIFC label snapshot JSON response for /reflect endpoints. -// Both the server and the proxy expose /reflect; sharing this helper ensures the output -// format evolves consistently as difc.BuildReflectResponse changes. -func WriteReflectResponse(w http.ResponseWriter, difcComponents difc.DIFCComponents) { - WriteJSONResponse(w, http.StatusOK, difc.BuildReflectResponse(difcComponents)) -} diff --git a/internal/httputil/httputil_test.go b/internal/httputil/httputil_test.go index 3fe2e129..3c79f175 100644 --- a/internal/httputil/httputil_test.go +++ b/internal/httputil/httputil_test.go @@ -6,12 +6,9 @@ import ( "net/http" "net/http/httptest" "testing" - "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" - - "github.com/github/gh-aw-mcpg/internal/difc" ) // errorResponseWriter is a minimal http.ResponseWriter whose Write method @@ -244,58 +241,3 @@ func TestWriteErrorResponse(t *testing.T) { }) } } - -// TestWriteSimpleHealthResponse verifies that the helper writes a {"status":"ok"} health response. -func TestWriteSimpleHealthResponse(t *testing.T) { - rec := httptest.NewRecorder() - WriteSimpleHealthResponse(rec) - - assert.Equal(t, http.StatusOK, rec.Code) - assert.Equal(t, "application/json", rec.Header().Get("Content-Type")) - - var body map[string]string - require.NoError(t, json.NewDecoder(rec.Body).Decode(&body)) - assert.Equal(t, "ok", body["status"]) - assert.Len(t, body, 1, "response body should contain only the 'status' field") -} - -// TestWriteReflectResponse verifies that the helper serialises a DIFC label snapshot correctly. -func TestWriteReflectResponse(t *testing.T) { - t.Run("empty components returns valid response shape", func(t *testing.T) { - rec := httptest.NewRecorder() - WriteReflectResponse(rec, difc.DIFCComponents{}) - - assert.Equal(t, http.StatusOK, rec.Code) - assert.Equal(t, "application/json", rec.Header().Get("Content-Type")) - - var body difc.ReflectResponse - require.NoError(t, json.NewDecoder(rec.Body).Decode(&body)) - assert.NotNil(t, body.Agents) - assert.NotEmpty(t, body.Timestamp) - }) - - t.Run("strict mode is reflected in response body", func(t *testing.T) { - components, err := difc.NewComponents("strict", difc.EnforcementStrict) - require.NoError(t, err) - - rec := httptest.NewRecorder() - WriteReflectResponse(rec, components) - - assert.Equal(t, http.StatusOK, rec.Code) - - var body difc.ReflectResponse - require.NoError(t, json.NewDecoder(rec.Body).Decode(&body)) - assert.Equal(t, "strict", body.Mode) - assert.NotNil(t, body.Agents) - }) - - t.Run("timestamp is an RFC3339 string", func(t *testing.T) { - rec := httptest.NewRecorder() - WriteReflectResponse(rec, difc.DIFCComponents{}) - - var body difc.ReflectResponse - require.NoError(t, json.NewDecoder(rec.Body).Decode(&body)) - _, parseErr := time.Parse(time.RFC3339, body.Timestamp) - assert.NoError(t, parseErr, "Timestamp should be a valid RFC3339 string") - }) -} diff --git a/internal/mcp/helpers.go b/internal/mcp/helpers.go index 6ab57566..319a8141 100644 --- a/internal/mcp/helpers.go +++ b/internal/mcp/helpers.go @@ -3,7 +3,6 @@ package mcp import ( "encoding/json" "fmt" - "strings" "github.com/github/gh-aw-mcpg/internal/logger" ) @@ -66,10 +65,3 @@ func callParamMethod[P any](c *Connection, rawParams interface{}, fn func(P) (in } return marshalToResponse(result) } - -// IsSingularReadTool returns true when toolName refers to a tool expected to -// return a single resource (e.g. get_*, *_read). List/search tools are treated -// as collection tools even if they happen to return one item. -func IsSingularReadTool(toolName string) bool { - return !strings.HasPrefix(toolName, "list_") && !strings.HasPrefix(toolName, "search_") -} diff --git a/internal/mcpresult/mcpresult.go b/internal/mcpresult/mcpresult.go index f8e3df78..0b39cc2a 100644 --- a/internal/mcpresult/mcpresult.go +++ b/internal/mcpresult/mcpresult.go @@ -19,8 +19,18 @@ func IsErrorResult(result map[string]interface{}) bool { // maps. It supports both []interface{} values produced by json.Unmarshal and // []map[string]interface{} values produced by helper constructors. // -// Non-map items in []interface{} are skipped so callers can decide whether to -// ignore them or treat them as an error. +// Semantics: lenient — non-map items in []interface{} are silently skipped. +// This contrasts with the strict path in internal/mcp/tool_result.go +// (convertMapToCallToolResult) which returns an error for non-map content items. +// Use this function in contexts where partial results are acceptable (e.g. +// DIFC filtering, rate-limit detection); use the strict path when producing +// SDK-valid CallToolResult values where malformed items must surface as errors. +// +// Both normalization paths handle the same wire format variations: +// - []interface{} — produced by json.Unmarshal on the raw backend response +// - []map[string]interface{} — produced by helper constructors (e.g. BuildMCPTextResponse) +// +// When adding new content type variants, update both paths consistently. func NormalizeContentItems(contentVal interface{}) ([]map[string]interface{}, bool) { logMCPResult.Printf("Normalizing MCP content items from type %T", contentVal) diff --git a/internal/proxy/handler.go b/internal/proxy/handler.go index b9999807..04f9c86c 100644 --- a/internal/proxy/handler.go +++ b/internal/proxy/handler.go @@ -64,13 +64,13 @@ func (h *proxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { // Health check endpoint if rawPath == "/health" || rawPath == "/healthz" { - httputil.WriteSimpleHealthResponse(w) + httputil.WriteJSONResponse(w, http.StatusOK, map[string]string{"status": "ok"}) return } // Reflect endpoint exposes a live DIFC label snapshot. if r.Method == http.MethodGet && rawPath == "/reflect" { - httputil.WriteReflectResponse(w, h.server.DIFCComponents) + httputil.WriteJSONResponse(w, http.StatusOK, difc.BuildReflectResponse(h.server.DIFCComponents)) return } diff --git a/internal/proxy/tls.go b/internal/proxy/tls.go index 638308f8..2d790ffe 100644 --- a/internal/proxy/tls.go +++ b/internal/proxy/tls.go @@ -36,6 +36,7 @@ import ( "github.com/github/gh-aw-mcpg/internal/httputil" "github.com/github/gh-aw-mcpg/internal/logger" + "github.com/github/gh-aw-mcpg/internal/util" ) var logTLS = logger.ForFile() @@ -183,13 +184,10 @@ func GenerateSelfSignedTLS(dir string) (*TLSConfig, error) { } func randomSerial() (*big.Int, error) { - serial, err := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 128)) + serial, err := util.RandomBigInt(128) if err != nil { return nil, fmt.Errorf("failed to generate serial number: %w", err) } - if serial.Sign() == 0 { - serial = new(big.Int).Add(serial, big.NewInt(1)) - } logTLS.Printf("generated random serial number: %s", serial) return serial, nil } diff --git a/internal/server/difc_log_test.go b/internal/server/difc_log_test.go index dac82516..81b9d957 100644 --- a/internal/server/difc_log_test.go +++ b/internal/server/difc_log_test.go @@ -10,7 +10,6 @@ import ( "github.com/github/gh-aw-mcpg/internal/difc" "github.com/github/gh-aw-mcpg/internal/logger" - "github.com/github/gh-aw-mcpg/internal/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -658,34 +657,3 @@ func TestBuildDIFCSingleItemFilteredError_NoReason(t *testing.T) { // No trailing "()" should appear when reason is empty. assert.NotContains(t, err.Error(), "()") } - -// TestIsSingularReadTool verifies the heuristic that distinguishes singular-read tools -// (get_*, *_read) from collection tools (list_*, search_*). -func TestIsSingularReadTool(t *testing.T) { - tests := []struct { - toolName string - singular bool - }{ - {"issue_read", true}, - {"pull_request_read", true}, - {"get_issue", true}, - {"get_pull_request", true}, - {"get_file_contents", true}, - {"get_repository", true}, - {"get_commit", true}, - {"list_issues", false}, - {"list_pull_requests", false}, - {"list_commits", false}, - {"list_branches", false}, - {"search_issues", false}, - {"search_pull_requests", false}, - {"search_code", false}, - {"search_repositories", false}, - } - for _, tc := range tests { - t.Run(tc.toolName, func(t *testing.T) { - assert.Equal(t, tc.singular, mcp.IsSingularReadTool(tc.toolName), - "mcp.IsSingularReadTool(%q) should be %v", tc.toolName, tc.singular) - }) - } -} diff --git a/internal/server/guard_init.go b/internal/server/guard_init.go index f6dfd704..9faf7e83 100644 --- a/internal/server/guard_init.go +++ b/internal/server/guard_init.go @@ -67,7 +67,7 @@ func (us *UnifiedServer) registerGuard(serverID string) error { // Security-by-default: non-safe-outputs write-sink servers get // sink-visibility="public" when no explicit value is configured. // This assumes external sinks release data publicly unless exempted. - if effectiveVisibility == "" && !isSafeOutputsServer(serverID) && !us.isServerExemptFromSinkVisibility(serverID) { + if effectiveVisibility == "" && !guard.IsSafeOutputsServer(serverID) && !us.isServerExemptFromSinkVisibility(serverID) { effectiveVisibility = "public" logger.LogInfoToServer(serverID, "difc", "Defaulting sink-visibility to \"public\" for non-safe-outputs write-sink server (security-by-default)") @@ -77,7 +77,7 @@ func (us *UnifiedServer) registerGuard(serverID string) error { // sink-visibility but the workflow repo is public, force "public" to // prevent exfiltration. This makes the gateway self-defending even // without compiler cooperation. - if effectiveVisibility == "" && isSafeOutputsServer(serverID) { + if effectiveVisibility == "" && guard.IsSafeOutputsServer(serverID) { if vis, ok := us.resolveWorkflowRepoVisibility(); ok && vis == githubhttp.RepoVisibilityPublic { effectiveVisibility = "public" logger.LogWarnToServer(serverID, "difc", diff --git a/internal/server/guard_visibility.go b/internal/server/guard_visibility.go index a5daefcb..b93ebbda 100644 --- a/internal/server/guard_visibility.go +++ b/internal/server/guard_visibility.go @@ -155,12 +155,6 @@ func (us *UnifiedServer) computeForcePublicRepos() bool { return vis == githubhttp.RepoVisibilityPublic } -// isSafeOutputsServer returns true if the server ID identifies a safe-outputs -// server. Matches "safe-outputs" and the legacy "safeoutputs" form. -func isSafeOutputsServer(serverID string) bool { - return serverID == "safe-outputs" || serverID == "safeoutputs" -} - // isServerExemptFromSinkVisibility returns true if the given server should NOT // receive the default sink-visibility="public" enforcement. A server is exempt when: // - forcePublicRepos is explicitly disabled (blanket opt-out) diff --git a/internal/server/handlers.go b/internal/server/handlers.go index 91f30229..a60a9f7e 100644 --- a/internal/server/handlers.go +++ b/internal/server/handlers.go @@ -7,6 +7,7 @@ import ( "os" "time" + "github.com/github/gh-aw-mcpg/internal/difc" "github.com/github/gh-aw-mcpg/internal/httputil" "github.com/github/gh-aw-mcpg/internal/logger" ) @@ -18,7 +19,7 @@ func HandleReflect(unifiedServer *UnifiedServer) http.HandlerFunc { logHandlers.Print("Creating reflect handler") return func(w http.ResponseWriter, r *http.Request) { logHandlers.Printf("Reflect endpoint request: remote=%s, method=%s, path=%s", r.RemoteAddr, r.Method, r.URL.Path) - httputil.WriteReflectResponse(w, unifiedServer.DIFCComponents) + httputil.WriteJSONResponse(w, http.StatusOK, difc.BuildReflectResponse(unifiedServer.DIFCComponents)) } } diff --git a/internal/server/http_helpers.go b/internal/server/http_helpers.go index 9eb142b6..c7752fd8 100644 --- a/internal/server/http_helpers.go +++ b/internal/server/http_helpers.go @@ -81,24 +81,18 @@ func readAndRestoreRequestBody(r *http.Request) ([]byte, error) { return b, nil } -// peekRequestBody reads all bytes from a POST request body and restores it -// so downstream handlers can read it again. -// Returns nil, nil for non-POST requests or requests with no body. -func peekRequestBody(r *http.Request) ([]byte, error) { - if r.Method != http.MethodPost { - return nil, nil - } - - return readAndRestoreRequestBody(r) -} - // logHTTPRequestBody logs the request body for debugging purposes. // It reads the body, logs it, and restores it so it can be read again. // The backendID parameter is optional and can be empty for unified mode. func logHTTPRequestBody(r *http.Request, sessionID, backendID string) { logServerHelpers.Printf("Checking request body: method=%s, hasBody=%v, sessionID=%s", r.Method, r.Body != nil, util.FormatSessionIDForLog(sessionID)) - bodyBytes, err := peekRequestBody(r) + if r.Method != http.MethodPost { + logServerHelpers.Printf("Skipping body logging: not a POST request, no body present, or empty body") + return + } + + bodyBytes, err := readAndRestoreRequestBody(r) if err != nil { logServerHelpers.Printf("Body read failed: err=%v", err) return diff --git a/internal/server/peek_request_body_test.go b/internal/server/peek_request_body_test.go index 9abcd142..bed72df9 100644 --- a/internal/server/peek_request_body_test.go +++ b/internal/server/peek_request_body_test.go @@ -29,10 +29,10 @@ type errorOnCloseReader struct { func (r *errorOnCloseReader) Read(p []byte) (int, error) { return r.data.Read(p) } func (r *errorOnCloseReader) Close() error { return r.closeErr } -// TestPeekRequestBody verifies all branches of peekRequestBody: non-POST methods, +// TestReadAndRestoreRequestBody verifies all branches of readAndRestoreRequestBody: // nil/NoBody bodies, read errors, close errors, empty body, and non-empty body with // body-restoration behaviour. -func TestPeekRequestBody(t *testing.T) { +func TestReadAndRestoreRequestBody(t *testing.T) { t.Parallel() readErr := errors.New("simulated read error") @@ -47,31 +47,7 @@ func TestPeekRequestBody(t *testing.T) { wantBodyVal string }{ { - name: "GET request returns nil without touching body", - buildReq: func() *http.Request { - return httptest.NewRequest(http.MethodGet, "/", bytes.NewBufferString("hello")) - }, - wantBytes: nil, - wantErr: nil, - }, - { - name: "PUT request returns nil without touching body", - buildReq: func() *http.Request { - return httptest.NewRequest(http.MethodPut, "/", bytes.NewBufferString("hello")) - }, - wantBytes: nil, - wantErr: nil, - }, - { - name: "DELETE request returns nil", - buildReq: func() *http.Request { - return httptest.NewRequest(http.MethodDelete, "/", nil) - }, - wantBytes: nil, - wantErr: nil, - }, - { - name: "POST with nil body returns nil", + name: "nil body returns nil", buildReq: func() *http.Request { req := httptest.NewRequest(http.MethodPost, "/", nil) req.Body = nil @@ -81,7 +57,7 @@ func TestPeekRequestBody(t *testing.T) { wantErr: nil, }, { - name: "POST with http.NoBody returns nil", + name: "http.NoBody returns nil", buildReq: func() *http.Request { req := httptest.NewRequest(http.MethodPost, "/", nil) req.Body = http.NoBody @@ -91,7 +67,7 @@ func TestPeekRequestBody(t *testing.T) { wantErr: nil, }, { - name: "POST with non-empty body returns bytes and restores body", + name: "non-empty body returns bytes and restores body", buildReq: func() *http.Request { return httptest.NewRequest(http.MethodPost, "/", bytes.NewBufferString(`{"method":"tools/list"}`)) }, @@ -101,7 +77,7 @@ func TestPeekRequestBody(t *testing.T) { wantBodyVal: `{"method":"tools/list"}`, }, { - name: "POST with binary body restores body for re-reading", + name: "binary body restores body for re-reading", buildReq: func() *http.Request { content := []byte{0x00, 0x01, 0x02, 0xFF} return httptest.NewRequest(http.MethodPost, "/", bytes.NewReader(content)) @@ -111,7 +87,7 @@ func TestPeekRequestBody(t *testing.T) { checkBody: true, }, { - name: "POST with empty body (reader at EOF) returns empty slice", + name: "empty body (reader at EOF) returns empty slice", buildReq: func() *http.Request { req := httptest.NewRequest(http.MethodPost, "/", bytes.NewBufferString("")) // httptest.NewRequest wraps an empty buffer in a ReadCloser rather than @@ -122,7 +98,7 @@ func TestPeekRequestBody(t *testing.T) { wantErr: nil, }, { - name: "POST with read error propagates the error", + name: "read error propagates the error", buildReq: func() *http.Request { req := httptest.NewRequest(http.MethodPost, "/", nil) req.Body = &errorOnReadReader{readErr: readErr} @@ -132,7 +108,7 @@ func TestPeekRequestBody(t *testing.T) { wantErr: readErr, }, { - name: "POST with close error propagates the error", + name: "close error propagates the error", buildReq: func() *http.Request { req := httptest.NewRequest(http.MethodPost, "/", nil) req.Body = &errorOnCloseReader{ @@ -151,7 +127,7 @@ func TestPeekRequestBody(t *testing.T) { t.Parallel() req := tt.buildReq() - got, err := peekRequestBody(req) + got, err := readAndRestoreRequestBody(req) if tt.wantErr != nil { require.Error(t, err) @@ -163,8 +139,8 @@ func TestPeekRequestBody(t *testing.T) { assert.Equal(t, tt.wantBytes, got) if tt.checkBody { - // Verify that peekRequestBody restored the body so it can be read again. - require.NotNil(t, req.Body, "body should not be nil after peek") + // Verify that readAndRestoreRequestBody restored the body so it can be read again. + require.NotNil(t, req.Body, "body should not be nil after read") assert.NotEqual(t, http.NoBody, req.Body, "body should be readable, not NoBody") restored, readErr := io.ReadAll(req.Body) @@ -180,33 +156,34 @@ func TestPeekRequestBody(t *testing.T) { } } -// TestPeekRequestBody_BodyRestoredForMultipleReads confirms the fundamental contract: -// after peekRequestBody returns, downstream handlers can still read the full body. -func TestPeekRequestBody_BodyRestoredForMultipleReads(t *testing.T) { +// TestReadAndRestoreRequestBody_BodyRestoredForMultipleReads confirms the fundamental +// contract: after readAndRestoreRequestBody returns, downstream handlers can still +// read the full body. +func TestReadAndRestoreRequestBody_BodyRestoredForMultipleReads(t *testing.T) { t.Parallel() body := `{"jsonrpc":"2.0","method":"tools/call","id":1}` req := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewBufferString(body)) - // First peek - b1, err := peekRequestBody(req) + // First read + b1, err := readAndRestoreRequestBody(req) require.NoError(t, err) assert.Equal(t, body, string(b1)) - // Body must still be fully readable after the peek. + // Body must still be fully readable after the read. b2, err := io.ReadAll(req.Body) require.NoError(t, err) assert.Equal(t, body, string(b2), "downstream handler should receive the complete body") } -// TestPeekRequestBody_EmptyBodySetsNoBody confirms that when the body is empty the -// request body is replaced with http.NoBody (not a lingering empty reader). -func TestPeekRequestBody_EmptyBodySetsNoBody(t *testing.T) { +// TestReadAndRestoreRequestBody_EmptyBodySetsNoBody confirms that when the body is empty +// the request body is replaced with http.NoBody (not a lingering empty reader). +func TestReadAndRestoreRequestBody_EmptyBodySetsNoBody(t *testing.T) { t.Parallel() req := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewBufferString("")) - got, err := peekRequestBody(req) + got, err := readAndRestoreRequestBody(req) require.NoError(t, err) assert.Empty(t, got) assert.Equal(t, http.NoBody, req.Body, "empty body should be replaced with http.NoBody") diff --git a/internal/server/sdk_logging.go b/internal/server/sdk_logging.go index a9d9ad7f..71cce837 100644 --- a/internal/server/sdk_logging.go +++ b/internal/server/sdk_logging.go @@ -30,7 +30,11 @@ func WithSDKLogging(handler http.Handler, mode string) http.Handler { mode, util.FormatSessionIDForLog(sessionID), util.FormatSessionIDForLog(mcpSessionID), r.Method, r.URL.Path) // Capture and log request body for POST requests - requestBody, err := peekRequestBody(r) + var requestBody []byte + var err error + if r.Method == http.MethodPost { + requestBody, err = readAndRestoreRequestBody(r) + } var jsonrpcReq mcp.Request if err == nil && len(requestBody) > 0 { // Parse JSON-RPC request diff --git a/internal/server/session_auto_init.go b/internal/server/session_auto_init.go index 58181fbe..a43a5261 100644 --- a/internal/server/session_auto_init.go +++ b/internal/server/session_auto_init.go @@ -47,7 +47,7 @@ func WrapWithSessionAutoInit(streamableHandler http.Handler) http.Handler { } // Peek at the request body to detect tools/call. - bodyBytes, err := peekRequestBody(r) + bodyBytes, err := readAndRestoreRequestBody(r) if err != nil || len(bodyBytes) == 0 { streamableHandler.ServeHTTP(w, r) return diff --git a/internal/server/session_auto_init_test.go b/internal/server/session_auto_init_test.go index a3567aa8..e8e688b1 100644 --- a/internal/server/session_auto_init_test.go +++ b/internal/server/session_auto_init_test.go @@ -26,7 +26,7 @@ func newMockStreamableHandler() *mockStreamableHandler { } func (h *mockStreamableHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { - bodyBytes, _ := peekRequestBody(r) + bodyBytes, _ := readAndRestoreRequestBody(r) var rpcReq struct { ID interface{} `json:"id"` @@ -201,7 +201,7 @@ func TestWrapWithSessionAutoInit_AutoInitForToolsCall(t *testing.T) { func TestWrapWithSessionAutoInit_AutoInitPreservesAuthHeader(t *testing.T) { var initAuthHeader string handlerFunc := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - bodyBytes, _ := peekRequestBody(r) + bodyBytes, _ := readAndRestoreRequestBody(r) var rpcReq struct { Method string `json:"method"` } @@ -245,7 +245,7 @@ func TestWrapWithSessionAutoInit_AutoInitPreservesAuthHeader(t *testing.T) { func TestWrapWithSessionAutoInit_FallsBackOnAutoInitFailure(t *testing.T) { // This handler never returns a session ID for initialize. handlerFunc := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - bodyBytes, _ := peekRequestBody(r) + bodyBytes, _ := readAndRestoreRequestBody(r) var rpcReq struct { Method string `json:"method"` } @@ -285,7 +285,7 @@ func TestWrapWithSessionAutoInit_FallsBackOnAutoInitFailure(t *testing.T) { // header, regardless of the HTTP status code. func TestPerformSessionAutoInit_MissingSessionID(t *testing.T) { handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - bodyBytes, _ := peekRequestBody(r) + bodyBytes, _ := readAndRestoreRequestBody(r) var rpcReq struct { Method string `json:"method"` } @@ -313,7 +313,7 @@ func TestPerformSessionAutoInit_MissingSessionID(t *testing.T) { // session ID check in the implementation. func TestPerformSessionAutoInit_NonOKStatusWithSessionID(t *testing.T) { handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - bodyBytes, _ := peekRequestBody(r) + bodyBytes, _ := readAndRestoreRequestBody(r) var rpcReq struct { Method string `json:"method"` } @@ -345,7 +345,7 @@ func TestPerformSessionAutoInit_Success(t *testing.T) { var receivedSessionIDs []string handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - bodyBytes, _ := peekRequestBody(r) + bodyBytes, _ := readAndRestoreRequestBody(r) var rpcReq struct { Method string `json:"method"` } @@ -389,7 +389,7 @@ func TestPerformSessionAutoInit_AuthHeaderCopied(t *testing.T) { var capturedInitAgentID string handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - bodyBytes, _ := peekRequestBody(r) + bodyBytes, _ := readAndRestoreRequestBody(r) var rpcReq struct { Method string `json:"method"` } diff --git a/internal/server/tool_policy.go b/internal/server/tool_policy.go new file mode 100644 index 00000000..87f70ebe --- /dev/null +++ b/internal/server/tool_policy.go @@ -0,0 +1,15 @@ +package server + +import "strings" + +// IsSingularReadTool returns true when toolName refers to a tool expected to +// return a single resource (e.g. get_*, *_read). List/search tools are treated +// as collection tools even if they happen to return one item. +// +// This function encodes server-layer policy about how GitHub-specific MCP tool +// results should be handled in DIFC filtering; it lives here rather than in the +// mcp (wire protocol) package because it is a server-level policy decision, not +// a protocol-level concern. +func IsSingularReadTool(toolName string) bool { + return !strings.HasPrefix(toolName, "list_") && !strings.HasPrefix(toolName, "search_") +} diff --git a/internal/mcp/helpers_test.go b/internal/server/tool_policy_test.go similarity index 99% rename from internal/mcp/helpers_test.go rename to internal/server/tool_policy_test.go index 2f0195ae..d7dc846a 100644 --- a/internal/mcp/helpers_test.go +++ b/internal/server/tool_policy_test.go @@ -1,4 +1,4 @@ -package mcp +package server import ( "testing" diff --git a/internal/server/unified.go b/internal/server/unified.go index d96101dc..0e932b8d 100644 --- a/internal/server/unified.go +++ b/internal/server/unified.go @@ -573,7 +573,7 @@ func (us *UnifiedServer) callBackendTool(ctx context.Context, serverID, toolName // (list_*, search_*) may legitimately return exactly one item that gets filtered // and should still receive the notice-only behavior so agents see an empty list // rather than an unexpected error. - if mcp.IsSingularReadTool(toolName) && difcFiltered.GetAccessibleCount() == 0 && difcFiltered.GetFilteredCount() == 1 { + if IsSingularReadTool(toolName) && difcFiltered.GetAccessibleCount() == 0 && difcFiltered.GetFilteredCount() == 1 { filteredErr := buildDIFCSingleItemFilteredError(difcFiltered.Filtered[0]) logger.LogWarn("difc", "Single item filtered — returning MCP error: %v", filteredErr) httpStatusCode = 403 diff --git a/internal/util/netutil.go b/internal/util/netutil.go new file mode 100644 index 00000000..a6e86383 --- /dev/null +++ b/internal/util/netutil.go @@ -0,0 +1,22 @@ +package util + +import "net" + +// ClientAddr returns a client-friendly address from a listener address. +// When the host is a wildcard (0.0.0.0, ::, or empty), it substitutes +// "localhost" so the printed address is usable from a client. +// +// Note: output.go uses "127.0.0.1" for the same wildcard substitution in +// the gateway config output, while this function uses "localhost" because +// GH_HOST must be a resolvable hostname for the gh CLI. +func ClientAddr(addr string) string { + host, port, err := net.SplitHostPort(addr) + if err != nil { + return addr + } + switch host { + case "", "0.0.0.0", "::", "[::]": + return net.JoinHostPort("localhost", port) + } + return addr +} diff --git a/internal/util/netutil_test.go b/internal/util/netutil_test.go new file mode 100644 index 00000000..e8dc4033 --- /dev/null +++ b/internal/util/netutil_test.go @@ -0,0 +1,56 @@ +package util + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestClientAddr(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + input string + expected string + }{ + { + name: "IPv4 wildcard becomes localhost", + input: "0.0.0.0:8080", + expected: "localhost:8080", + }, + { + name: "IPv6 wildcard :: becomes localhost", + input: "[::]:8443", + expected: "localhost:8443", + }, + { + name: "explicit localhost unchanged", + input: "localhost:3000", + expected: "localhost:3000", + }, + { + name: "explicit 127.0.0.1 unchanged", + input: "127.0.0.1:9090", + expected: "127.0.0.1:9090", + }, + { + name: "non-loopback host unchanged", + input: "192.168.1.1:8080", + expected: "192.168.1.1:8080", + }, + { + name: "invalid address returned as-is", + input: "not-an-addr", + expected: "not-an-addr", + }, + } + + for _, tc := range tests { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + assert.Equal(t, tc.expected, ClientAddr(tc.input)) + }) + } +} diff --git a/internal/util/normalize.go b/internal/util/normalize.go deleted file mode 100644 index f12819d5..00000000 --- a/internal/util/normalize.go +++ /dev/null @@ -1,9 +0,0 @@ -package util - -import "strings" - -// NormalizeStringCI trims surrounding whitespace and lowercases a string for -// case-insensitive comparisons. -func NormalizeStringCI(value string) string { - return strings.ToLower(strings.TrimSpace(value)) -} diff --git a/internal/util/normalize_test.go b/internal/util/normalize_test.go deleted file mode 100644 index b0442a1b..00000000 --- a/internal/util/normalize_test.go +++ /dev/null @@ -1,30 +0,0 @@ -package util - -import ( - "testing" - - "github.com/stretchr/testify/assert" -) - -func TestNormalizeStringCI(t *testing.T) { - t.Parallel() - - tests := []struct { - name string - input string - want string - }{ - {name: "empty", input: "", want: ""}, - {name: "already normalized", input: "public", want: "public"}, - {name: "trimmed and lowercased", input: " Public ", want: "public"}, - {name: "whitespace only", input: " ", want: ""}, - } - - for _, tt := range tests { - tt := tt - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - assert.Equal(t, tt.want, NormalizeStringCI(tt.input)) - }) - } -} diff --git a/internal/util/random.go b/internal/util/random.go index d5b55d3a..ea0492ed 100644 --- a/internal/util/random.go +++ b/internal/util/random.go @@ -6,6 +6,7 @@ import ( "encoding/hex" "fmt" "io" + "math/big" "os" "time" ) @@ -62,3 +63,23 @@ func RandomHexWithFallback(n int) string { } return s } + +// RandomBigInt returns a cryptographically random non-negative integer with +// the given bit width. The result is guaranteed to be strictly positive (≥ 1). +// This centralises crypto/rand.Int usage for callers that need a *big.Int +// (e.g. X.509 certificate serial numbers). +func RandomBigInt(bits uint) (*big.Int, error) { + if bits == 0 { + return nil, fmt.Errorf("failed to generate random big.Int: bits must be > 0") + } + max := new(big.Int).Lsh(big.NewInt(1), bits) + n, err := rand.Int(rand.Reader, max) + if err != nil { + return nil, fmt.Errorf("failed to generate random big.Int(%d bits): %w", bits, err) + } + // Ensure strictly positive: a zero result (astronomically rare) becomes 1. + if n.Sign() == 0 { + n.SetInt64(1) + } + return n, nil +} diff --git a/internal/util/random_test.go b/internal/util/random_test.go index ab670fcb..42fb9d51 100644 --- a/internal/util/random_test.go +++ b/internal/util/random_test.go @@ -310,3 +310,44 @@ func TestRandomBytes_Uniqueness(t *testing.T) { seen[key] = true } } + +// TestRandomBigInt verifies that RandomBigInt produces a positive integer within +// the expected bit range and that successive calls produce distinct values. +func TestRandomBigInt(t *testing.T) { + t.Parallel() + + t.Run("128-bit produces positive integer in range", func(t *testing.T) { + t.Parallel() + n, err := RandomBigInt(128) + require.NoError(t, err) + assert.Positive(t, n.Sign(), "result should be positive") + assert.LessOrEqual(t, n.BitLen(), 128, "result should fit in 128 bits") + }) + + t.Run("64-bit produces positive integer in range", func(t *testing.T) { + t.Parallel() + n, err := RandomBigInt(64) + require.NoError(t, err) + assert.Positive(t, n.Sign()) + assert.LessOrEqual(t, n.BitLen(), 64) + }) + + t.Run("successive calls produce distinct values", func(t *testing.T) { + t.Parallel() + seen := make(map[string]bool) + for i := 0; i < 20; i++ { + n, err := RandomBigInt(128) + require.NoError(t, err) + key := n.String() + assert.False(t, seen[key], "RandomBigInt should produce unique values") + seen[key] = true + } + }) + + t.Run("zero bits returns error", func(t *testing.T) { + t.Parallel() + _, err := RandomBigInt(0) + require.Error(t, err) + assert.Contains(t, err.Error(), "bits must be > 0") + }) +} diff --git a/internal/util/util.go b/internal/util/util.go index c2e5ee49..49cbf5bd 100644 --- a/internal/util/util.go +++ b/internal/util/util.go @@ -5,6 +5,7 @@ import ( "fmt" "math" "sort" + "strings" ) // SortedSetKeys returns the keys of a string set (map[string]struct{}) as a sorted slice. @@ -97,3 +98,9 @@ func InterfaceToIntString(v interface{}) (string, bool) { } return "", false } + +// NormalizeStringCI trims surrounding whitespace and lowercases a string for +// case-insensitive comparisons. +func NormalizeStringCI(value string) string { + return strings.ToLower(strings.TrimSpace(value)) +} diff --git a/internal/util/util_test.go b/internal/util/util_test.go index d174c3ac..3ecba3dd 100644 --- a/internal/util/util_test.go +++ b/internal/util/util_test.go @@ -480,3 +480,26 @@ func TestInterfaceToIntString(t *testing.T) { assert.Equal(t, "", s) }) } + +func TestNormalizeStringCI(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + input string + want string + }{ + {name: "empty", input: "", want: ""}, + {name: "already normalized", input: "public", want: "public"}, + {name: "trimmed and lowercased", input: " Public ", want: "public"}, + {name: "whitespace only", input: " ", want: ""}, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + assert.Equal(t, tt.want, NormalizeStringCI(tt.input)) + }) + } +}