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) { diff --git a/internal/probe/mcp/checks.go b/internal/probe/mcp/checks.go index 71e6e71..b0376d0 100644 --- a/internal/probe/mcp/checks.go +++ b/internal/probe/mcp/checks.go @@ -59,6 +59,27 @@ 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() +} + +// 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 @@ -364,7 +385,12 @@ 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 + } + defer closeSession(unauthSess) + 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 +757,14 @@ 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 + } + 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 } @@ -1085,8 +1115,13 @@ 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 + } + defer closeSession(unauthSess) 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..ba8a474 100644 --- a/internal/probe/mcp/checks_test.go +++ b/internal/probe/mcp/checks_test.go @@ -1,7 +1,9 @@ package mcp import ( + "bufio" "context" + "encoding/json" "fmt" "net/http" "net/http/httptest" @@ -9,9 +11,92 @@ 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, 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 { + 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 + } + 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 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..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" @@ -31,6 +32,7 @@ import ( type SSESession struct { sseURL string authHeader string + timeout time.Duration httpClient *http.Client connOnce sync.Once @@ -39,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 } @@ -47,6 +51,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 +61,28 @@ 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 +} + +// 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 @@ -64,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 3aa1190..608df6d 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,17 @@ 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) +} + +// 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") + } +}