diff --git a/pkg/mcp/server.go b/pkg/mcp/server.go index 5153bf00..8471922b 100644 --- a/pkg/mcp/server.go +++ b/pkg/mcp/server.go @@ -20,6 +20,7 @@ import ( "fmt" "log" "net/http" + "regexp" "strings" "sync" "time" @@ -105,9 +106,10 @@ func (s *Server) Handler() http.Handler { // process restarts. The sweep runs every SessionIdleTTL/2, so state is removed // 30 to 45 minutes after a session's last request. // -// The transport is stateless: session IDs are random, not recorded, and any -// well-formed ID is accepted, so sweeping (or DELETE) frees memory but does -// not invalidate an ID. The privacy policy on gosqlx.dev cites this value. +// The transport is stateless (see newSessionIdManager): session IDs are +// random, not recorded, and any well-formed ID is accepted, so sweeping (or +// DELETE) frees memory but does not invalidate an ID. The privacy policy on +// gosqlx.dev cites this value. const SessionIdleTTL = 30 * time.Minute // streamable returns the server's streamable HTTP transport, building it with @@ -117,6 +119,7 @@ func (s *Server) streamable() *mcpserver.StreamableHTTPServer { defer s.streamMu.Unlock() if s.streamSrv == nil { s.streamSrv = mcpserver.NewStreamableHTTPServer(s.mcpSrv, + mcpserver.WithSessionIdManager(newSessionIdManager()), mcpserver.WithSessionIdleTTL(s.sessionIdleTTL), mcpserver.WithLogger(s.transportLogger), ) @@ -124,6 +127,25 @@ func (s *Server) streamable() *mcpserver.StreamableHTTPServer { return s.streamSrv } +// newSessionIdManager returns the session ID manager for the HTTP transport. +// +// It is pinned explicitly rather than left to mcp-go's default so that an +// upstream default change cannot silently switch the public, unauthenticated +// server to mcp-go's InsecureStatefulSessionIdManager. That manager records +// every terminated session ID in a sync.Map that is never pruned +// (server/streamable_http.go in mcp-go v0.45.0, the terminated field and its +// Store in Terminate), so memory would grow with every session a client opens +// and closes, and the IDs would be retained for the life of the process. +// +// StatelessGeneratingSessionIdManager issues random IDs, checks only their +// format, and keeps no record of issued or terminated IDs. None of the GoSQLX +// tools needs per-session state, so nothing depends on a terminated ID being +// rejected; the transport's own per-session state is still freed by DELETE +// and by the idle sweeper (see SessionIdleTTL). +func newSessionIdManager() mcpserver.SessionIdManager { + return &mcpserver.StatelessGeneratingSessionIdManager{} +} + // Shutdown stops the background session sweeper of the HTTP transport. It // does not close an http.Server the caller built around Handler; shut that // down separately. Handlers returned earlier keep serving requests, without @@ -142,13 +164,41 @@ func (s *Server) Shutdown(ctx context.Context) error { return nil } -// transportLogger adapts mcp-go's transport logging. Informational messages -// are dropped because they carry session IDs (e.g. "Sweeping expired session: -// "); errors, such as failed response writes, are still logged. +// transportLogger adapts mcp-go's transport logging so that session IDs are +// never logged, as the privacy policy on gosqlx.dev promises. +// +// Informational messages are dropped because they carry session IDs (e.g. +// "Sweeping expired session: "). Error messages are logged, but their +// arguments are printed only for formats known to carry no session ID (I/O +// failures and fixed sentinel errors); for any other format the arguments are +// replaced with a placeholder. This matters because some mcp-go errors embed +// the client-supplied Mcp-Session-Id header, e.g. "Failed to deliver sampling +// response: no active session found for session ", which any client can +// trigger by POSTing a response-shaped JSON-RPC message. type transportLogger struct { logger *log.Logger } +// transportErrorsWithSafeArgs lists the mcp-go v0.45.0 Errorf formats whose +// arguments never include a session ID. Formats not listed here (including +// ones added by future mcp-go releases) are logged with redacted arguments. +var transportErrorsWithSafeArgs = map[string]bool{ + "Failed to write SSE event: %v": true, + "Failed to write SSE event during drain: %v": true, + "Failed to write final SSE response event: %v": true, + "Failed to write response: %v": true, + "Failed to register POST session: %v": true, // ErrSessionExists only + "Failed to write JSONRPCError: %v": true, +} + +// redactedArg replaces each argument of a transport error whose format is not +// in transportErrorsWithSafeArgs. +const redactedArg = "[redacted]" + +// sessionIDPattern matches the IDs mcp-go generates. It is applied to every +// logged error as a second line of defence. +var sessionIDPattern = regexp.MustCompile(`mcp-session-[0-9A-Fa-f-]+`) + func newTransportLogger() util.Logger { return transportLogger{logger: log.Default()} } @@ -156,7 +206,15 @@ func newTransportLogger() util.Logger { func (transportLogger) Infof(string, ...any) {} func (l transportLogger) Errorf(format string, v ...any) { - l.logger.Printf("gosqlx-mcp: transport error: "+format, v...) + args := v + if !transportErrorsWithSafeArgs[format] { + args = make([]any, len(v)) + for i := range args { + args[i] = redactedArg + } + } + msg := sessionIDPattern.ReplaceAllString(fmt.Sprintf(format, args...), "mcp-session-"+redactedArg) + l.logger.Printf("gosqlx-mcp: transport error: %s", msg) } // Cfg returns the server configuration. diff --git a/pkg/mcp/session_retention_test.go b/pkg/mcp/session_retention_test.go new file mode 100644 index 00000000..84d6919b --- /dev/null +++ b/pkg/mcp/session_retention_test.go @@ -0,0 +1,221 @@ +// Copyright 2026 GoSQLX Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package mcp + +import ( + "bytes" + "context" + "errors" + "fmt" + "io" + "log" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + + mcpserver "github.com/mark3labs/mcp-go/server" +) + +// syncBuffer is a goroutine-safe bytes.Buffer for capturing log output +// written from HTTP handler goroutines. +type syncBuffer struct { + mu sync.Mutex + buf bytes.Buffer +} + +func (b *syncBuffer) Write(p []byte) (int, error) { + b.mu.Lock() + defer b.mu.Unlock() + return b.buf.Write(p) +} + +func (b *syncBuffer) String() string { + b.mu.Lock() + defer b.mu.Unlock() + return b.buf.String() +} + +// mcpRequest sends method to url with an optional Mcp-Session-Id header and +// returns the status code and body. +func mcpRequest(t *testing.T, method, url, sessionID, body string) (int, string) { + t.Helper() + req, err := http.NewRequest(method, url, strings.NewReader(body)) + if err != nil { + t.Fatalf("new request: %v", err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + if sessionID != "" { + req.Header.Set("Mcp-Session-Id", sessionID) + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("%s: %v", method, err) + } + defer resp.Body.Close() + data, _ := io.ReadAll(resp.Body) + return resp.StatusCode, string(data) +} + +func TestNewSessionIdManager_KeepsNoTerminatedSet(t *testing.T) { + mgr := newSessionIdManager() + if _, ok := mgr.(*mcpserver.StatelessGeneratingSessionIdManager); !ok { + t.Fatalf("newSessionIdManager() = %T, want *StatelessGeneratingSessionIdManager "+ + "(InsecureStatefulSessionIdManager never prunes terminated IDs)", mgr) + } + id := mgr.Generate() + if !strings.HasPrefix(id, "mcp-session-") { + t.Fatalf("Generate() = %q, want an mcp-session- ID", id) + } + if notAllowed, err := mgr.Terminate(id); notAllowed || err != nil { + t.Fatalf("Terminate() = (%v, %v), want (false, nil)", notAllowed, err) + } + // A manager that remembered terminated IDs would report this one as + // terminated; a stateless one has nothing to remember. + if terminated, err := mgr.Validate(id); terminated || err != nil { + t.Errorf("Validate(terminated ID) = (%v, %v), want (false, nil): manager must not track terminations", terminated, err) + } +} + +// TestServer_CreateTerminateCyclesRetainNothing opens and DELETEs many +// sessions over HTTP. If the transport tracked terminated IDs, the first +// session would be rejected with "Session terminated" afterwards; instead it +// is still accepted, proving no per-ID record is kept. +func TestServer_CreateTerminateCyclesRetainNothing(t *testing.T) { + srv := New(DefaultConfig()) + t.Cleanup(func() { _ = srv.Shutdown(context.Background()) }) + ts := httptest.NewServer(srv.Handler()) + t.Cleanup(ts.Close) + url := ts.URL + "/mcp" + + const cycles = 500 + var first string + for i := 0; i < cycles; i++ { + status, sid := postInitialize(t, url) + if status != http.StatusOK || sid == "" { + t.Fatalf("cycle %d initialize: status %d, session %q", i, status, sid) + } + if first == "" { + first = sid + } + if status, body := mcpRequest(t, http.MethodDelete, url, sid, ""); status != http.StatusOK { + t.Fatalf("cycle %d DELETE: status %d, body %q", i, status, body) + } + } + + status, body := mcpRequest(t, http.MethodPost, url, first, `{"jsonrpc":"2.0","id":2,"method":"ping"}`) + if status != http.StatusOK { + t.Errorf("ping with terminated session ID: status %d, body %q; want 200 "+ + "(a 404 means the transport records terminated IDs, which grows without bound)", status, body) + } +} + +func TestTransportLogger_RedactsSessionIDsInErrors(t *testing.T) { + const sid = "mcp-session-8c3e0a7e-58a4-4f43-9e1b-5d1b2c3f4a5b" + tests := []struct { + name string + format string + args []any + want string + mustSkip []string + }{ + { + name: "unknown format arguments are withheld", + format: "Failed to handle sampling response: %v", + args: []any{errors.New("invalid session id: attacker-chosen-value")}, + want: "Failed to handle sampling response: [redacted]", + mustSkip: []string{"attacker-chosen-value"}, + }, + { + name: "deliver error with generated ID", + format: "Failed to deliver sampling response: %v", + args: []any{fmt.Errorf("no active session found for session %s", sid)}, + want: "Failed to deliver sampling response: [redacted]", + mustSkip: []string{sid, "8c3e0a7e"}, + }, + { + name: "safe format keeps its error detail", + format: "Failed to write response: %v", + args: []any{errors.New("write: broken pipe")}, + want: "Failed to write response: write: broken pipe", + mustSkip: nil, + }, + { + name: "generated ID is scrubbed even from a safe format", + format: "Failed to write response: %v", + args: []any{fmt.Errorf("stream for %s closed", sid)}, + want: "stream for mcp-session-[redacted] closed", + mustSkip: []string{sid, "8c3e0a7e"}, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var buf bytes.Buffer + l := transportLogger{logger: log.New(&buf, "", 0)} + l.Errorf(tt.format, tt.args...) + got := buf.String() + if !strings.Contains(got, tt.want) { + t.Errorf("Errorf output = %q, want it to contain %q", got, tt.want) + } + for _, s := range tt.mustSkip { + if strings.Contains(got, s) { + t.Errorf("Errorf output = %q leaks %q", got, s) + } + } + }) + } +} + +// TestServer_SessionIDsNotLoggedEndToEnd drives the mcp-go error paths that +// embed the client's Mcp-Session-Id (a response-shaped JSON-RPC POST with no +// pending sampling request) through the real transport and checks that the +// log carries the error but not the ID. +func TestServer_SessionIDsNotLoggedEndToEnd(t *testing.T) { + logBuf := &syncBuffer{} + srv := New(DefaultConfig()) + srv.transportLogger = transportLogger{logger: log.New(logBuf, "", 0)} + t.Cleanup(func() { _ = srv.Shutdown(context.Background()) }) + ts := httptest.NewServer(srv.Handler()) + t.Cleanup(ts.Close) + url := ts.URL + "/mcp" + + status, sid := postInitialize(t, url) + if status != http.StatusOK || sid == "" { + t.Fatalf("initialize: status %d, session %q", status, sid) + } + const forged = "client-forged-session-id" + samplingResponse := `{"jsonrpc":"2.0","id":7,"result":{"role":"assistant"}}` + for _, id := range []string{sid, forged} { + if status, _ := mcpRequest(t, http.MethodPost, url, id, samplingResponse); status == http.StatusOK { + t.Fatalf("sampling response for %q unexpectedly accepted", id) + } + } + if status, _ := mcpRequest(t, http.MethodDelete, url, sid, ""); status != http.StatusOK { + t.Fatalf("DELETE: status %d", status) + } + + got := logBuf.String() + if !strings.Contains(got, "sampling response") { + t.Fatalf("expected the sampling error path to be logged, got %q", got) + } + uuidPart := strings.TrimPrefix(sid, "mcp-session-") + for _, leak := range []string{sid, uuidPart, forged} { + if strings.Contains(got, leak) { + t.Errorf("transport log leaks session ID %q:\n%s", leak, got) + } + } +}