diff options
Diffstat (limited to 'internal/proxy/websocket_test.go')
| -rw-r--r-- | internal/proxy/websocket_test.go | 764 |
1 files changed, 764 insertions, 0 deletions
diff --git a/internal/proxy/websocket_test.go b/internal/proxy/websocket_test.go new file mode 100644 index 0000000..6ea6024 --- /dev/null +++ b/internal/proxy/websocket_test.go @@ -0,0 +1,764 @@ +package proxy_test + +import ( + "bufio" + "bytes" + "context" + "crypto/tls" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/gorilla/websocket" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// ── helpers ─────────────────────────────────────────────────────────────────── + +// wsUpstreamServer starts a TLS WebSocket echo server. +func wsUpstreamServer(t *testing.T) *httptest.Server { + t.Helper() + upgrader := websocket.Upgrader{ + CheckOrigin: func(r *http.Request) bool { return true }, + } + srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + return + } + defer conn.Close() + for { + msgType, payload, err := conn.ReadMessage() + if err != nil { + return + } + if err := conn.WriteMessage(msgType, payload); err != nil { + return + } + } + })) + t.Cleanup(srv.Close) + return srv +} + +// wsUpstreamServerWithHandler starts a TLS WebSocket server with a custom handler. +func wsUpstreamServerWithHandler(t *testing.T, handler func(*websocket.Conn)) *httptest.Server { + t.Helper() + upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }} + srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + return + } + defer conn.Close() + handler(conn) + })) + t.Cleanup(srv.Close) + return srv +} + +// wsClientOverTunnel dials a WebSocket through the proxy CONNECT tunnel. +func wsClientOverTunnel(t *testing.T, proxyAddr, targetHost string) *websocket.Conn { + t.Helper() + conn := doConnectTunnel(t, proxyAddr, targetHost) + + tlsConn := tls.Client(conn, &tls.Config{ + ServerName: strings.Split(targetHost, ":")[0], + InsecureSkipVerify: true, //nolint:gosec + NextProtos: []string{"http/1.1"}, + }) + require.NoError(t, tlsConn.Handshake()) + t.Cleanup(func() { tlsConn.Close() }) + + dialer := websocket.Dialer{ + NetDialTLSContext: func(ctx context.Context, network, addr string) (net.Conn, error) { + return tlsConn, nil + }, + } + wsConn, _, err := dialer.Dial("wss://"+targetHost+"/", nil) + require.NoError(t, err) + t.Cleanup(func() { wsConn.Close() }) + return wsConn +} + +// sendHTTPKeepAlive writes an HTTP/1.1 request with Connection: keep-alive +// and reads the response, leaving the connection open for reuse. +func sendHTTPKeepAlive(t *testing.T, conn io.ReadWriter, method, path, host string) *http.Response { + t.Helper() + var req strings.Builder + fmt.Fprintf(&req, "%s %s HTTP/1.1\r\n", method, path) + fmt.Fprintf(&req, "Host: %s\r\n", host) + req.WriteString("Connection: keep-alive\r\n\r\n") + _, err := io.WriteString(conn, req.String()) + require.NoError(t, err) + resp, err := http.ReadResponse(bufio.NewReader(conn), nil) + require.NoError(t, err) + return resp +} +func wsClientOverTunnelWithPath(t *testing.T, proxyAddr, targetHost, path string) *websocket.Conn { + t.Helper() + conn := doConnectTunnel(t, proxyAddr, targetHost) + tlsConn := tls.Client(conn, &tls.Config{ + ServerName: strings.Split(targetHost, ":")[0], + InsecureSkipVerify: true, //nolint:gosec + NextProtos: []string{"http/1.1"}, + }) + require.NoError(t, tlsConn.Handshake()) + t.Cleanup(func() { tlsConn.Close() }) + + dialer := websocket.Dialer{ + NetDialTLSContext: func(ctx context.Context, network, addr string) (net.Conn, error) { + return tlsConn, nil + }, + } + wsConn, _, err := dialer.Dial("wss://"+targetHost+path, nil) + require.NoError(t, err) + t.Cleanup(func() { wsConn.Close() }) + return wsConn +} + +// waitForFrames polls until the store has at least n frames for the first +// connection, or fails after 3 seconds. +func waitForFrames(t *testing.T, st interface { + AllWebSocketConnections() (interface{ Len() int }, error) +}, n int) { + t.Helper() + // implemented inline in each test for flexibility +} + +// ── basic correctness ───────────────────────────────────────────────────────── + +func TestMITM_WebSocket_EchoesMessages(t *testing.T) { + st := newTestStore(t) + upstream := wsUpstreamServer(t) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + wsConn := wsClientOverTunnel(t, proxyAddr, targetHost) + + require.NoError(t, wsConn.WriteMessage(websocket.TextMessage, []byte("hello"))) + msgType, payload, err := wsConn.ReadMessage() + require.NoError(t, err) + assert.Equal(t, websocket.TextMessage, msgType) + assert.Equal(t, "hello", string(payload)) +} + +func TestMITM_WebSocket_LogsConnection(t *testing.T) { + st := newTestStore(t) + upstream := wsUpstreamServer(t) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + wsConn := wsClientOverTunnel(t, proxyAddr, targetHost) + defer wsConn.Close() + + require.NoError(t, wsConn.WriteMessage(websocket.TextMessage, []byte("ping"))) + _, _, err = wsConn.ReadMessage() + require.NoError(t, err) + + deadline := time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + loaded, err := st.AllWebSocketConnections() + require.NoError(t, err) + if len(loaded) > 0 { + assert.Equal(t, targetHost, loaded[0].Host) + assert.True(t, loaded[0].TLS) + return + } + time.Sleep(50 * time.Millisecond) + } + t.Fatal("timed out waiting for WebSocket connection to be logged") +} + +func TestMITM_WebSocket_LogsFrames(t *testing.T) { + st := newTestStore(t) + upstream := wsUpstreamServer(t) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + wsConn := wsClientOverTunnel(t, proxyAddr, targetHost) + defer wsConn.Close() + + for _, msg := range []string{"frame1", "frame2"} { + require.NoError(t, wsConn.WriteMessage(websocket.TextMessage, []byte(msg))) + _, _, err := wsConn.ReadMessage() + require.NoError(t, err) + } + + deadline := time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + conns, err := st.AllWebSocketConnections() + require.NoError(t, err) + if len(conns) == 0 { + time.Sleep(50 * time.Millisecond) + continue + } + frames, err := st.FramesForConnection(conns[0].ID) + require.NoError(t, err) + if len(frames) >= 4 { + assert.Equal(t, 4, len(frames)) + return + } + time.Sleep(50 * time.Millisecond) + } + t.Fatal("timed out waiting for WebSocket frames to be logged") +} + +// ── hijackWriter / gorilla coupling canary ──────────────────────────────────── + +// TestMITM_WebSocket_HijackCalledByUpgrader verifies that gorilla's Upgrader +// calls Hijack() on our hijackWriter. If gorilla ever changes this internal +// behaviour the test will fail before it hits production. +func TestMITM_WebSocket_HijackCalledByUpgrader(t *testing.T) { + hijackCalled := false + + // trackingWriter wraps a real http.ResponseWriter and records Hijack calls. + type trackingWriter struct { + http.ResponseWriter + conn net.Conn + brw *bufio.ReadWriter + called *bool + } + tw := &trackingWriter{} + tw.called = &hijackCalled + + // We test this via a real proxy round-trip — if the WebSocket connection + // succeeds, Hijack() was called. If gorilla stops calling Hijack(), the + // upgrade fails and the test fails. + st := newTestStore(t) + upstream := wsUpstreamServer(t) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + + // If the connection succeeds, Hijack was called correctly. + wsConn := wsClientOverTunnel(t, proxyAddr, targetHost) + defer wsConn.Close() + + require.NoError(t, wsConn.WriteMessage(websocket.TextMessage, []byte("hijack-test"))) + _, payload, err := wsConn.ReadMessage() + require.NoError(t, err) + assert.Equal(t, "hijack-test", string(payload), "gorilla Upgrader must use Hijack to take over the connection") +} + +// ── buffered reader integrity ───────────────────────────────────────────────── + +// TestMITM_WebSocket_BufferedReaderNotLost verifies that the proxy correctly +// uses the existing bufio.Reader when handing off to the WebSocket handler. +// This is tested indirectly: if browserReader were discarded and a new one +// created over the raw conn, the WS framing would desync on high-traffic +// connections. We verify correctness by sending many rapid messages immediately +// after connection, which exercises the read buffer aggressively. +func TestMITM_WebSocket_BufferedReaderNotLost(t *testing.T) { + st := newTestStore(t) + upstream := wsUpstreamServer(t) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + wsConn := wsClientOverTunnel(t, proxyAddr, targetHost) + defer wsConn.Close() + + // Send 50 messages immediately without waiting for echo. + // If the bufio.Reader is recreated (losing buffered bytes), framing + // desyncs and ReadMessage returns an error. + const n = 50 + done := make(chan error, 1) + go func() { + for i := 0; i < n; i++ { + if err := wsConn.WriteMessage(websocket.TextMessage, []byte("buf")); err != nil { + done <- err + return + } + } + done <- nil + }() + + wsConn.SetReadDeadline(time.Now().Add(10 * time.Second)) + for i := 0; i < n; i++ { + _, payload, err := wsConn.ReadMessage() + require.NoError(t, err, "frame %d: bufio.Reader must not lose bytes", i) + assert.Equal(t, "buf", string(payload)) + } + require.NoError(t, <-done) +} + +// ── message types ───────────────────────────────────────────────────────────── + +func TestMITM_WebSocket_BinaryMessage(t *testing.T) { + st := newTestStore(t) + upstream := wsUpstreamServer(t) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + wsConn := wsClientOverTunnel(t, proxyAddr, targetHost) + defer wsConn.Close() + + payload := []byte{0x00, 0x01, 0x02, 0xFE, 0xFF} + require.NoError(t, wsConn.WriteMessage(websocket.BinaryMessage, payload)) + msgType, received, err := wsConn.ReadMessage() + require.NoError(t, err) + assert.Equal(t, websocket.BinaryMessage, msgType) + assert.Equal(t, payload, received) +} + +func TestMITM_WebSocket_EmptyPayload(t *testing.T) { + st := newTestStore(t) + upstream := wsUpstreamServer(t) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + wsConn := wsClientOverTunnel(t, proxyAddr, targetHost) + defer wsConn.Close() + + require.NoError(t, wsConn.WriteMessage(websocket.TextMessage, []byte(""))) + msgType, received, err := wsConn.ReadMessage() + require.NoError(t, err) + assert.Equal(t, websocket.TextMessage, msgType) + assert.Equal(t, []byte{}, received) +} + +func TestMITM_WebSocket_LargePayload(t *testing.T) { + st := newTestStore(t) + upstream := wsUpstreamServer(t) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + wsConn := wsClientOverTunnel(t, proxyAddr, targetHost) + defer wsConn.Close() + + // 512KB — well under wsFrameSizeLimit but large enough to exercise buffering. + payload := bytes.Repeat([]byte("A"), 512*1024) + require.NoError(t, wsConn.WriteMessage(websocket.TextMessage, payload)) + _, received, err := wsConn.ReadMessage() + require.NoError(t, err) + assert.Equal(t, payload, received) +} + +func TestMITM_WebSocket_PayloadAtSizeLimit(t *testing.T) { + st := newTestStore(t) + upstream := wsUpstreamServer(t) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + wsConn := wsClientOverTunnel(t, proxyAddr, targetHost) + defer wsConn.Close() + + // Exactly at wsFrameSizeLimit (1MB) — should proxy fine, and logged payload + // should be exactly wsFrameSizeLimit bytes (no truncation needed). + payload := bytes.Repeat([]byte("B"), 1<<20) + require.NoError(t, wsConn.WriteMessage(websocket.TextMessage, payload)) + _, received, err := wsConn.ReadMessage() + require.NoError(t, err) + assert.Equal(t, len(payload), len(received), "payload at size limit must be proxied intact") +} + +// ── message ordering ────────────────────────────────────────────────────────── + +func TestMITM_WebSocket_MessageOrderingPreserved(t *testing.T) { + st := newTestStore(t) + upstream := wsUpstreamServer(t) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + wsConn := wsClientOverTunnel(t, proxyAddr, targetHost) + defer wsConn.Close() + + messages := []string{"first", "second", "third", "fourth", "fifth"} + for _, msg := range messages { + require.NoError(t, wsConn.WriteMessage(websocket.TextMessage, []byte(msg))) + } + for _, expected := range messages { + _, payload, err := wsConn.ReadMessage() + require.NoError(t, err) + assert.Equal(t, expected, string(payload), "messages must arrive in send order") + } +} + +// ── frame logging correctness ───────────────────────────────────────────────── + +func TestMITM_WebSocket_FrameDirectionsCorrect(t *testing.T) { + st := newTestStore(t) + upstream := wsUpstreamServer(t) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + wsConn := wsClientOverTunnel(t, proxyAddr, targetHost) + defer wsConn.Close() + + require.NoError(t, wsConn.WriteMessage(websocket.TextMessage, []byte("dirtest"))) + _, _, err = wsConn.ReadMessage() + require.NoError(t, err) + + deadline := time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + conns, err := st.AllWebSocketConnections() + require.NoError(t, err) + if len(conns) == 0 { + time.Sleep(50 * time.Millisecond) + continue + } + frames, err := st.FramesForConnection(conns[0].ID) + require.NoError(t, err) + if len(frames) >= 2 { + // First frame must be client→server, second server→client. + assert.Equal(t, 0, int(frames[0].Direction), "first frame must be client→server") + assert.Equal(t, 1, int(frames[1].Direction), "second frame must be server→client") + return + } + time.Sleep(50 * time.Millisecond) + } + t.Fatal("timed out waiting for frames with correct directions") +} + +func TestMITM_WebSocket_FramePayloadsLogged(t *testing.T) { + st := newTestStore(t) + upstream := wsUpstreamServer(t) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + wsConn := wsClientOverTunnel(t, proxyAddr, targetHost) + defer wsConn.Close() + + require.NoError(t, wsConn.WriteMessage(websocket.TextMessage, []byte("payload-check"))) + _, _, err = wsConn.ReadMessage() + require.NoError(t, err) + + deadline := time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + conns, err := st.AllWebSocketConnections() + require.NoError(t, err) + if len(conns) == 0 { + time.Sleep(50 * time.Millisecond) + continue + } + frames, err := st.FramesForConnection(conns[0].ID) + require.NoError(t, err) + if len(frames) >= 1 { + assert.Equal(t, "payload-check", string(frames[0].Payload)) + return + } + time.Sleep(50 * time.Millisecond) + } + t.Fatal("timed out waiting for frame payload to be logged") +} + +// ── connection lifecycle ────────────────────────────────────────────────────── + +func TestMITM_WebSocket_CleanClose(t *testing.T) { + st := newTestStore(t) + upstream := wsUpstreamServer(t) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + wsConn := wsClientOverTunnel(t, proxyAddr, targetHost) + + require.NoError(t, wsConn.WriteMessage(websocket.TextMessage, []byte("bye"))) + _, payload, err := wsConn.ReadMessage() + require.NoError(t, err) + assert.Equal(t, "bye", string(payload)) + + err = wsConn.WriteMessage(websocket.CloseMessage, + websocket.FormatCloseMessage(websocket.CloseNormalClosure, "done")) + require.NoError(t, err) + wsConn.Close() +} + +func TestMITM_WebSocket_UpstreamDropsConnection(t *testing.T) { + st := newTestStore(t) + + // Upstream closes the connection after receiving the first message. + upstream := wsUpstreamServerWithHandler(t, func(conn *websocket.Conn) { + _, _, _ = conn.ReadMessage() + // Close without sending a close frame — simulates abrupt upstream drop. + conn.Close() + }) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + wsConn := wsClientOverTunnel(t, proxyAddr, targetHost) + defer wsConn.Close() + + require.NoError(t, wsConn.WriteMessage(websocket.TextMessage, []byte("drop-me"))) + + // Browser side should get an error when upstream drops. + wsConn.SetReadDeadline(time.Now().Add(3 * time.Second)) + _, _, err = wsConn.ReadMessage() + assert.Error(t, err, "browser conn must get an error when upstream drops") +} + +func TestMITM_WebSocket_BrowserDropsConnection(t *testing.T) { + st := newTestStore(t) + + received := make(chan struct{}) + upstream := wsUpstreamServerWithHandler(t, func(conn *websocket.Conn) { + // Signal that upstream is ready, then try to read — should fail when + // browser drops. + close(received) + conn.SetReadDeadline(time.Now().Add(3 * time.Second)) + _, _, _ = conn.ReadMessage() + }) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + wsConn := wsClientOverTunnel(t, proxyAddr, targetHost) + + // Wait for upstream to be ready, then abruptly close browser side. + <-received + wsConn.Close() + + // Upstream read should fail — the proxy must propagate the close. + // The test passes if it completes without hanging (the upstream handler returns). +} + +// ── concurrent connections ──────────────────────────────────────────────────── + +func TestMITM_WebSocket_ConcurrentConnections(t *testing.T) { + st := newTestStore(t) + upstream := wsUpstreamServer(t) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + + const nConns = 5 + var wg sync.WaitGroup + errors := make(chan error, nConns) + + for i := 0; i < nConns; i++ { + wg.Add(1) + go func(id int) { + defer wg.Done() + wsConn := wsClientOverTunnel(t, proxyAddr, targetHost) + defer wsConn.Close() + + msg := []byte("concurrent") + if err := wsConn.WriteMessage(websocket.TextMessage, msg); err != nil { + errors <- err + return + } + _, payload, err := wsConn.ReadMessage() + if err != nil { + errors <- err + return + } + if string(payload) != string(msg) { + errors <- nil + } + }(i) + } + + wg.Wait() + close(errors) + for err := range errors { + require.NoError(t, err) + } +} + +// ── HTTP and WebSocket through same proxy ───────────────────────────────────── + +// TestMITM_WebSocket_AfterHTTPOnSameProxy verifies that the proxy correctly +// handles both HTTP and WebSocket connections. HTTP and WS use separate CONNECT +// tunnels (which is how browsers work — each is a separate TCP connection). +func TestMITM_WebSocket_AfterHTTPOnSameProxy(t *testing.T) { + st := newTestStore(t) + upstream := wsUpstreamServer(t) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + + // First: a plain HTTPS request through the proxy. + httpConn := doConnectTunnel(t, proxyAddr, targetHost) + httpTLS := tlsOverTunnel(t, httpConn, host) + resp := sendHTTPOverConn(t, httpTLS, "GET", "/", targetHost, "") + resp.Body.Close() + // The upstream is a WS-only server so it won't return 200, + // but the proxy must forward the request and get a response. + assert.NotZero(t, resp.StatusCode) + + // Then: a WebSocket connection through the same proxy. + wsConn := wsClientOverTunnel(t, proxyAddr, targetHost) + defer wsConn.Close() + + require.NoError(t, wsConn.WriteMessage(websocket.TextMessage, []byte("after-http"))) + _, payload, err := wsConn.ReadMessage() + require.NoError(t, err) + assert.Equal(t, "after-http", string(payload), "WebSocket must work correctly after a prior HTTP request through the same proxy") +} + +// ── URL and path preservation ───────────────────────────────────────────────── + +func TestMITM_WebSocket_PathPreserved(t *testing.T) { + st := newTestStore(t) + + receivedPath := make(chan string, 1) + var wsUpgrader = websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }} + srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + receivedPath <- r.URL.Path + "?" + r.URL.RawQuery + conn, err := wsUpgrader.Upgrade(w, r, nil) + if err != nil { + return + } + conn.Close() + })) + t.Cleanup(srv.Close) + + upstreamAddr := strings.TrimPrefix(srv.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + _ = wsClientOverTunnelWithPath(t, proxyAddr, targetHost, "/chat/room?token=abc123") + + select { + case path := <-receivedPath: + assert.Equal(t, "/chat/room?token=abc123", path, "path and query must be preserved through proxy") + case <-time.After(3 * time.Second): + t.Fatal("timed out waiting for upstream to receive request") + } +} + +// ── store URL logged correctly ──────────────────────────────────────────────── + +func TestMITM_WebSocket_URLLoggedCorrectly(t *testing.T) { + st := newTestStore(t) + upstream := wsUpstreamServer(t) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + wsConn := wsClientOverTunnelWithPath(t, proxyAddr, targetHost, "/api/stream") + defer wsConn.Close() + + require.NoError(t, wsConn.WriteMessage(websocket.TextMessage, []byte("url-test"))) + _, _, err = wsConn.ReadMessage() + require.NoError(t, err) + + deadline := time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + conns, err := st.AllWebSocketConnections() + require.NoError(t, err) + if len(conns) > 0 { + assert.Contains(t, conns[0].URL, "/api/stream", "URL must be logged with path") + assert.Contains(t, conns[0].URL, "wss://", "URL must use wss scheme") + return + } + time.Sleep(50 * time.Millisecond) + } + t.Fatal("timed out waiting for connection URL to be logged") +} + +// ── rapid fire messages ─────────────────────────────────────────────────────── + +func TestMITM_WebSocket_RapidFireMessages(t *testing.T) { + st := newTestStore(t) + upstream := wsUpstreamServer(t) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + wsConn := wsClientOverTunnel(t, proxyAddr, targetHost) + defer wsConn.Close() + + const n = 100 + var sent atomic.Int64 + + // Send 100 messages without waiting for echo. + go func() { + for range n { + if err := wsConn.WriteMessage(websocket.TextMessage, []byte("rapid")); err != nil { + return + } + sent.Add(1) + } + }() + + // Receive all 100 echoes. + wsConn.SetReadDeadline(time.Now().Add(10 * time.Second)) + for range n { + _, payload, err := wsConn.ReadMessage() + require.NoError(t, err) + assert.Equal(t, "rapid", string(payload)) + } +} |