1 files changed, 83 insertions, 0 deletions
diff --git a/internal/proxy/capture.go b/internal/proxy/capture.go
new file mode 100644
index 0000000..ccc95e9
--- /dev/null
+++ b/internal/proxy/capture.go
@@ -0,0 +1,83 @@
+package proxy
+
+import (
+	"bytes"
+	"io"
+	"net/http"
+)
+
+// cappedTee wraps an io.Reader, copying up to maxCaptureBytes of what
+// passes through into an internal buffer while still passing everything
+// through unmodified and unbounded to the real reader. Used to capture a
+// bounded sample of a body for reconstruction when exact wire capture
+// isn't available (the HTTP/2 leg - see below).
+type cappedTee struct {
+	r   io.Reader
+	buf bytes.Buffer
+}
+
+func newCappedTee(r io.Reader) *cappedTee {
+	return &cappedTee{r: r}
+}
+
+func (c *cappedTee) Read(p []byte) (int, error) {
+	n, err := c.r.Read(p)
+	if n > 0 {
+		if room := maxCaptureBytes - c.buf.Len(); room > 0 {
+			end := n
+			if end > room {
+				end = room
+			}
+			c.buf.Write(p[:end])
+		}
+	}
+	return n, err
+}
+
+// captureRequest returns the raw bytes of r for storage. When tee is
+// non-nil (an HTTP/1.1 client connection), the bytes are exactly what
+// was read off the wire. Otherwise (HTTP/2, which has no single "raw
+// bytes" representation - it's multiplexed, HPACK-compressed framing)
+// it's a reconstruction from the parsed request, exact=false.
+func captureRequest(r *http.Request, tee *teeConn, bodyCap *cappedTee) (raw []byte, exact bool) {
+	if tee != nil {
+		return tee.Take(), true
+	}
+
+	dump := r.Clone(r.Context())
+	if bodyCap != nil {
+		dump.Body = io.NopCloser(bytes.NewReader(bodyCap.buf.Bytes()))
+		dump.ContentLength = int64(bodyCap.buf.Len())
+	} else {
+		dump.Body = http.NoBody
+		dump.ContentLength = 0
+	}
+	var buf bytes.Buffer
+	if err := dump.Write(&buf); err != nil {
+		return nil, false
+	}
+	return buf.Bytes(), false
+}
+
+// captureResponse mirrors captureRequest for the upstream leg: exact
+// wire bytes when tee is non-nil (upstream negotiated HTTP/1.1),
+// otherwise a reconstruction.
+func captureResponse(resp *http.Response, tee *teeConn, bodyCap *cappedTee) (raw []byte, exact bool) {
+	if tee != nil {
+		return tee.Take(), true
+	}
+
+	dump := *resp
+	if bodyCap != nil {
+		dump.Body = io.NopCloser(bytes.NewReader(bodyCap.buf.Bytes()))
+		dump.ContentLength = int64(bodyCap.buf.Len())
+	} else {
+		dump.Body = http.NoBody
+		dump.ContentLength = 0
+	}
+	var buf bytes.Buffer
+	if err := dump.Write(&buf); err != nil {
+		return nil, false
+	}
+	return buf.Bytes(), false
+}