1 files changed, 92 insertions, 0 deletions
diff --git a/internal/proxy/tee.go b/internal/proxy/tee.go
new file mode 100644
index 0000000..f6a385a
--- /dev/null
+++ b/internal/proxy/tee.go
@@ -0,0 +1,92 @@
+package proxy
+
+import (
+	"context"
+	"net"
+	"sync"
+)
+
+// maxCaptureBytes bounds how much of any single request or response
+// mitmux buffers for history storage, independent of how much data
+// actually flows through the proxy. Proxying itself always streams the
+// full body regardless of this limit - only what gets stored is capped,
+// so a multi-gigabyte download can't be turned into a memory exhaustion
+// vector just because the history view wants to remember it.
+const maxCaptureBytes = 10 << 20 // 10 MiB
+
+// teeConn wraps a net.Conn, recording every byte read off the wire (up
+// to maxCaptureBytes) so it can be attributed to a specific request or
+// response later. Take returns everything recorded since the last call
+// and resets the buffer, so callers must take exactly once per message
+// they want attributed correctly - see forward() for why that's safe
+// here (call sites synchronize on the request/response boundary itself).
+type teeConn struct {
+	net.Conn
+	mu  sync.Mutex
+	buf []byte
+}
+
+func newTeeConn(c net.Conn) *teeConn {
+	return &teeConn{Conn: c}
+}
+
+func (c *teeConn) Read(p []byte) (int, error) {
+	n, err := c.Conn.Read(p)
+	if n > 0 {
+		c.mu.Lock()
+		if room := maxCaptureBytes - len(c.buf); room > 0 {
+			end := n
+			if end > room {
+				end = room
+			}
+			c.buf = append(c.buf, p[:end]...)
+		}
+		c.mu.Unlock()
+	}
+	return n, err
+}
+
+// Take returns the bytes read since the last Take call (or since the
+// connection was created) and resets the buffer.
+func (c *teeConn) Take() []byte {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	out := c.buf
+	c.buf = nil
+	return out
+}
+
+// teeListener wraps a net.Listener so every accepted connection is
+// tee-captured.
+type teeListener struct {
+	net.Listener
+}
+
+func (l *teeListener) Accept() (net.Conn, error) {
+	c, err := l.Listener.Accept()
+	if err != nil {
+		return nil, err
+	}
+	return newTeeConn(c), nil
+}
+
+type contextKey int
+
+const clientTeeKey contextKey = iota
+
+// teeConnFromContext returns the teeConn wrapping the client connection
+// the current request was read from, as attached via http.Server's
+// ConnContext hook. Returns nil for HTTP/2 client connections, which
+// aren't tee-captured (see capture.go).
+func teeConnFromContext(ctx context.Context) *teeConn {
+	tc, _ := ctx.Value(clientTeeKey).(*teeConn)
+	return tc
+}
+
+func withClientTee(ctx context.Context, c net.Conn) context.Context {
+	tc, ok := c.(*teeConn)
+	if !ok {
+		return ctx
+	}
+	return context.WithValue(ctx, clientTeeKey, tc)
+}