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 reconstructs raw response bytes from the parsed
// response for the HTTP/2 upstream case - the exact-capture path
// (HTTP/1.1 upstream) is handled directly in forward() via the
// teeConn's own Take(), which is the only reason this one doesn't also
// need a *teeConn parameter.
func captureResponse(resp *http.Response, bodyCap *cappedTee) (raw []byte, exact bool) {
	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
}