From 5dc3e32bed69ec895b58b28510c3c23a28f9ce82 Mon Sep 17 00:00:00 2001 From: sanjayy0612 Date: Sun, 16 Aug 2026 18:16:53 +0530 Subject: [PATCH 1/6] Add anonymous session handling and tests for unauthenticated WebSocket connections --- internal/probe/mcp/checks.go | 35 +++++++++++++++---- internal/probe/mcp/checks_test.go | 56 +++++++++++++++++++++++++++++++ 2 files changed, 85 insertions(+), 6 deletions(-) diff --git a/internal/probe/mcp/checks.go b/internal/probe/mcp/checks.go index 71e6e71..72b7484 100644 --- a/internal/probe/mcp/checks.go +++ b/internal/probe/mcp/checks.go @@ -59,6 +59,18 @@ func reproBody(method string, params any) string { var httpOnlyTransports = []string{"http-streamable", "http-sse-legacy"} var anyTransport = []string{"*"} +type anonymousSessionProvider interface { + AnonymousSession() (probe.Session, error) +} + +func anonymousSession(s probe.Session) (probe.Session, error) { + provider, ok := s.(anonymousSessionProvider) + if !ok { + return nil, fmt.Errorf("session does not support a separate anonymous connection") + } + return provider.AnonymousSession() +} + // streamableHTTPOnly is for probes that depend on a mechanism specific to // the streamable-HTTP session implementation (e.g. the Mcp-Session-Id // response header it captures) that has no equivalent in legacy-SSE or @@ -364,7 +376,11 @@ func (p *oauthMetadataPostureProbe) Transports() []string { return httpOnlyTrans // signal that must not be reported as "published but incomplete." func (p *oauthMetadataPostureProbe) Run(ctx context.Context, s probe.Session, r *report.Report) error { // --- Bearer challenge: observed on the resource's own 401, not metadata. --- - unauthRaw, unauthErr := s.Do(ctx, "tools/list", map[string]any{}, probe.WithNoAuth()) + unauthSess, err := anonymousSession(s) + if err != nil { + return nil + } + unauthRaw, unauthErr := unauthSess.Do(ctx, "tools/list", map[string]any{}) sawChallenge := unauthErr == nil && unauthRaw != nil && unauthRaw.StatusCode == http.StatusUnauthorized if sawChallenge { wwwAuth := unauthRaw.Headers.Get("WWW-Authenticate") @@ -731,10 +747,13 @@ func (p *unauthToolsListProbe) Protocol() string { return "mcp" } func (p *unauthToolsListProbe) Transports() []string { return anyTransport } func (p *unauthToolsListProbe) Run(ctx context.Context, s probe.Session, r *report.Report) error { - // Re-issue tools/list explicitly WITHOUT the auth header, regardless of - // whether the initial handshake used one. This answers the specific - // question: "can an anonymous caller enumerate tools?" - raw, err := s.Do(ctx, "tools/list", map[string]any{}, probe.WithNoAuth()) + // Use a fresh connection so no Authorization header, MCP session ID, or + // authenticated persistent transport state can affect this observation. + unauthSess, err := anonymousSession(s) + if err != nil { + return nil + } + raw, err := unauthSess.Do(ctx, "tools/list", map[string]any{}) if err != nil { return nil // network failure is not a finding; leave silent, CLI logs errors separately } @@ -1085,8 +1104,12 @@ func (p *resourcesPromptsExposureProbe) Protocol() string { return "mcp" } func (p *resourcesPromptsExposureProbe) Transports() []string { return anyTransport } func (p *resourcesPromptsExposureProbe) Run(ctx context.Context, s probe.Session, r *report.Report) error { + unauthSess, err := anonymousSession(s) + if err != nil { + return nil + } for _, method := range []string{"resources/list", "prompts/list"} { - raw, err := s.Do(ctx, method, map[string]any{}, probe.WithNoAuth()) + raw, err := unauthSess.Do(ctx, method, map[string]any{}) if err != nil || raw.StatusCode != 200 { continue } diff --git a/internal/probe/mcp/checks_test.go b/internal/probe/mcp/checks_test.go index 80d48b8..1b52d00 100644 --- a/internal/probe/mcp/checks_test.go +++ b/internal/probe/mcp/checks_test.go @@ -1,6 +1,7 @@ package mcp import ( + "bufio" "context" "fmt" "net/http" @@ -9,9 +10,64 @@ import ( "time" "github.com/hackwither/reap/internal/probe" + "github.com/hackwither/reap/internal/probe/common" "github.com/hackwither/reap/internal/report" ) +func authRequiredWSServer(t *testing.T) *httptest.Server { + t.Helper() + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Authorization") != "Bearer secret" { + w.WriteHeader(http.StatusUnauthorized) + return + } + key := r.Header.Get("Sec-WebSocket-Key") + hj, _ := w.(http.Hijacker) + conn, buf, err := hj.Hijack() + if err != nil { return } + defer conn.Close() + buf.WriteString("HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: " + common.ExpectedWebSocketAccept(key) + "\r\n\r\n") + _ = buf.Flush() + reader := bufio.NewReader(buf) + for { + fin, opcode, payload, err := readWSFrame(reader) + if err != nil { return } + if !fin || opcode != wsOpText { continue } + var req struct { ID int `json:"id"`; Method string `json:"method"` } + if json.Unmarshal(payload, &req) != nil { continue } + result := map[string]any{} + switch req.Method { + case "initialize": + result = map[string]any{"protocolVersion": mcpProtocolVersion, "serverInfo": map[string]any{"name": "auth-ws"}, "capabilities": map[string]any{}} + case "tools/list": + result = map[string]any{"tools": []map[string]any{{"name": "secret_tool"}}} + } + resp, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": req.ID, "result": result}) + if writeServerWSFrame(conn, wsOpText, resp) != nil { return } + } + })) +} + +func TestUnauthToolsListProbeDoesNotReuseAuthenticatedWebSocket(t *testing.T) { + srv := authRequiredWSServer(t) + defer srv.Close() + wsURL := "ws" + srv.URL[len("http"):] + sess, err := NewWSSession(wsURL, "Bearer secret", 5*time.Second) + if err != nil { t.Fatalf("authenticated websocket handshake failed: %v", err) } + defer sess.conn.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + init, _, err := InitializeSession(ctx, sess) + if err != nil || init == nil { t.Fatalf("authenticated initialize failed: %v", err) } + raw, err := sess.Do(ctx, "tools/list", map[string]any{}) + if err != nil || raw.StatusCode != http.StatusOK { t.Fatalf("authenticated tools/list failed: status=%d err=%v", raw.StatusCode, err) } + + rep := &report.Report{Target: report.Target{URL: wsURL, Protocol: "mcp"}} + if err := (&unauthToolsListProbe{}).Run(ctx, sess, rep); err != nil { t.Fatalf("probe failed: %v", err) } + if len(rep.Findings) != 0 { t.Fatalf("authenticated websocket response was incorrectly reported as anonymous exposure: %+v", rep.Findings) } +} + type fakeSession struct { responses map[string]*probe.RawResult lastHost string From ac684df81ef6c064f78d8be9c883df2556e5fb1c Mon Sep 17 00:00:00 2001 From: sanjayy0612 Date: Sun, 16 Aug 2026 18:17:03 +0530 Subject: [PATCH 2/6] Add timeout support and AnonymousSession method to WSSession --- internal/probe/mcp/session_ws.go | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/internal/probe/mcp/session_ws.go b/internal/probe/mcp/session_ws.go index 3aa1190..58203c8 100644 --- a/internal/probe/mcp/session_ws.go +++ b/internal/probe/mcp/session_ws.go @@ -40,6 +40,7 @@ const ( // frame's JSON-RPC id. type WSSession struct { targetURL string + timeout time.Duration conn net.Conn writeMu sync.Mutex @@ -64,6 +65,7 @@ func NewWSSession(targetURL, authHeader string, timeout time.Duration) (*WSSessi } s := &WSSession{ targetURL: targetURL, + timeout: timeout, conn: conn, pending: make(map[int]chan *probe.RawResult), } @@ -71,6 +73,12 @@ func NewWSSession(targetURL, authHeader string, timeout time.Duration) (*WSSessi return s, nil } +// AnonymousSession establishes a separate WebSocket handshake without auth. +// WithNoAuth cannot change the credentials of an already-upgraded socket. +func (s *WSSession) AnonymousSession() (probe.Session, error) { + return NewWSSession(s.targetURL, "", s.timeout) +} + func (s *WSSession) TargetURL() string { return s.targetURL } // Do sends one JSON-RPC request as a text frame and waits for a response From d7e9dbd40be53e4ad48f8fdcb72ad357b1004c0e Mon Sep 17 00:00:00 2001 From: sanjayy0612 Date: Sun, 16 Aug 2026 18:17:12 +0530 Subject: [PATCH 3/6] Add timeout support and AnonymousSession methods to Session and SSESession --- internal/probe/mcp/session.go | 10 ++++++++++ internal/probe/mcp/session_sse.go | 8 ++++++++ 2 files changed, 18 insertions(+) diff --git a/internal/probe/mcp/session.go b/internal/probe/mcp/session.go index 5e4eff7..6eef6b7 100644 --- a/internal/probe/mcp/session.go +++ b/internal/probe/mcp/session.go @@ -35,6 +35,7 @@ type rpcRequest struct { type Session struct { url string httpClient *http.Client + timeout time.Duration authHeader string // e.g. "Bearer xyz", set via --auth-header; empty if none supplied sessionID string // captured from Mcp-Session-Id response header, if the server issues one reqID int @@ -44,10 +45,19 @@ func NewSession(url, authHeader string, timeout time.Duration) *Session { return &Session{ url: url, authHeader: authHeader, + timeout: timeout, httpClient: &http.Client{Timeout: timeout}, } } +// AnonymousSession returns a fresh streamable-HTTP session with no +// authentication or inherited MCP session ID. It is intentionally separate +// from Do(WithNoAuth): an anonymous probe must not reuse authenticated +// transport state. +func (s *Session) AnonymousSession() (probe.Session, error) { + return NewSession(s.url, "", s.timeout), nil +} + func (s *Session) TargetURL() string { return s.url } func (s *Session) Do(ctx context.Context, method string, params any, opts ...probe.ReqOption) (*probe.RawResult, error) { diff --git a/internal/probe/mcp/session_sse.go b/internal/probe/mcp/session_sse.go index e9b307a..af952d6 100644 --- a/internal/probe/mcp/session_sse.go +++ b/internal/probe/mcp/session_sse.go @@ -31,6 +31,7 @@ import ( type SSESession struct { sseURL string authHeader string + timeout time.Duration httpClient *http.Client connOnce sync.Once @@ -47,6 +48,7 @@ func NewSSESession(sseURL, authHeader string, timeout time.Duration) *SSESession return &SSESession{ sseURL: sseURL, authHeader: authHeader, + timeout: timeout, // The GET stream is intentionally long-lived — per-request // deadlines are enforced via the ctx each Do() call receives, not // a client-wide timeout that would kill the SSE connection itself. @@ -56,6 +58,12 @@ func NewSSESession(sseURL, authHeader string, timeout time.Duration) *SSESession } } +// AnonymousSession creates a new SSE stream and POST channel without auth. +// The authenticated stream cannot be made anonymous after its handshake. +func (s *SSESession) AnonymousSession() (probe.Session, error) { + return NewSSESession(s.sseURL, "", s.timeout), nil +} + func (s *SSESession) TargetURL() string { return s.sseURL } // connect opens the persistent GET SSE connection on first use and blocks From 93f5f7568868ef4c9e66ed1e91fc942d2273331e Mon Sep 17 00:00:00 2001 From: sanjayy0612 Date: Sun, 16 Aug 2026 18:23:40 +0530 Subject: [PATCH 4/6] Enhance mcpWebSocketDetector to support host:port candidates and add corresponding tests --- internal/discovery/websocket_detector.go | 74 ++++++++++++++----- internal/discovery/websocket_detector_test.go | 20 +++++ 2 files changed, 75 insertions(+), 19 deletions(-) diff --git a/internal/discovery/websocket_detector.go b/internal/discovery/websocket_detector.go index d9b255f..f971234 100644 --- a/internal/discovery/websocket_detector.go +++ b/internal/discovery/websocket_detector.go @@ -3,6 +3,7 @@ package discovery import ( "context" "net" + "strconv" "github.com/hackwither/reap/internal/probe/common" ) @@ -27,29 +28,64 @@ func (d *mcpWebSocketDetector) Kinds() []CandidateKind { } func (d *mcpWebSocketDetector) Detect(ctx context.Context, c Candidate, opts DetectOptions) (*Fingerprint, error) { - if c.URL == "" { - return nil, nil - } dialer := &net.Dialer{Timeout: opts.Timeout} var headers map[string]string if opts.AuthHeader != "" { headers = map[string]string{"Authorization": opts.AuthHeader} } - conn, _, err := common.DialWebSocket(dialer, c.URL, headers) - if err != nil { - return nil, nil // unreachable candidate, or not a conformant WebSocket server — not a Detector failure + for _, targetURL := range websocketCandidateURLs(c) { + conn, _, err := common.DialWebSocket(dialer, targetURL, headers) + if err != nil { + continue // unreachable candidate, or not a conformant WebSocket server + } + conn.Close() + + matchedCandidate := c + matchedCandidate.URL = targetURL + return &Fingerprint{ + Candidate: matchedCandidate, + Protocol: "mcp", + Transport: "websocket", + Confidence: "medium", // capped: non-standard transport, upgrade-only confirmation, no JSON-RPC round trip yet + Evidence: map[string]any{ + "upgrade_confirmed": true, + "note": "WebSocket is not part of the official MCP spec; this only confirms a conformant upgrade handshake, not an MCP JSON-RPC round trip", + }, + DetectorID: d.ID(), + }, nil + } + return nil, nil +} + +func websocketCandidateURLs(c Candidate) []string { + if c.Kind == KindURL { + if c.URL == "" { + return nil + } + return []string{c.URL} + } + if c.Kind != KindHostPort { + return nil + } + host, port := c.Host, c.Port + if host == "" && c.RawInput != "" { + if parsedHost, parsedPort, err := net.SplitHostPort(c.RawInput); err == nil { + host = parsedHost + if port == 0 { + port, _ = strconv.Atoi(parsedPort) + } + } else { + host = c.RawInput + } + } + if port != 0 { + host = net.JoinHostPort(host, strconv.Itoa(port)) + } + var urls []string + for _, scheme := range []string{"http", "https"} { + for _, path := range MCPWellKnownPaths { + urls = append(urls, scheme+"://"+host+path) + } } - defer conn.Close() - - return &Fingerprint{ - Candidate: c, - Protocol: "mcp", - Transport: "websocket", - Confidence: "medium", // capped: non-standard transport, upgrade-only confirmation, no JSON-RPC round trip yet - Evidence: map[string]any{ - "upgrade_confirmed": true, - "note": "WebSocket is not part of the official MCP spec; this only confirms a conformant upgrade handshake, not an MCP JSON-RPC round trip", - }, - DetectorID: d.ID(), - }, nil + return urls } diff --git a/internal/discovery/websocket_detector_test.go b/internal/discovery/websocket_detector_test.go index 5bbf977..925f5f1 100644 --- a/internal/discovery/websocket_detector_test.go +++ b/internal/discovery/websocket_detector_test.go @@ -49,6 +49,26 @@ func TestWebSocketDetector_MatchesConformantUpgrade(t *testing.T) { } } +func TestWebSocketDetector_MatchesHostPortCandidate(t *testing.T) { + srv := wsUpgradeServer() + defer srv.Close() + + det := &mcpWebSocketDetector{} + fp, err := det.Detect(context.Background(), Candidate{ + Kind: KindHostPort, + RawInput: srv.URL[len("http://"):], + }, DetectOptions{Timeout: 5 * time.Second}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if fp == nil { + t.Fatal("expected a fingerprint match against a host:port WS candidate, got nil") + } + if fp.Candidate.URL == "" { + t.Fatal("expected fingerprint to record the URL that upgraded") + } +} + // TestWebSocketDetector_NoFalsePositiveOnPlainHTTP is the false-positive // discipline check: a server that never upgrades at all must not match. func TestWebSocketDetector_NoFalsePositiveOnPlainHTTP(t *testing.T) { From 91c6cedd02dd114bfe2ff613ee53ed18f452c3b0 Mon Sep 17 00:00:00 2001 From: sanjayy0612 Date: Sun, 16 Aug 2026 18:38:06 +0530 Subject: [PATCH 5/6] Add closeSession function to release transport state and improve session management --- internal/probe/mcp/checks.go | 12 +++++++ internal/probe/mcp/checks_test.go | 53 ++++++++++++++++++++++++------- 2 files changed, 53 insertions(+), 12 deletions(-) diff --git a/internal/probe/mcp/checks.go b/internal/probe/mcp/checks.go index 72b7484..b0376d0 100644 --- a/internal/probe/mcp/checks.go +++ b/internal/probe/mcp/checks.go @@ -71,6 +71,15 @@ func anonymousSession(s probe.Session) (probe.Session, error) { return provider.AnonymousSession() } +// closeSession releases persistent transport state when a probe created a +// separate anonymous session. Streamable HTTP sessions need no explicit +// cleanup, while WebSocket and legacy-SSE sessions keep connections open. +func closeSession(s probe.Session) { + if closer, ok := s.(io.Closer); ok { + _ = closer.Close() + } +} + // streamableHTTPOnly is for probes that depend on a mechanism specific to // the streamable-HTTP session implementation (e.g. the Mcp-Session-Id // response header it captures) that has no equivalent in legacy-SSE or @@ -380,6 +389,7 @@ func (p *oauthMetadataPostureProbe) Run(ctx context.Context, s probe.Session, r if err != nil { return nil } + defer closeSession(unauthSess) unauthRaw, unauthErr := unauthSess.Do(ctx, "tools/list", map[string]any{}) sawChallenge := unauthErr == nil && unauthRaw != nil && unauthRaw.StatusCode == http.StatusUnauthorized if sawChallenge { @@ -753,6 +763,7 @@ func (p *unauthToolsListProbe) Run(ctx context.Context, s probe.Session, r *repo if err != nil { return nil } + defer closeSession(unauthSess) raw, err := unauthSess.Do(ctx, "tools/list", map[string]any{}) if err != nil { return nil // network failure is not a finding; leave silent, CLI logs errors separately @@ -1108,6 +1119,7 @@ func (p *resourcesPromptsExposureProbe) Run(ctx context.Context, s probe.Session if err != nil { return nil } + defer closeSession(unauthSess) for _, method := range []string{"resources/list", "prompts/list"} { raw, err := unauthSess.Do(ctx, method, map[string]any{}) if err != nil || raw.StatusCode != 200 { diff --git a/internal/probe/mcp/checks_test.go b/internal/probe/mcp/checks_test.go index 1b52d00..ba8a474 100644 --- a/internal/probe/mcp/checks_test.go +++ b/internal/probe/mcp/checks_test.go @@ -3,6 +3,7 @@ package mcp import ( "bufio" "context" + "encoding/json" "fmt" "net/http" "net/http/httptest" @@ -22,19 +23,35 @@ func authRequiredWSServer(t *testing.T) *httptest.Server { return } key := r.Header.Get("Sec-WebSocket-Key") - hj, _ := w.(http.Hijacker) + hj, ok := w.(http.Hijacker) + if !ok { + http.Error(w, "webserver does not support hijacking", http.StatusInternalServerError) + return + } conn, buf, err := hj.Hijack() - if err != nil { return } + if err != nil { + http.Error(w, "websocket hijack failed", http.StatusInternalServerError) + return + } defer conn.Close() buf.WriteString("HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: " + common.ExpectedWebSocketAccept(key) + "\r\n\r\n") _ = buf.Flush() reader := bufio.NewReader(buf) for { fin, opcode, payload, err := readWSFrame(reader) - if err != nil { return } - if !fin || opcode != wsOpText { continue } - var req struct { ID int `json:"id"`; Method string `json:"method"` } - if json.Unmarshal(payload, &req) != nil { continue } + if err != nil { + return + } + if !fin || opcode != wsOpText { + continue + } + var req struct { + ID int `json:"id"` + Method string `json:"method"` + } + if json.Unmarshal(payload, &req) != nil { + continue + } result := map[string]any{} switch req.Method { case "initialize": @@ -43,7 +60,9 @@ func authRequiredWSServer(t *testing.T) *httptest.Server { result = map[string]any{"tools": []map[string]any{{"name": "secret_tool"}}} } resp, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": req.ID, "result": result}) - if writeServerWSFrame(conn, wsOpText, resp) != nil { return } + if writeServerWSFrame(conn, wsOpText, resp) != nil { + return + } } })) } @@ -53,19 +72,29 @@ func TestUnauthToolsListProbeDoesNotReuseAuthenticatedWebSocket(t *testing.T) { defer srv.Close() wsURL := "ws" + srv.URL[len("http"):] sess, err := NewWSSession(wsURL, "Bearer secret", 5*time.Second) - if err != nil { t.Fatalf("authenticated websocket handshake failed: %v", err) } + if err != nil { + t.Fatalf("authenticated websocket handshake failed: %v", err) + } defer sess.conn.Close() ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() init, _, err := InitializeSession(ctx, sess) - if err != nil || init == nil { t.Fatalf("authenticated initialize failed: %v", err) } + if err != nil || init == nil { + t.Fatalf("authenticated initialize failed: %v", err) + } raw, err := sess.Do(ctx, "tools/list", map[string]any{}) - if err != nil || raw.StatusCode != http.StatusOK { t.Fatalf("authenticated tools/list failed: status=%d err=%v", raw.StatusCode, err) } + if err != nil || raw.StatusCode != http.StatusOK { + t.Fatalf("authenticated tools/list failed: status=%d err=%v", raw.StatusCode, err) + } rep := &report.Report{Target: report.Target{URL: wsURL, Protocol: "mcp"}} - if err := (&unauthToolsListProbe{}).Run(ctx, sess, rep); err != nil { t.Fatalf("probe failed: %v", err) } - if len(rep.Findings) != 0 { t.Fatalf("authenticated websocket response was incorrectly reported as anonymous exposure: %+v", rep.Findings) } + if err := (&unauthToolsListProbe{}).Run(ctx, sess, rep); err != nil { + t.Fatalf("probe failed: %v", err) + } + if len(rep.Findings) != 0 { + t.Fatalf("authenticated websocket response was incorrectly reported as anonymous exposure: %+v", rep.Findings) + } } type fakeSession struct { From 79fc5405a52b5c15af2a9e06c6ae5e6570381618 Mon Sep 17 00:00:00 2001 From: sanjayy0612 Date: Sun, 16 Aug 2026 18:38:32 +0530 Subject: [PATCH 6/6] Implement Close methods for SSESession and WSSession to manage connection termination --- internal/probe/mcp/session_sse.go | 30 ++++++++++++++++++++- internal/probe/mcp/session_sse_test.go | 36 ++++++++++++++++++++++++++ internal/probe/mcp/session_ws.go | 5 ++++ internal/probe/mcp/session_ws_test.go | 20 ++++++++++++++ 4 files changed, 90 insertions(+), 1 deletion(-) diff --git a/internal/probe/mcp/session_sse.go b/internal/probe/mcp/session_sse.go index af952d6..6ddf40a 100644 --- a/internal/probe/mcp/session_sse.go +++ b/internal/probe/mcp/session_sse.go @@ -16,6 +16,7 @@ import ( "bytes" "context" "encoding/json" + "errors" "fmt" "io" "net/http" @@ -40,6 +41,8 @@ type SSESession struct { mu sync.Mutex postURL string connErr error + closed bool + cancel context.CancelFunc reqID int pending map[int]chan *probe.RawResult } @@ -64,6 +67,22 @@ func (s *SSESession) AnonymousSession() (probe.Session, error) { return NewSSESession(s.sseURL, "", s.timeout), nil } +// Close stops the persistent SSE stream. It is safe to call more than once. +func (s *SSESession) Close() error { + s.mu.Lock() + if s.closed { + s.mu.Unlock() + return nil + } + s.closed = true + cancel := s.cancel + s.mu.Unlock() + if cancel != nil { + cancel() + } + return nil +} + func (s *SSESession) TargetURL() string { return s.sseURL } // connect opens the persistent GET SSE connection on first use and blocks @@ -72,7 +91,16 @@ func (s *SSESession) TargetURL() string { return s.sseURL } // first caller does any work, everyone else just waits on s.ready. func (s *SSESession) connect(ctx context.Context) error { s.connOnce.Do(func() { - go s.runStream(ctx) + s.mu.Lock() + if s.closed { + s.mu.Unlock() + s.failConnect(errors.New("SSE session is closed")) + return + } + streamCtx, cancel := context.WithCancel(ctx) + s.cancel = cancel + s.mu.Unlock() + go s.runStream(streamCtx) }) select { case <-s.ready: diff --git a/internal/probe/mcp/session_sse_test.go b/internal/probe/mcp/session_sse_test.go index 4af95be..1b68e77 100644 --- a/internal/probe/mcp/session_sse_test.go +++ b/internal/probe/mcp/session_sse_test.go @@ -130,3 +130,39 @@ func TestSSESession_SynchronousPOSTResponsePreferred(t *testing.T) { t.Fatalf("expected the synchronous POST response body to be used, got %s", raw.Body) } } + +func TestSSESession_CloseStopsStream(t *testing.T) { + streamClosed := make(chan struct{}) + mux := http.NewServeMux() + mux.HandleFunc("/sse", func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, "event: endpoint\ndata: /messages\n\n") + w.(http.Flusher).Flush() + <-r.Context().Done() + close(streamClosed) + }) + mux.HandleFunc("/messages", func(w http.ResponseWriter, r *http.Request) { + var body struct { + ID int `json:"id"` + } + _ = json.NewDecoder(r.Body).Decode(&body) + _ = json.NewEncoder(w).Encode(map[string]any{"jsonrpc": "2.0", "id": body.ID, "result": map[string]any{}}) + }) + srv := httptest.NewServer(mux) + defer srv.Close() + + sess := NewSSESession(srv.URL+"/sse", "", 5*time.Second) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if _, err := sess.Do(ctx, "initialize", map[string]any{}); err != nil { + t.Fatalf("Do failed: %v", err) + } + if err := sess.Close(); err != nil { + t.Fatalf("Close failed: %v", err) + } + select { + case <-streamClosed: + case <-time.After(time.Second): + t.Fatal("Close did not stop the SSE stream") + } +} diff --git a/internal/probe/mcp/session_ws.go b/internal/probe/mcp/session_ws.go index 58203c8..608df6d 100644 --- a/internal/probe/mcp/session_ws.go +++ b/internal/probe/mcp/session_ws.go @@ -79,6 +79,11 @@ func (s *WSSession) AnonymousSession() (probe.Session, error) { return NewWSSession(s.targetURL, "", s.timeout) } +// Close releases the upgraded connection and stops the background reader. +func (s *WSSession) Close() error { + return s.conn.Close() +} + func (s *WSSession) TargetURL() string { return s.targetURL } // Do sends one JSON-RPC request as a text frame and waits for a response diff --git a/internal/probe/mcp/session_ws_test.go b/internal/probe/mcp/session_ws_test.go index de5bc33..fc28140 100644 --- a/internal/probe/mcp/session_ws_test.go +++ b/internal/probe/mcp/session_ws_test.go @@ -152,3 +152,23 @@ func TestWSSession_MultipleSequentialRequests(t *testing.T) { } } } + +func TestWSSession_ClosePreventsFurtherRequests(t *testing.T) { + srv := wsEchoServer(t) + defer srv.Close() + wsURL := "ws" + srv.URL[len("http"):] + + sess, err := NewWSSession(wsURL, "", 5*time.Second) + if err != nil { + t.Fatalf("NewWSSession failed: %v", err) + } + if err := sess.Close(); err != nil { + t.Fatalf("Close failed: %v", err) + } + + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + if _, err := sess.Do(ctx, "tools/list", map[string]any{}); err == nil { + t.Fatal("Do succeeded after Close") + } +}