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)
}