diff options
Diffstat (limited to 'internal/proxy')
| -rw-r--r-- | internal/proxy/forward.go | 29 | ||||
| -rw-r--r-- | internal/proxy/forward_test.go | 197 | ||||
| -rw-r--r-- | internal/proxy/mitm.go | 214 | ||||
| -rw-r--r-- | internal/proxy/mitm2.go | 211 | ||||
| -rw-r--r-- | internal/proxy/mitm2_test.go | 280 | ||||
| -rw-r--r-- | internal/proxy/mitm_test.go | 281 | ||||
| -rw-r--r-- | internal/proxy/proxy.go | 200 | ||||
| -rw-r--r-- | internal/proxy/proxy_test.go | 218 | ||||
| -rw-r--r-- | internal/proxy/recover.go | 13 | ||||
| -rw-r--r-- | internal/proxy/websocket.go | 225 | ||||
| -rw-r--r-- | internal/proxy/websocket_test.go | 764 |
11 files changed, 2632 insertions, 0 deletions
diff --git a/internal/proxy/forward.go b/internal/proxy/forward.go new file mode 100644 index 0000000..7e40d8e --- /dev/null +++ b/internal/proxy/forward.go @@ -0,0 +1,29 @@ +package proxy + +import ( + "fmt" + "net/http" + + internalhttp "github.com/shiv/internal/http" +) + +func forward(req *http.Request, client *http.Client) (*http.Response, error) { + out, err := http.NewRequestWithContext(req.Context(), req.Method, req.URL.String(), req.Body) + if err != nil { + return nil, fmt.Errorf("forward: build request: %w", err) + } + for k, vals := range req.Header { + if internalhttp.HopByHop[k] { + continue + } + out.Header[k] = vals + } + internalhttp.StripRequestCacheHeaders(out.Header) + + resp, err := client.Do(out) + if err != nil { + return nil, fmt.Errorf("forward: upstream request: %w", err) + } + internalhttp.StripResponseCacheHeaders(resp.Header) + return resp, nil +} diff --git a/internal/proxy/forward_test.go b/internal/proxy/forward_test.go new file mode 100644 index 0000000..88d4db4 --- /dev/null +++ b/internal/proxy/forward_test.go @@ -0,0 +1,197 @@ +package proxy + +import ( + "bytes" + "compress/flate" + "compress/gzip" + "net/http" + "testing" + + "github.com/andybalholm/brotli" + "github.com/klauspost/compress/zstd" + internalhttp "github.com/shiv/internal/http" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func decompressBody(h http.Header, body []byte) []byte { return internalhttp.Decompress(h, body) } +func isBinary(h http.Header) bool { return internalhttp.IsBinary(h) } +func stripRequestCacheHeaders(h http.Header) { internalhttp.StripRequestCacheHeaders(h) } +func stripResponseCacheHeaders(h http.Header) { internalhttp.StripResponseCacheHeaders(h) } + +// ── decompressBody ──────────────────────────────────────────────────────────── + +func TestDecompressBody_NoEncoding(t *testing.T) { + body := []byte("plain text") + out := decompressBody(http.Header{}, body) + assert.Equal(t, body, out) +} + +func TestDecompressBody_Gzip(t *testing.T) { + var buf bytes.Buffer + w := gzip.NewWriter(&buf) + _, err := w.Write([]byte("gzip content")) + require.NoError(t, err) + w.Close() + + h := http.Header{"Content-Encoding": []string{"gzip"}} + out := decompressBody(h, buf.Bytes()) + assert.Equal(t, "gzip content", string(out)) +} + +func TestDecompressBody_Gzip_CorruptData(t *testing.T) { + body := []byte("not gzip") + h := http.Header{"Content-Encoding": []string{"gzip"}} + out := decompressBody(h, body) + assert.Equal(t, body, out, "corrupt gzip should return original bytes") +} + +func TestDecompressBody_Deflate(t *testing.T) { + var buf bytes.Buffer + w, err := flate.NewWriter(&buf, flate.DefaultCompression) + require.NoError(t, err) + _, err = w.Write([]byte("deflate content")) + require.NoError(t, err) + w.Close() + + h := http.Header{"Content-Encoding": []string{"deflate"}} + out := decompressBody(h, buf.Bytes()) + assert.Equal(t, "deflate content", string(out)) +} + +func TestDecompressBody_Brotli(t *testing.T) { + var buf bytes.Buffer + w := brotli.NewWriter(&buf) + _, err := w.Write([]byte("brotli content")) + require.NoError(t, err) + w.Close() + + h := http.Header{"Content-Encoding": []string{"br"}} + out := decompressBody(h, buf.Bytes()) + assert.Equal(t, "brotli content", string(out)) +} + +func TestDecompressBody_Zstd(t *testing.T) { + enc, err := zstd.NewWriter(nil) + require.NoError(t, err) + compressed := enc.EncodeAll([]byte("zstd content"), nil) + + h := http.Header{"Content-Encoding": []string{"zstd"}} + out := decompressBody(h, compressed) + assert.Equal(t, "zstd content", string(out)) +} + +func TestDecompressBody_UnknownEncoding(t *testing.T) { + body := []byte("mystery") + h := http.Header{"Content-Encoding": []string{"identity"}} + out := decompressBody(h, body) + assert.Equal(t, body, out) +} + +// ── isBinary ───────────────────────────────────────────────────────────────── + +func TestIsBinary_PlainText(t *testing.T) { + h := http.Header{"Content-Type": []string{"text/plain"}} + assert.False(t, isBinary(h)) +} + +func TestIsBinary_HTML(t *testing.T) { + h := http.Header{"Content-Type": []string{"text/html; charset=utf-8"}} + assert.False(t, isBinary(h)) +} + +func TestIsBinary_JSON(t *testing.T) { + h := http.Header{"Content-Type": []string{"application/json"}} + assert.False(t, isBinary(h)) +} + +func TestIsBinary_XML(t *testing.T) { + h := http.Header{"Content-Type": []string{"application/xml"}} + assert.False(t, isBinary(h)) +} + +func TestIsBinary_JavaScript(t *testing.T) { + h := http.Header{"Content-Type": []string{"application/javascript"}} + assert.False(t, isBinary(h)) +} + +func TestIsBinary_FormEncoded(t *testing.T) { + h := http.Header{"Content-Type": []string{"application/x-www-form-urlencoded"}} + assert.False(t, isBinary(h)) +} + +func TestIsBinary_Image(t *testing.T) { + h := http.Header{"Content-Type": []string{"image/png"}} + assert.True(t, isBinary(h)) +} + +func TestIsBinary_OctetStream(t *testing.T) { + h := http.Header{"Content-Type": []string{"application/octet-stream"}} + assert.True(t, isBinary(h)) +} + +func TestIsBinary_NoContentType(t *testing.T) { + assert.False(t, isBinary(http.Header{})) +} + +func TestIsBinary_ZstdEncoding(t *testing.T) { + h := http.Header{ + "Content-Encoding": []string{"zstd"}, + "Content-Type": []string{"text/plain"}, + } + assert.False(t, isBinary(h)) +} + +// ── stripRequestCacheHeaders ────────────────────────────────────────────────── + +func TestStripRequestCacheHeaders(t *testing.T) { + h := http.Header{ + "If-None-Match": []string{`"etag123"`}, + "If-Modified-Since": []string{"Mon, 01 Jan 2024 00:00:00 GMT"}, + "If-Range": []string{`"etag123"`}, + "If-Match": []string{`"etag123"`}, + "If-Unmodified-Since": []string{"Mon, 01 Jan 2024 00:00:00 GMT"}, + "Authorization": []string{"Bearer token"}, + } + stripRequestCacheHeaders(h) + + assert.Empty(t, h.Get("If-None-Match")) + assert.Empty(t, h.Get("If-Modified-Since")) + assert.Empty(t, h.Get("If-Range")) + assert.Empty(t, h.Get("If-Match")) + assert.Empty(t, h.Get("If-Unmodified-Since")) + assert.Equal(t, "Bearer token", h.Get("Authorization")) +} + +func TestStripRequestCacheHeaders_AlreadyAbsent(t *testing.T) { + h := http.Header{"Content-Type": []string{"application/json"}} + assert.NotPanics(t, func() { stripRequestCacheHeaders(h) }) + assert.Equal(t, "application/json", h.Get("Content-Type")) +} + +// ── stripResponseCacheHeaders ───────────────────────────────────────────────── + +func TestStripResponseCacheHeaders(t *testing.T) { + h := http.Header{ + "Cache-Control": []string{"max-age=3600"}, + "Expires": []string{"Mon, 01 Jan 2025 00:00:00 GMT"}, + "ETag": []string{`"etag123"`}, + "Last-Modified": []string{"Mon, 01 Jan 2024 00:00:00 GMT"}, + "Pragma": []string{"no-cache"}, + "Content-Type": []string{"text/html"}, + } + stripResponseCacheHeaders(h) + + assert.Equal(t, "no-store, no-cache, must-revalidate", h.Get("Cache-Control")) + assert.Empty(t, h.Get("Expires")) + assert.Empty(t, h.Get("ETag")) + assert.Empty(t, h.Get("Last-Modified")) + assert.Empty(t, h.Get("Pragma")) + assert.Equal(t, "text/html", h.Get("Content-Type")) +} + +func TestStripResponseCacheHeaders_SetsNoCacheDirective(t *testing.T) { + h := http.Header{} + stripResponseCacheHeaders(h) + assert.Equal(t, "no-store, no-cache, must-revalidate", h.Get("Cache-Control")) +} diff --git a/internal/proxy/mitm.go b/internal/proxy/mitm.go new file mode 100644 index 0000000..5ed63a1 --- /dev/null +++ b/internal/proxy/mitm.go @@ -0,0 +1,214 @@ +package proxy + +import ( + "bufio" + "bytes" + "context" + "crypto/tls" + "errors" + "fmt" + "io" + "net" + "net/http" + "time" + + "github.com/shiv/internal/events" + internalhttp "github.com/shiv/internal/http" + "github.com/shiv/internal/logger" +) + +func (p *Proxy) handleConnect(w http.ResponseWriter, r *http.Request) { + defer recoverPanic("ServeHTTP " + r.URL.String()) + w.WriteHeader(http.StatusOK) + if flusher, ok := w.(http.Flusher); ok { + flusher.Flush() + } + + hijacker, ok := w.(http.Hijacker) + if !ok { + logger.Error("mitm: server does not support hijacking") + return + } + clientConn, _, err := hijacker.Hijack() + if err != nil { + logger.Error("mitm: hijack: %v", err) + return + } + p.trackConn(clientConn) + defer p.untrackConn(clientConn) + defer clientConn.Close() + + bareHost, _, err := net.SplitHostPort(r.Host) + if err != nil { + bareHost = r.Host + } + + tlsCert, err := p.certAuth.TLSCertForHost(bareHost) + if err != nil { + logger.Error("mitm: cert for %s: %v", bareHost, err) + return + } + + browserTLS := tls.Server(clientConn, &tls.Config{ + Certificates: []tls.Certificate{*tlsCert}, + NextProtos: []string{"h2", "http/1.1"}, + }) + if err := browserTLS.Handshake(); err != nil { + logger.Debug("mitm: browser TLS handshake for %s: %v", bareHost, err) + return + } + defer browserTLS.Close() + logger.Debug("mitm: negotiated protocol for %s: %q", bareHost, browserTLS.ConnectionState().NegotiatedProtocol) + + if browserTLS.ConnectionState().NegotiatedProtocol == "h2" { + p.handleConnectH2(browserTLS, r, bareHost) + return + } + + transport := &http.Transport{ + DialTLSContext: func(ctx context.Context, network, addr string) (net.Conn, error) { + conn, err := (&net.Dialer{Timeout: 10 * time.Second}).DialContext(ctx, "tcp", addr) + if err != nil { + return nil, fmt.Errorf("mitm: dial %s: %w", addr, err) + } + tlsConn := tls.Client(conn, &tls.Config{ + ServerName: bareHost, + InsecureSkipVerify: true, //nolint:gosec + }) + if err := tlsConn.Handshake(); err != nil { + conn.Close() + return nil, fmt.Errorf("mitm: upstream TLS handshake for %s: %w", bareHost, err) + } + return tlsConn, nil + }, + ResponseHeaderTimeout: 30 * time.Second, + } + defer transport.CloseIdleConnections() + + // No client-level Timeout — a per-connection timeout would terminate the + // entire keep-alive session after the first long request. Individual + // requests are bounded by ResponseHeaderTimeout + read deadlines instead. + upstreamClient := &http.Client{ + Transport: transport, + CheckRedirect: func(req *http.Request, via []*http.Request) error { + return http.ErrUseLastResponse + }, + } + + browserReader := bufio.NewReader(browserTLS) + for { + req, err := http.ReadRequest(browserReader) + if err != nil { + logger.Debug("mitm: connection closed for %s: %v", bareHost, err) + return + } + + start := time.Now() + + reqBody, err := io.ReadAll(io.LimitReader(req.Body, maxBodySize)) + if err != nil { + logger.Error("mitm: read request body for %s: %v", bareHost, err) + return + } + req.Body.Close() + + req.URL.Scheme = "https" + // Preserve r.Host (with port) so non-443 HTTPS targets work correctly. + req.URL.Host = r.Host + + // Detect WebSocket upgrade before intercept. WebSocket connections are + // long-lived bidirectional streams — they cannot be intercepted the same + // way as regular HTTP requests. Hand off to the WebSocket handler which + // manages the entire session and returns when the connection closes. + if isWebSocketUpgrade(req) { + p.handleWebSocketTLS(browserTLS, browserReader, req, bareHost, r.Host) + return + } + + result := p.bus.EmitRequest(events.RequestEvent{Request: req, Body: reqBody}) + if result.Drop { + resp := &http.Response{ + Status: "403 Forbidden", + StatusCode: 403, + Proto: "HTTP/1.1", + ProtoMajor: 1, + ProtoMinor: 1, + Body: http.NoBody, + Header: make(http.Header), + ContentLength: 0, + } + if err := resp.Write(browserTLS); err != nil { + logger.Error("mitm: write 403 to browser for %s: %v", bareHost, err) + } + continue + } + interceptedReq, reqBody := result.Request, result.Body + + outReq, err := http.NewRequest(interceptedReq.Method, interceptedReq.URL.String(), bytes.NewReader(reqBody)) + if err != nil { + logger.Error("mitm: build upstream request for %s: %v", bareHost, err) + return + } + outReq.Header = interceptedReq.Header.Clone() + outReq.Host = internalhttp.NormalizeHost(bareHost, true) + internalhttp.StripRequestCacheHeaders(outReq.Header) + + resp, err := upstreamClient.Do(outReq) + if err != nil { + // Distinguish client-cancelled requests (normal) from real upstream + // failures. Context cancellation here means the browser closed the + // connection before we could forward — log at Debug. + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + logger.Debug("mitm: stream cancelled for %s: %v", bareHost, err) + } else { + logger.Error("mitm: upstream request for %s: %v", bareHost, err) + } + return + } + + var respBuf bytes.Buffer + reader := io.TeeReader(io.LimitReader(resp.Body, maxBodySize), &respBuf) + + internalhttp.StripResponseCacheHeaders(resp.Header) + resp.Body = io.NopCloser(reader) + + writeErr := resp.Write(browserTLS) + _, _ = io.Copy(io.Discard, resp.Body) + resp.Body.Close() + + if writeErr != nil { + logger.Debug("mitm: write response to browser for %s: %v", bareHost, writeErr) + return + } + + respBody := respBuf.Bytes() + + if len(respBody) >= int(maxBodySize) { + logger.Debug("mitm: response body truncated at %d bytes for %s %s", + maxBodySize, interceptedReq.Method, interceptedReq.URL) + } + + elapsed := time.Since(start).Milliseconds() + logger.Info("%s %s %d %db %dms", interceptedReq.Method, interceptedReq.URL, resp.StatusCode, len(respBody), elapsed) + + logBody := internalhttp.Decompress(resp.Header, respBody) + if internalhttp.IsBinary(resp.Header) { + logBody = nil + } + + p.bus.EmitResponse(events.ResponseEvent{ + Timestamp: start, + Host: r.Host, + Proto: "HTTP/1.1", + Method: interceptedReq.Method, + URL: interceptedReq.URL.String(), + ReqHeaders: interceptedReq.Header, + ReqBody: reqBody, + StatusCode: resp.StatusCode, + RespHeaders: resp.Header, + RespBody: logBody, + DurationMs: elapsed, + TLS: true, + }) + } +} diff --git a/internal/proxy/mitm2.go b/internal/proxy/mitm2.go new file mode 100644 index 0000000..fcf8e0f --- /dev/null +++ b/internal/proxy/mitm2.go @@ -0,0 +1,211 @@ +package proxy + +import ( + "bytes" + "context" + "crypto/tls" + "crypto/x509" + "errors" + "fmt" + "io" + "net" + "net/http" + "time" + + "golang.org/x/net/http2" + + "github.com/shiv/internal/events" + internalhttp "github.com/shiv/internal/http" + "github.com/shiv/internal/logger" +) + +const ( + h2UpstreamDialTimeout = 10 * time.Second + h2UpstreamHeaderTimeout = 30 * time.Second + // h2StreamRequestTimeout is a per-stream deadline, not a connection-level + // timeout. Using a context timeout per stream instead of a client-level + // Timeout means a single slow stream cannot kill the entire H2 connection. + h2StreamRequestTimeout = 60 * time.Second + // h2MaxConcurrentStreams matches the HTTP/2 spec default (RFC 7540 §6.5.2). + h2MaxConcurrentStreams = 250 +) + +func (p *Proxy) handleConnectH2(browserTLS *tls.Conn, connectReq *http.Request, bareHost string) { + defer recoverPanic("handleConnectH2 " + bareHost) + + // Use the port-aware host from the original CONNECT request for upstream + // dialling so non-443 HTTPS ports are preserved. + upstreamHost := connectReq.Host + + upstreamTransport := &http.Transport{ + DialTLSContext: func(ctx context.Context, network, addr string) (net.Conn, error) { + conn, err := (&net.Dialer{Timeout: h2UpstreamDialTimeout}).DialContext(ctx, "tcp", addr) + if err != nil { + return nil, fmt.Errorf("mitm/h2: dial %s: %w", addr, err) + } + tlsConn := tls.Client(conn, &tls.Config{ + ServerName: bareHost, + InsecureSkipVerify: true, //nolint:gosec + NextProtos: []string{"h2", "http/1.1"}, + VerifyPeerCertificate: func(rawCerts [][]byte, _ [][]*x509.Certificate) error { + if len(rawCerts) == 0 { + return nil + } + cert, err := x509.ParseCertificate(rawCerts[0]) + if err != nil { + return nil + } + now := time.Now() + if now.After(cert.NotAfter) { + logger.Info("mitm/h2: EXPIRED cert for %s (expired %s, issuer: %s)", + bareHost, cert.NotAfter.Format("2006-01-02"), cert.Issuer.CommonName) + } else if now.Before(cert.NotBefore) { + logger.Info("mitm/h2: NOT YET VALID cert for %s (valid from %s, issuer: %s)", + bareHost, cert.NotBefore.Format("2006-01-02"), cert.Issuer.CommonName) + } + if err := cert.VerifyHostname(bareHost); err != nil { + logger.Info("mitm/h2: hostname mismatch for %s: %v (CN: %s, SANs: %v)", + bareHost, err, cert.Subject.CommonName, cert.DNSNames) + } + return nil + }, + }) + if err := tlsConn.Handshake(); err != nil { + conn.Close() + return nil, fmt.Errorf("mitm/h2: upstream TLS handshake for %s: %w", addr, err) + } + return tlsConn, nil + }, + ResponseHeaderTimeout: h2UpstreamHeaderTimeout, + } + if err := http2.ConfigureTransport(upstreamTransport); err != nil { + logger.Error("mitm/h2: configure upstream transport: %v", err) + } + defer upstreamTransport.CloseIdleConnections() + + // No client-level Timeout — that would kill the entire H2 connection on + // any slow stream. Per-stream timeouts are applied via request contexts. + upstreamClient := &http.Client{ + Transport: upstreamTransport, + CheckRedirect: func(req *http.Request, via []*http.Request) error { + return http.ErrUseLastResponse + }, + } + + streamHandler := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { + defer recoverPanic("h2 stream " + bareHost) + + req.URL.Scheme = "https" + req.URL.Host = upstreamHost + + reqBody, err := io.ReadAll(io.LimitReader(req.Body, maxBodySize)) + if err != nil { + logger.Error("mitm/h2: read request body for %s: %v", bareHost, err) + http.Error(w, "bad request", http.StatusBadRequest) + return + } + req.Body.Close() + + start := time.Now() + + result := p.bus.EmitRequest(events.RequestEvent{Request: req, Body: reqBody}) + if result.Drop { + // Stream is dropped by intercept rules; the H2 connection stays alive. + http.Error(w, "request dropped", http.StatusForbidden) + return + } + interceptedReq, reqBody := result.Request, result.Body + + streamCtx, cancel := context.WithTimeout(req.Context(), h2StreamRequestTimeout) + defer cancel() + + outReq, err := http.NewRequestWithContext(streamCtx, interceptedReq.Method, interceptedReq.URL.String(), bytes.NewReader(reqBody)) + if err != nil { + logger.Error("mitm/h2: build upstream request for %s: %v", bareHost, err) + http.Error(w, "internal error", http.StatusInternalServerError) + return + } + outReq.Header = interceptedReq.Header.Clone() + outReq.Host = internalhttp.NormalizeHost(bareHost, true) + // Strip H2-illegal headers before forwarding. + outReq.Header.Del("Transfer-Encoding") + internalhttp.StripRequestCacheHeaders(outReq.Header) + + resp, err := upstreamClient.Do(outReq) + if err != nil { + // context.Canceled means the browser cancelled the stream (RST_STREAM) + // or we hit our per-stream deadline. This is normal browser behaviour + // (prefetch, autocomplete, navigation away) — log at Debug, not Error. + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + logger.Debug("mitm/h2: stream cancelled/timed out for %s: %v", bareHost, err) + } else { + logger.Error("mitm/h2: upstream request for %s: %v", bareHost, err) + } + http.Error(w, "bad gateway", http.StatusBadGateway) + return + } + + var respBuf bytes.Buffer + reader := io.TeeReader(io.LimitReader(resp.Body, maxBodySize), &respBuf) + + // Strip Transfer-Encoding — illegal in HTTP/2 responses (RFC 7540 §8.1.2.2). + // Strip caching headers so the browser always fetches fresh content. + resp.Header.Del("Transfer-Encoding") + internalhttp.StripResponseCacheHeaders(resp.Header) + + for k, vals := range resp.Header { + for _, v := range vals { + w.Header().Add(k, v) + } + } + w.WriteHeader(resp.StatusCode) + + if _, err := io.Copy(w, reader); err != nil { + logger.Debug("mitm/h2: write response to browser for %s: %v", bareHost, err) + // Drain whatever remains into respBuf so the log body is as complete + // as possible even if the browser disconnected mid-stream. + _, _ = io.Copy(&respBuf, resp.Body) + } + + respBody := respBuf.Bytes() + resp.Body.Close() + + if len(respBody) >= int(maxBodySize) { + logger.Debug("mitm/h2: response body truncated at %d bytes for %s %s", + maxBodySize, interceptedReq.Method, interceptedReq.URL) + } + + elapsed := time.Since(start).Milliseconds() + logger.Info("%s %s %d %db %dms (h2)", interceptedReq.Method, interceptedReq.URL, resp.StatusCode, len(respBody), elapsed) + + logBody := internalhttp.Decompress(resp.Header, respBody) + if internalhttp.IsBinary(resp.Header) { + logBody = nil + } + + p.bus.EmitResponse(events.ResponseEvent{ + Timestamp: start, + Host: connectReq.Host, + Proto: "HTTP/2", + Method: interceptedReq.Method, + URL: interceptedReq.URL.String(), + ReqHeaders: interceptedReq.Header, + ReqBody: reqBody, + StatusCode: resp.StatusCode, + RespHeaders: resp.Header, + RespBody: logBody, + DurationMs: elapsed, + TLS: true, + }) + }) + + h2srv := &http2.Server{ + MaxConcurrentStreams: h2MaxConcurrentStreams, + } + // ServeConn blocks until the browser closes the underlying TLS connection. + // It handles all stream multiplexing internally. No idle timer is needed — + // the connection lifetime is naturally tied to browserTLS. + h2srv.ServeConn(browserTLS, &http2.ServeConnOpts{ + Handler: streamHandler, + }) +} diff --git a/internal/proxy/mitm2_test.go b/internal/proxy/mitm2_test.go new file mode 100644 index 0000000..a1d4c40 --- /dev/null +++ b/internal/proxy/mitm2_test.go @@ -0,0 +1,280 @@ +package proxy_test + +import ( + "context" + "crypto/tls" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "golang.org/x/net/http2" + + "github.com/shiv/internal/store" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// h2UpstreamServer starts an HTTP/2-enabled TLS test server. +func h2UpstreamServer(t *testing.T, handler http.Handler) *httptest.Server { + t.Helper() + srv := httptest.NewUnstartedServer(handler) + srv.TLS = &tls.Config{ + NextProtos: []string{"h2", "http/1.1"}, + } + h2s := &http2.Server{} + require.NoError(t, http2.ConfigureServer(srv.Config, h2s)) + srv.StartTLS() + t.Cleanup(srv.Close) + return srv +} + +// h2ClientOverTunnel wraps a raw conn in a TLS client that negotiates HTTP/2 +// and returns an *http.Client that routes all requests over that connection. +func h2ClientOverTunnel(t *testing.T, conn net.Conn, serverName string) *http.Client { + t.Helper() + tlsConn := tls.Client(conn, &tls.Config{ + ServerName: serverName, + InsecureSkipVerify: true, //nolint:gosec — test only + NextProtos: []string{"h2"}, + }) + require.NoError(t, tlsConn.Handshake()) + assert.Equal(t, "h2", tlsConn.ConnectionState().NegotiatedProtocol, "proxy must negotiate h2 with the client") + t.Cleanup(func() { tlsConn.Close() }) + + transport := &http2.Transport{ + DialTLSContext: func(ctx context.Context, network, addr string, cfg *tls.Config) (net.Conn, error) { + return tlsConn, nil + }, + TLSClientConfig: &tls.Config{InsecureSkipVerify: true}, //nolint:gosec + } + + return &http.Client{ + Transport: transport, + CheckRedirect: func(req *http.Request, via []*http.Request) error { + return http.ErrUseLastResponse + }, + } +} + +// ── tests ───────────────────────────────────────────────────────────────────── + +func TestMITM_H2_ForwardsGET(t *testing.T) { + st := newTestStore(t) + + upstream := h2UpstreamServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + io.WriteString(w, "h2 works") + })) + + 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) + conn := doConnectTunnel(t, proxyAddr, targetHost) + client := h2ClientOverTunnel(t, conn, host) + + resp, err := client.Get("https://" + targetHost + "/") + require.NoError(t, err) + defer resp.Body.Close() + + assert.Equal(t, http.StatusOK, resp.StatusCode) + b, _ := io.ReadAll(resp.Body) + assert.Equal(t, "h2 works", string(b)) +} + +func TestMITM_H2_LogsTransaction(t *testing.T) { + st := newTestStore(t) + + upstream := h2UpstreamServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + io.WriteString(w, "logged") + })) + + 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) + conn := doConnectTunnel(t, proxyAddr, targetHost) + client := h2ClientOverTunnel(t, conn, host) + + resp, err := client.Get("https://" + targetHost + "/log-test") + require.NoError(t, err) + resp.Body.Close() + + tx := waitForTransaction(t, st) + assert.True(t, tx.TLS) + assert.Equal(t, "GET", tx.Method) + assert.Equal(t, http.StatusOK, tx.StatusCode) +} + +func TestMITM_H2_PostWithBody(t *testing.T) { + st := newTestStore(t) + + var receivedBody string + upstream := h2UpstreamServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + b, _ := io.ReadAll(r.Body) + receivedBody = string(b) + w.WriteHeader(http.StatusCreated) + })) + + 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) + conn := doConnectTunnel(t, proxyAddr, targetHost) + client := h2ClientOverTunnel(t, conn, host) + + body := `{"key":"value"}` + req, err := http.NewRequest(http.MethodPost, "https://"+targetHost+"/api", strings.NewReader(body)) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + + resp, err := client.Do(req) + require.NoError(t, err) + resp.Body.Close() + + assert.Equal(t, http.StatusCreated, resp.StatusCode) + assert.Equal(t, body, receivedBody) +} + +func TestMITM_H2_CacheHeadersStripped(t *testing.T) { + st := newTestStore(t) + + upstream := h2UpstreamServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("ETag", `"abc"`) + w.Header().Set("Cache-Control", "max-age=3600") + w.WriteHeader(http.StatusOK) + })) + + 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) + conn := doConnectTunnel(t, proxyAddr, targetHost) + client := h2ClientOverTunnel(t, conn, host) + + resp, err := client.Get("https://" + targetHost + "/") + require.NoError(t, err) + resp.Body.Close() + + assert.Empty(t, resp.Header.Get("ETag")) + assert.Equal(t, "no-store, no-cache, must-revalidate", resp.Header.Get("Cache-Control")) +} + +func TestMITM_H2_InterceptDrop_Returns403(t *testing.T) { + st := newTestStore(t) + + upstream := h2UpstreamServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + io.WriteString(w, "should not reach") + })) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + st.Intercept.SetEnabled(true) + go func() { + pending := <-st.Intercept.Queue() + pending.Reply <- store.Decision{Forward: false} + }() + + proxyAddr := startProxyServer(t, st) + conn := doConnectTunnel(t, proxyAddr, targetHost) + client := h2ClientOverTunnel(t, conn, host) + + resp, err := client.Get("https://" + targetHost + "/secret") + require.NoError(t, err) + resp.Body.Close() + + assert.Equal(t, http.StatusForbidden, resp.StatusCode) +} + +func TestMITM_H2_InterceptForward_ModifiesRequest(t *testing.T) { + st := newTestStore(t) + + var receivedHeader string + upstream := h2UpstreamServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + receivedHeader = r.Header.Get("X-Injected") + w.WriteHeader(http.StatusOK) + })) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + st.Intercept.SetEnabled(true) + go func() { + pending := <-st.Intercept.Queue() + pending.Request.Header.Set("X-Injected", "h2-injected") + pending.Reply <- store.Decision{ + Forward: true, + Request: pending.Request, + Body: pending.Body, + } + }() + + proxyAddr := startProxyServer(t, st) + conn := doConnectTunnel(t, proxyAddr, targetHost) + client := h2ClientOverTunnel(t, conn, host) + + resp, err := client.Get("https://" + targetHost + "/inject") + require.NoError(t, err) + resp.Body.Close() + + assert.Equal(t, http.StatusOK, resp.StatusCode) + assert.Equal(t, "h2-injected", receivedHeader) +} + +func TestMITM_H2_MultipleConcurrentStreams(t *testing.T) { + st := newTestStore(t) + + upstream := h2UpstreamServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + fmt.Fprintf(w, "path:%s", r.URL.Path) + })) + + 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) + conn := doConnectTunnel(t, proxyAddr, targetHost) + client := h2ClientOverTunnel(t, conn, host) + + results := make(chan string, 3) + for _, path := range []string{"/a", "/b", "/c"} { + p := path + go func() { + resp, err := client.Get("https://" + targetHost + p) + if err != nil { + results <- "error" + return + } + defer resp.Body.Close() + b, _ := io.ReadAll(resp.Body) + results <- string(b) + }() + } + + got := make(map[string]bool) + for i := 0; i < 3; i++ { + got[<-results] = true + } + assert.True(t, got["path:/a"]) + assert.True(t, got["path:/b"]) + assert.True(t, got["path:/c"]) +} diff --git a/internal/proxy/mitm_test.go b/internal/proxy/mitm_test.go new file mode 100644 index 0000000..c683c70 --- /dev/null +++ b/internal/proxy/mitm_test.go @@ -0,0 +1,281 @@ +package proxy_test + +import ( + "bufio" + "bytes" + "crypto/tls" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/shiv/internal/store" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// startProxyServer starts the proxy as a real TCP listener and returns its address. +// httptest.NewRecorder does not support hijacking so we need a real listener for CONNECT. +func startProxyServer(t *testing.T, st *store.Store) string { + t.Helper() + bus := newTestBus(t, st) + p := newTestProxy(t, bus) + + ln, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + + srv := &http.Server{Handler: p} + go srv.Serve(ln) + t.Cleanup(func() { srv.Close() }) + + return ln.Addr().String() +} + +// tlsUpstream starts an httptest TLS server. +func tlsUpstream(t *testing.T, status int, body string) *httptest.Server { + t.Helper() + srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(status) + io.WriteString(w, body) + })) + t.Cleanup(srv.Close) + return srv +} + +// doConnectTunnel sends CONNECT to the proxy and returns the raw conn after the 200 response. +func doConnectTunnel(t *testing.T, proxyAddr, targetHost string) net.Conn { + t.Helper() + conn, err := net.DialTimeout("tcp", proxyAddr, 5*time.Second) + require.NoError(t, err) + t.Cleanup(func() { conn.Close() }) + + fmt.Fprintf(conn, "CONNECT %s HTTP/1.1\r\nHost: %s\r\n\r\n", targetHost, targetHost) + + resp, err := http.ReadResponse(bufio.NewReader(conn), nil) + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode, "proxy must respond 200 to CONNECT") + + return conn +} + +// tlsOverTunnel wraps conn in a TLS client trusting the proxy's self-signed cert. +func tlsOverTunnel(t *testing.T, conn net.Conn, serverName string) *tls.Conn { + t.Helper() + tlsConn := tls.Client(conn, &tls.Config{ + ServerName: serverName, + InsecureSkipVerify: true, //nolint:gosec + }) + require.NoError(t, tlsConn.Handshake()) + t.Cleanup(func() { tlsConn.Close() }) + return tlsConn +} + +// sendHTTPOverConn writes a raw HTTP/1.1 request and reads the response. +func sendHTTPOverConn(t *testing.T, conn io.ReadWriter, method, path, host, body 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) + if body != "" { + fmt.Fprintf(&req, "Content-Length: %d\r\n", len(body)) + } + req.WriteString("Connection: close\r\n\r\n") + if body != "" { + req.WriteString(body) + } + + _, err := io.WriteString(conn, req.String()) + require.NoError(t, err) + + resp, err := http.ReadResponse(bufio.NewReader(conn), nil) + require.NoError(t, err) + return resp +} + +// waitForTransaction blocks until a transaction appears on st.Updates or times out. +func waitForTransaction(t *testing.T, st *store.Store) store.Transaction { + t.Helper() + select { + case tx := <-st.Updates: + return tx + case <-time.After(3 * time.Second): + t.Fatal("timed out waiting for transaction to be logged") + return store.Transaction{} + } +} + +// ── tests ───────────────────────────────────────────────────────────────────── + +func TestMITM_ConnectTunnel_ForwardsRequest(t *testing.T) { + st := newTestStore(t) + upstream := tlsUpstream(t, http.StatusOK, "mitm works") + + 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) + conn := doConnectTunnel(t, proxyAddr, targetHost) + tlsConn := tlsOverTunnel(t, conn, host) + + resp := sendHTTPOverConn(t, tlsConn, "GET", "/", targetHost, "") + defer resp.Body.Close() + + assert.Equal(t, http.StatusOK, resp.StatusCode) + b, _ := io.ReadAll(resp.Body) + assert.Equal(t, "mitm works", string(b)) +} + +func TestMITM_ConnectTunnel_LogsTransactionWithTLSTrue(t *testing.T) { + st := newTestStore(t) + upstream := tlsUpstream(t, http.StatusOK, "logged") + + 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) + conn := doConnectTunnel(t, proxyAddr, targetHost) + tlsConn := tlsOverTunnel(t, conn, host) + + resp := sendHTTPOverConn(t, tlsConn, "GET", "/log-test", targetHost, "") + resp.Body.Close() + + tx := waitForTransaction(t, st) + assert.True(t, tx.TLS, "transaction must be flagged as TLS") + assert.Equal(t, "GET", tx.Method) + assert.Equal(t, http.StatusOK, tx.StatusCode) +} + +func TestMITM_ConnectTunnel_CacheHeadersStripped(t *testing.T) { + st := newTestStore(t) + + upstream := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("ETag", `"abc"`) + w.Header().Set("Cache-Control", "max-age=3600") + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(upstream.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) + conn := doConnectTunnel(t, proxyAddr, targetHost) + tlsConn := tlsOverTunnel(t, conn, host) + + resp := sendHTTPOverConn(t, tlsConn, "GET", "/", targetHost, "") + resp.Body.Close() + + assert.Empty(t, resp.Header.Get("ETag")) + assert.Equal(t, "no-store, no-cache, must-revalidate", resp.Header.Get("Cache-Control")) +} + +func TestMITM_ConnectTunnel_InterceptDrop_Returns403(t *testing.T) { + st := newTestStore(t) + upstream := tlsUpstream(t, http.StatusOK, "should not reach") + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + st.Intercept.SetEnabled(true) + go func() { + pending := <-st.Intercept.Queue() + pending.Reply <- store.Decision{Forward: false} + }() + + proxyAddr := startProxyServer(t, st) + conn := doConnectTunnel(t, proxyAddr, targetHost) + tlsConn := tlsOverTunnel(t, conn, host) + + resp := sendHTTPOverConn(t, tlsConn, "GET", "/secret", targetHost, "") + resp.Body.Close() + + assert.Equal(t, http.StatusForbidden, resp.StatusCode) +} + +func TestMITM_ConnectTunnel_InterceptForward_ModifiesRequest(t *testing.T) { + st := newTestStore(t) + + var receivedHeader string + upstream := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + receivedHeader = r.Header.Get("X-Injected") + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(upstream.Close) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + st.Intercept.SetEnabled(true) + go func() { + pending := <-st.Intercept.Queue() + pending.Request.Header.Set("X-Injected", "mitm-injected") + pending.Reply <- store.Decision{ + Forward: true, + Request: pending.Request, + Body: pending.Body, + } + }() + + proxyAddr := startProxyServer(t, st) + conn := doConnectTunnel(t, proxyAddr, targetHost) + tlsConn := tlsOverTunnel(t, conn, host) + + resp := sendHTTPOverConn(t, tlsConn, "GET", "/inject", targetHost, "") + resp.Body.Close() + + assert.Equal(t, http.StatusOK, resp.StatusCode) + assert.Equal(t, "mitm-injected", receivedHeader) +} + +func TestMITM_ConnectTunnel_PostWithBody(t *testing.T) { + st := newTestStore(t) + + var receivedBody string + upstream := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + b, _ := io.ReadAll(r.Body) + receivedBody = string(b) + w.WriteHeader(http.StatusCreated) + })) + t.Cleanup(upstream.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) + conn := doConnectTunnel(t, proxyAddr, targetHost) + tlsConn := tlsOverTunnel(t, conn, host) + + body := `{"key":"value"}` + var req bytes.Buffer + fmt.Fprintf(&req, "POST /api HTTP/1.1\r\n") + fmt.Fprintf(&req, "Host: %s\r\n", targetHost) + fmt.Fprintf(&req, "Content-Type: application/json\r\n") + fmt.Fprintf(&req, "Content-Length: %d\r\n", len(body)) + fmt.Fprintf(&req, "Connection: close\r\n\r\n") + fmt.Fprintf(&req, "%s", body) + _, err = tlsConn.Write(req.Bytes()) + require.NoError(t, err) + + resp, err := http.ReadResponse(bufio.NewReader(tlsConn), nil) + require.NoError(t, err) + resp.Body.Close() + + assert.Equal(t, http.StatusCreated, resp.StatusCode) + assert.Equal(t, body, receivedBody) +} diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go new file mode 100644 index 0000000..3e9845a --- /dev/null +++ b/internal/proxy/proxy.go @@ -0,0 +1,200 @@ +package proxy + +import ( + "bytes" + "context" + "fmt" + "io" + "net" + "net/http" + "sync" + "time" + + "github.com/shiv/internal/cert" + "github.com/shiv/internal/events" + internalhttp "github.com/shiv/internal/http" + "github.com/shiv/internal/logger" +) + +type Proxy struct { + addr string + certAuth *cert.CA + bus *events.Bus + client *http.Client + + mu sync.Mutex + srv *http.Server + conns sync.Map // active hijacked net.Conn +} + +const maxBodySize = 10 << 20 + +func New(addr string, bus *events.Bus) (*Proxy, error) { + ca, err := cert.Load() + if err != nil { + return nil, fmt.Errorf("proxy: load CA: %w", err) + } + p := &Proxy{ + addr: addr, + certAuth: ca, + bus: bus, + client: &http.Client{ + Transport: &http.Transport{ + ResponseHeaderTimeout: 30 * time.Second, + }, + CheckRedirect: func(req *http.Request, via []*http.Request) error { + return http.ErrUseLastResponse + }, + }, + } + bus.Register(p) + return p, nil +} + +func (p *Proxy) CA() *cert.CA { + return p.certAuth +} + +func (p *Proxy) trackConn(c net.Conn) { + p.conns.Store(c, struct{}{}) +} + +func (p *Proxy) untrackConn(c net.Conn) { + p.conns.Delete(c) +} + +func (p *Proxy) closeAllConns() { + p.conns.Range(func(k, _ any) bool { + k.(net.Conn).SetDeadline(time.Now()) + return true + }) +} + +func (p *Proxy) Start(newAddr string) error { + ln, err := net.Listen("tcp", newAddr) + if err != nil { + return fmt.Errorf("proxy: listen on %s: %w", newAddr, err) + } + + newSrv := &http.Server{ + Handler: p, + ReadTimeout: 60 * time.Second, + WriteTimeout: 60 * time.Second, + IdleTimeout: 120 * time.Second, + } + + p.mu.Lock() + oldSrv := p.srv + p.addr = newAddr + p.srv = newSrv + p.mu.Unlock() + + if oldSrv != nil { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + _ = oldSrv.Shutdown(ctx) + } + p.closeAllConns() + + go func() { + if err := newSrv.Serve(ln); err != nil && err != http.ErrServerClosed { + logger.Error("proxy: Serve: %v", err) + } + }() + + return nil +} + +func (p *Proxy) Stop() { + p.mu.Lock() + srv := p.srv + p.mu.Unlock() + + if srv != nil { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + _ = srv.Shutdown(ctx) + + p.mu.Lock() + p.srv = nil + p.mu.Unlock() + } + + p.closeAllConns() +} + +// ObserveProxyRestart implements events.ProxyCommandObserver. +func (p *Proxy) ObserveProxyRestart(e events.ProxyRestartEvent) { + _ = p.Start(e.Addr) +} + +// ObserveProxyStop implements events.ProxyCommandObserver. +func (p *Proxy) ObserveProxyStop(_ events.ProxyStopEvent) { + p.Stop() +} + +func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodConnect { + p.handleConnect(w, r) + return + } + + reqBody, err := io.ReadAll(io.LimitReader(r.Body, maxBodySize)) + if err != nil { + http.Error(w, "failed to read request body", http.StatusBadRequest) + return + } + r.Body = io.NopCloser(bytes.NewReader(reqBody)) + + result := p.bus.EmitRequest(events.RequestEvent{Request: r, Body: reqBody}) + if result.Drop { + http.Error(w, "request dropped", http.StatusForbidden) + return + } + interceptedReq, reqBody := result.Request, result.Body + interceptedReq.Body = io.NopCloser(bytes.NewReader(reqBody)) + + start := time.Now() + + resp, err := forward(interceptedReq, p.client) + if err != nil { + http.Error(w, "bad gateway", http.StatusBadGateway) + return + } + defer resp.Body.Close() + + var respBuf bytes.Buffer + reader := io.TeeReader(io.LimitReader(resp.Body, maxBodySize), &respBuf) + + for k, vals := range resp.Header { + for _, v := range vals { + w.Header().Add(k, v) + } + } + w.WriteHeader(resp.StatusCode) + + _, _ = io.Copy(w, reader) + + respBody := respBuf.Bytes() + elapsed := time.Since(start).Milliseconds() + + logBody := internalhttp.Decompress(resp.Header, respBody) + if internalhttp.IsBinary(resp.Header) { + logBody = nil + } + + p.bus.EmitResponse(events.ResponseEvent{ + Timestamp: start, + Host: interceptedReq.Host, + Proto: "HTTP/1.1", + Method: interceptedReq.Method, + URL: interceptedReq.URL.String(), + ReqHeaders: interceptedReq.Header, + ReqBody: reqBody, + StatusCode: resp.StatusCode, + RespHeaders: resp.Header, + RespBody: logBody, + DurationMs: elapsed, + TLS: false, + }) +} diff --git a/internal/proxy/proxy_test.go b/internal/proxy/proxy_test.go new file mode 100644 index 0000000..ab5d69e --- /dev/null +++ b/internal/proxy/proxy_test.go @@ -0,0 +1,218 @@ +package proxy_test + +import ( + "io" + "net/http" + "net/http/httptest" + "os" + "strings" + "testing" + "time" + + "github.com/shiv/internal/events" + "github.com/shiv/internal/proxy" + "github.com/shiv/internal/store" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// newTestStore opens a temp SQLite store and registers cleanup. +func newTestStore(t *testing.T) *store.Store { + t.Helper() + f, err := os.CreateTemp("", "shiv-proxy-test-*.shiv") + require.NoError(t, err) + f.Close() + t.Cleanup(func() { os.Remove(f.Name()) }) + projectStore, err := store.Open(f.Name()) + require.NoError(t, err) + t.Cleanup(func() { projectStore.Close() }) + return projectStore +} + +// newTestBus creates a bus wired to the given store — mirrors production wiring. +func newTestBus(t *testing.T, st *store.Store) *events.Bus { + t.Helper() + bus := events.NewBus() + bus.Register(st.Intercept) + bus.Register(st) + return bus +} + +// newTestProxy creates a Proxy wired to an events.Bus. +// The proxy is wired as an http.Handler so we can call it via httptest without +// binding a real port. +func newTestProxy(t *testing.T, bus *events.Bus) *proxy.Proxy { + t.Helper() + p, err := proxy.New("127.0.0.1:0", bus) + require.NoError(t, err) + return p +} + +// upstreamServer starts a simple upstream that echoes status + body. +func upstreamServer(t *testing.T, status int, body string) *httptest.Server { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(status) + io.WriteString(w, body) + })) + t.Cleanup(srv.Close) + return srv +} + +// ── HTTP forwarding ─────────────────────────────────────────────────────────── + +func TestServeHTTP_ForwardsRequestAndReturnsResponse(t *testing.T) { + st := newTestStore(t) + upstream := upstreamServer(t, http.StatusOK, "hello proxy") + p := newTestProxy(t, newTestBus(t, st)) + + req := httptest.NewRequest(http.MethodGet, upstream.URL+"/test", nil) + req.RequestURI = upstream.URL + "/test" + w := httptest.NewRecorder() + + p.ServeHTTP(w, req) + + assert.Equal(t, http.StatusOK, w.Code) + assert.Equal(t, "hello proxy", w.Body.String()) +} + +func TestServeHTTP_ForwardsPostWithBody(t *testing.T) { + st := newTestStore(t) + + var receivedBody string + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + b, _ := io.ReadAll(r.Body) + receivedBody = string(b) + w.WriteHeader(http.StatusCreated) + })) + t.Cleanup(upstream.Close) + + p := newTestProxy(t, newTestBus(t, st)) + body := `{"name":"test"}` + req := httptest.NewRequest(http.MethodPost, upstream.URL+"/api", strings.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + p.ServeHTTP(w, req) + + assert.Equal(t, http.StatusCreated, w.Code) + assert.Equal(t, body, receivedBody) +} + +func TestServeHTTP_LogsTransactionToStore(t *testing.T) { + st := newTestStore(t) + upstream := upstreamServer(t, http.StatusOK, "logged") + p := newTestProxy(t, newTestBus(t, st)) + + req := httptest.NewRequest(http.MethodGet, upstream.URL+"/logged", nil) + w := httptest.NewRecorder() + p.ServeHTTP(w, req) + + // Give the store a moment to write (it's async via channel) + time.Sleep(50 * time.Millisecond) + + txs, err := st.TransactionsPage(0, store.TransactionFilter{ShowOutScope: true}) + require.NoError(t, err) + require.Len(t, txs, 1) + assert.Equal(t, "GET", txs[0].Method) + assert.Equal(t, http.StatusOK, txs[0].StatusCode) +} + +func TestServeHTTP_ResponseHeadersPassedThrough(t *testing.T) { + st := newTestStore(t) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("X-Custom-Header", "custom-value") + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(upstream.Close) + + p := newTestProxy(t, newTestBus(t, st)) + req := httptest.NewRequest(http.MethodGet, upstream.URL, nil) + w := httptest.NewRecorder() + p.ServeHTTP(w, req) + + assert.Equal(t, "custom-value", w.Header().Get("X-Custom-Header")) +} + +func TestServeHTTP_CacheHeadersStripped(t *testing.T) { + st := newTestStore(t) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("ETag", `"abc123"`) + w.Header().Set("Cache-Control", "max-age=3600") + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(upstream.Close) + + p := newTestProxy(t, newTestBus(t, st)) + req := httptest.NewRequest(http.MethodGet, upstream.URL, nil) + w := httptest.NewRecorder() + p.ServeHTTP(w, req) + + assert.Empty(t, w.Header().Get("ETag")) + assert.Equal(t, "no-store, no-cache, must-revalidate", w.Header().Get("Cache-Control")) +} + +func TestServeHTTP_UpstreamError_ReturnsBadGateway(t *testing.T) { + st := newTestStore(t) + p := newTestProxy(t, newTestBus(t, st)) + + // Point at a port that is not listening + req := httptest.NewRequest(http.MethodGet, "http://127.0.0.1:1/", nil) + w := httptest.NewRecorder() + p.ServeHTTP(w, req) + + assert.Equal(t, http.StatusBadGateway, w.Code) +} + +// ── Intercept gate ──────────────────────────────────────────────────────────── + +func TestServeHTTP_InterceptDrop_ReturnsForbidden(t *testing.T) { + st := newTestStore(t) + upstream := upstreamServer(t, http.StatusOK, "should not reach") + p := newTestProxy(t, newTestBus(t, st)) + + st.Intercept.SetEnabled(true) + + // Drain the intercept queue and drop the request in a goroutine + go func() { + pending := <-st.Intercept.Queue() + pending.Reply <- store.Decision{Forward: false} + }() + + req := httptest.NewRequest(http.MethodGet, upstream.URL, nil) + w := httptest.NewRecorder() + p.ServeHTTP(w, req) + + assert.Equal(t, http.StatusForbidden, w.Code) +} + +func TestServeHTTP_InterceptForward_ModifiedRequest(t *testing.T) { + st := newTestStore(t) + + var receivedHeader string + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + receivedHeader = r.Header.Get("X-Injected") + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(upstream.Close) + + p := newTestProxy(t, newTestBus(t, st)) + st.Intercept.SetEnabled(true) + + go func() { + pending := <-st.Intercept.Queue() + pending.Request.Header.Set("X-Injected", "intercepted") + pending.Reply <- store.Decision{ + Forward: true, + Request: pending.Request, + Body: pending.Body, + } + }() + + req := httptest.NewRequest(http.MethodGet, upstream.URL, nil) + w := httptest.NewRecorder() + p.ServeHTTP(w, req) + + assert.Equal(t, http.StatusOK, w.Code) + assert.Equal(t, "intercepted", receivedHeader) +} diff --git a/internal/proxy/recover.go b/internal/proxy/recover.go new file mode 100644 index 0000000..27338e9 --- /dev/null +++ b/internal/proxy/recover.go @@ -0,0 +1,13 @@ +package proxy + +import ( + "runtime/debug" + + "github.com/shiv/internal/logger" +) + +func recoverPanic(context string) { + if r := recover(); r != nil { + logger.Error("panic in %s: %v\n%s", context, r, debug.Stack()) + } +} diff --git a/internal/proxy/websocket.go b/internal/proxy/websocket.go new file mode 100644 index 0000000..b02470e --- /dev/null +++ b/internal/proxy/websocket.go @@ -0,0 +1,225 @@ +package proxy + +import ( + "bufio" + "crypto/tls" + "fmt" + "net" + "net/http" + "net/url" + "time" + + "github.com/gorilla/websocket" + + "github.com/shiv/internal/events" + "github.com/shiv/internal/logger" +) + +// isWebSocketUpgrade returns true if the request is a WebSocket upgrade. +func isWebSocketUpgrade(req *http.Request) bool { + return websocket.IsWebSocketUpgrade(req) +} + +// handleWebSocketTLS handles a WebSocket upgrade over an already-established +// TLS connection. +// +// browserReader is the *bufio.Reader that was used to parse the upgrade +// request in the H1 loop. It must be passed here — not discarded — because +// bufio.Reader reads ahead and may have buffered bytes beyond the end of the +// HTTP request. Discarding it and creating a new bufio.Reader over the raw +// conn would silently lose those bytes, corrupting the WebSocket stream. +// +// Design: +// 1. Dial the upstream as a WebSocket client. +// 2. Give gorilla's Upgrader a hijackWriter that wraps (tlsConn, browserReader) +// so gorilla writes the 101 and takes back the conn via Hijack() — +// the read side of Hijack uses the existing browserReader, preserving +// any buffered bytes. Zero replay. Zero data loss. +// 3. Proxy frames bidirectionally with gorilla on both sides. +func (p *Proxy) handleWebSocketTLS( + tlsConn *tls.Conn, + browserReader *bufio.Reader, + req *http.Request, + bareHost string, + hostWithPort string, +) { + defer recoverPanic("websocket " + bareHost) + + upstreamURL := &url.URL{ + Scheme: "wss", + Host: hostWithPort, + Path: req.URL.Path, + RawQuery: req.URL.RawQuery, + } + + // ── 1. Dial upstream ────────────────────────────────────────────────────── + + upstreamHeaders := http.Header{} + for k, vals := range req.Header { + switch k { + case "Upgrade", "Connection", "Sec-Websocket-Key", + "Sec-Websocket-Version", "Sec-Websocket-Extensions", + "Sec-Websocket-Protocol": + // Let gorilla's dialer handle these. + default: + upstreamHeaders[k] = vals + } + } + + dialer := websocket.Dialer{ + TLSClientConfig: &tls.Config{ + ServerName: bareHost, + InsecureSkipVerify: true, //nolint:gosec + }, + HandshakeTimeout: 10 * time.Second, + } + + upstreamConn, upstreamResp, err := dialer.Dial(upstreamURL.String(), upstreamHeaders) + if err != nil { + logger.Error("ws: dial upstream %s: %v", upstreamURL, err) + writeWSError(tlsConn, http.StatusBadGateway) + return + } + defer upstreamConn.Close() + + // ── 2. Upgrade browser conn ─────────────────────────────────────────────── + // + // hijackWriter implements http.ResponseWriter + http.Hijacker over + // (tlsConn, browserReader). When gorilla's Upgrader calls Hijack(), it + // gets back the existing bufio.ReadWriter — the read side IS browserReader, + // so any bytes already buffered by the H1 loop are not lost. + + responseHeader := http.Header{} + if proto := upstreamResp.Header.Get("Sec-Websocket-Protocol"); proto != "" { + responseHeader.Set("Sec-WebSocket-Protocol", proto) + } + + hw := newHijackWriter(tlsConn, browserReader) + upgrader := websocket.Upgrader{ + CheckOrigin: func(r *http.Request) bool { return true }, + ReadBufferSize: 1 << 20, + WriteBufferSize: 1 << 20, + } + browserConn, err := upgrader.Upgrade(hw, req, responseHeader) + if err != nil { + logger.Error("ws: upgrade browser for %s: %v", bareHost, err) + return + } + defer browserConn.Close() + const wsMaxMessageSize = 16 << 20 // 16 MB + browserConn.SetReadLimit(wsMaxMessageSize) + upstreamConn.SetReadLimit(wsMaxMessageSize) + + // ── 3. Emit connection event ────────────────────────────────────────────── + + connID := p.bus.EmitWebSocketConnection(events.WebSocketConnectionEvent{ + Host: hostWithPort, + URL: upstreamURL.String(), + TLS: true, + Timestamp: time.Now(), + }) + + logger.Info("ws: connected %s%s", bareHost, req.URL.Path) + + // ── 4. Bidirectional frame proxy ────────────────────────────────────────── + + done := make(chan struct{}, 2) + + go func() { + defer func() { done <- struct{}{} }() + for { + msgType, payload, err := browserConn.ReadMessage() + if err != nil { + if !websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) { + logger.Debug("ws: browser→upstream read for %s: %v", bareHost, err) + } + upstreamConn.Close() + return + } + result := p.bus.EmitWebSocketFrame(events.WebSocketFrameEvent{ + ConnectionID: connID, + Timestamp: time.Now(), + Direction: events.WebSocketClient, + Opcode: events.WebSocketOpcode(msgType), + Payload: payload, + }) + if err := upstreamConn.WriteMessage(msgType, result.Payload); err != nil { + logger.Debug("ws: browser→upstream write for %s: %v", bareHost, err) + return + } + } + }() + + go func() { + defer func() { done <- struct{}{} }() + for { + msgType, payload, err := upstreamConn.ReadMessage() + if err != nil { + if !websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) { + logger.Debug("ws: upstream→browser read for %s: %v", bareHost, err) + } + browserConn.Close() + return + } + result := p.bus.EmitWebSocketFrame(events.WebSocketFrameEvent{ + ConnectionID: connID, + Timestamp: time.Now(), + Direction: events.WebSocketServer, + Opcode: events.WebSocketOpcode(msgType), + Payload: payload, + }) + if err := browserConn.WriteMessage(msgType, result.Payload); err != nil { + logger.Debug("ws: upstream→browser write for %s: %v", bareHost, err) + return + } + } + }() + + <-done + <-done + logger.Info("ws: closed %s%s", bareHost, req.URL.Path) +} + +// writeWSError writes a minimal HTTP error response directly to a raw conn. +// Used when the upstream dial fails before the 101 is sent. +func writeWSError(conn net.Conn, code int) { + fmt.Fprintf(conn, "HTTP/1.1 %d %s\r\nContent-Length: 0\r\nConnection: close\r\n\r\n", + code, http.StatusText(code)) +} + +// hijackWriter implements http.ResponseWriter and http.Hijacker over a raw +// net.Conn + an existing *bufio.Reader. +// +// The critical invariant: the bufio.Reader passed here MUST be the same one +// used to parse the HTTP upgrade request. gorilla's Upgrader calls Hijack() +// to take ownership of the connection — the read side of the returned +// bufio.ReadWriter is this same reader, so any bytes already buffered in it +// (read-ahead past the end of the HTTP request) are not lost. +// +// The write side is a fresh bufio.Writer over the raw conn — gorilla flushes +// it after writing the 101 response. +type hijackWriter struct { + conn net.Conn + reader *bufio.Reader + header http.Header +} + +func newHijackWriter(conn net.Conn, reader *bufio.Reader) *hijackWriter { + return &hijackWriter{ + conn: conn, + reader: reader, + header: make(http.Header), + } +} + +func (hw *hijackWriter) Header() http.Header { return hw.header } +func (hw *hijackWriter) WriteHeader(_ int) {} +func (hw *hijackWriter) Write(b []byte) (int, error) { return hw.conn.Write(b) } + +// Hijack returns the raw conn and a bufio.ReadWriter whose read side is the +// existing browserReader — preserving any buffered bytes — and whose write +// side is a fresh writer over the conn. +func (hw *hijackWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { + brw := bufio.NewReadWriter(hw.reader, bufio.NewWriter(hw.conn)) + return hw.conn, brw, nil +} 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)) + } +} |