// Package proxy is the mitmux proxy engine: a forward HTTP proxy.
// Plain HTTP requests pass through unmodified. CONNECT requests (HTTPS)
// are intercepted: mitmux terminates TLS with the client using a leaf
// certificate signed by its own CA, and separately terminates TLS with
// the real server, forwarding requests between the two. ALPN is
// negotiated independently on each side (see handleConnect) so HTTP/2
// stays HTTP/2 end to end without one side being forced to match the
// other. Every request/response pair is captured to the history store -
// exactly, byte for byte, on HTTP/1.1 legs; reconstructed on HTTP/2 legs,
// which have no meaningful "raw bytes" of their own (see capture.go).
//
// Upstream requests are round-tripped manually (write the request,
// read the response off the same connection) rather than through
// http.Transport: Transport's automatic HTTP/2 dispatch keys off a
// literal *tls.Conn type assertion on the connection it dials, which a
// capturing wrapper around that connection defeats - the request would
// silently be parsed as HTTP/1.1 over what is actually HTTP/2 framing.
// Handling both protocols explicitly here, per request, avoids that and
// also removes any ambiguity about which connection served which
// request, since each request gets its own connection either way.
package proxy

import (
	"bufio"
	"context"
	"crypto/tls"
	"errors"
	"io"
	"log"
	"net"
	"net/http"
	"sync"
	"time"

	"golang.org/x/net/http2"

	"mitmux/internal/ca"
	"mitmux/internal/rules"
	"mitmux/internal/store"
)

// upstreamTimeout bounds the write-request/read-response phase of an
// upstream exchange, once dialing has already succeeded.
const upstreamTimeout = 60 * time.Second

// hopByHopHeaders are stripped before forwarding a request or response,
// per RFC 7230 6.1 - they are meaningful only between a client and its
// immediate next hop, not end-to-end.
var hopByHopHeaders = []string{
	"Connection",
	"Proxy-Connection",
	"Keep-Alive",
	"Proxy-Authenticate",
	"Proxy-Authorization",
	"TE",
	"Trailers",
	"Transfer-Encoding",
	"Upgrade",
}

// Server is a forward proxy listener.
type Server struct {
	Addr string

	// OnEntry, if set, is called after each request/response pair is
	// stored, so a daemon can broadcast it to live TUI subscribers.
	OnEntry func(store.Summary)

	ca     *ca.CA
	store  *store.Store
	server *http.Server
}

// New creates a proxy Server bound to addr (e.g. "127.0.0.1:8080"),
// signing intercepted TLS connections with root and recording history to
// db.
func New(addr string, root *ca.CA, db *store.Store) *Server {
	s := &Server{Addr: addr, ca: root, store: db}
	s.server = &http.Server{
		Addr:        addr,
		Handler:     http.HandlerFunc(s.handle),
		ConnContext: withClientTee,
	}
	return s
}

// ListenAndServe starts the proxy and blocks until it stops.
func (s *Server) ListenAndServe() error {
	ln, err := net.Listen("tcp", s.Addr)
	if err != nil {
		return err
	}
	log.Printf("proxy listening on %s", s.Addr)
	return s.server.Serve(&teeListener{Listener: ln})
}

// Shutdown gracefully stops the proxy.
func (s *Server) Shutdown(ctx context.Context) error {
	return s.server.Shutdown(ctx)
}

func (s *Server) handle(w http.ResponseWriter, r *http.Request) {
	if r.Method == http.MethodConnect {
		s.handleConnect(w, r)
		return
	}
	s.handleHTTP(w, r)
}

// dialer resolves a fresh upstream connection for one request, along
// with the ALPN protocol negotiated for it ("http/1.1", "h2", or "" if
// not applicable/negotiated).
type dialer func(ctx context.Context) (conn net.Conn, negotiated string, err error)

// handleConnect intercepts a CONNECT request: it terminates TLS with the
// client using a leaf certificate signed by mitmux's CA, then forwards
// each request upstream over its own independently negotiated TLS
// connection. Client-side and upstream-side ALPN are negotiated
// separately (each offering both HTTP/2 and HTTP/1.1) rather than one
// being forced to match the other, so e.g. an HTTP/1.1-only client
// reaching an HTTP/2-preferring server doesn't fail to connect.
func (s *Server) handleConnect(w http.ResponseWriter, r *http.Request) {
	hostPort := r.Host
	hostname, _, err := net.SplitHostPort(hostPort)
	if err != nil {
		hostname = hostPort
		hostPort = net.JoinHostPort(hostPort, "443")
	}

	hijacker, ok := w.(http.Hijacker)
	if !ok {
		http.Error(w, "hijacking not supported", http.StatusInternalServerError)
		return
	}
	client, _, err := hijacker.Hijack()
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	if _, err := client.Write([]byte("HTTP/1.1 200 Connection Established\r\n\r\n")); err != nil {
		client.Close()
		return
	}

	clientTLS := tls.Server(client, &tls.Config{
		GetCertificate: func(hello *tls.ClientHelloInfo) (*tls.Certificate, error) {
			name := hello.ServerName
			if name == "" {
				name = hostname
			}
			return s.ca.LeafFor(name)
		},
		NextProtos: []string{http2.NextProtoTLS, "http/1.1"},
		MinVersion: tls.VersionTLS12,
	})
	if err := clientTLS.Handshake(); err != nil {
		log.Printf("mitm handshake with client for %s: %v", hostname, err)
		client.Close()
		return
	}

	dial := func(ctx context.Context) (net.Conn, string, error) {
		return dialUpstreamTLS(ctx, hostPort, hostname)
	}
	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		s.forward(dial, "https", hostname, w, r)
	})

	if clientTLS.ConnectionState().NegotiatedProtocol == http2.NextProtoTLS {
		(&http2.Server{}).ServeConn(clientTLS, &http2.ServeConnOpts{Handler: handler})
		return
	}

	h1 := &http.Server{Handler: handler, ConnContext: withClientTee}
	err = h1.Serve(newSingleConnListener(clientTLS))
	if err != nil && !errors.Is(err, io.EOF) {
		log.Printf("h1 serve for %s: %v", hostname, err)
	}
}

// dialUpstreamTLS connects to the real server, offering both HTTP/2 and
// HTTP/1.1 over ALPN and letting the server pick.
func dialUpstreamTLS(ctx context.Context, hostPort, sni string) (net.Conn, string, error) {
	nd := &net.Dialer{Timeout: 10 * time.Second}
	raw, err := nd.DialContext(ctx, "tcp", hostPort)
	if err != nil {
		return nil, "", err
	}
	conn := tls.Client(raw, &tls.Config{
		ServerName: sni,
		NextProtos: []string{http2.NextProtoTLS, "http/1.1"},
	})
	if err := conn.HandshakeContext(ctx); err != nil {
		raw.Close()
		return nil, "", err
	}
	return conn, conn.ConnectionState().NegotiatedProtocol, nil
}

// dialUpstreamPlain connects to a plain (non-TLS) upstream for the
// non-CONNECT proxy path, which is always HTTP/1.1.
func dialUpstreamPlain(ctx context.Context, host string) (net.Conn, string, error) {
	if _, _, err := net.SplitHostPort(host); err != nil {
		host = net.JoinHostPort(host, "80")
	}
	nd := &net.Dialer{Timeout: 10 * time.Second}
	conn, err := nd.DialContext(ctx, "tcp", host)
	return conn, "http/1.1", err
}

// roundTripH1 writes outReq directly to conn and reads the response back
// off the same connection, wrapping conn in a teeConn so the exact wire
// bytes of both can be captured.
func roundTripH1(conn net.Conn, outReq *http.Request) (*http.Response, *teeConn, error) {
	tee := newTeeConn(conn)
	if err := outReq.Write(tee); err != nil {
		return nil, nil, err
	}
	resp, err := http.ReadResponse(bufio.NewReader(tee), outReq)
	if err != nil {
		return nil, nil, err
	}
	return resp, tee, nil
}

// roundTripH2 sends outReq over a new single-connection HTTP/2 client.
func roundTripH2(conn net.Conn, outReq *http.Request) (*http.Response, error) {
	cc, err := (&http2.Transport{}).NewClientConn(conn)
	if err != nil {
		return nil, err
	}
	return cc.RoundTrip(outReq)
}

// forward dials upstream, sends r, copies the response back to w, and
// records the exchange to history. r's URL is rewritten from
// origin-form (as read off the terminated connection) to absolute-form
// for the round trip.
func (s *Server) forward(dial dialer, scheme, hostname string, w http.ResponseWriter, r *http.Request) {
	clientTee := teeConnFromContext(r.Context())

	outReq := r.Clone(r.Context())
	outReq.URL.Scheme = scheme
	outReq.URL.Host = hostname
	outReq.RequestURI = ""
	stripHopByHop(outReq.Header)
	// Header rules are applied to outReq only, after cloning and header
	// stripping - history's request_raw keeps showing what the client
	// actually sent (clientTee/reqBodyCap already capture from r, not
	// outReq), while what actually reaches the upstream server reflects
	// the rules. That split is deliberate: match-and-replace is a wire
	// transform, not a rewrite of the audit trail.
	if reqRules, err := s.enabledRules("request"); err != nil {
		log.Printf("load request rules: %v", err)
	} else if len(reqRules) > 0 {
		outReq.Header = rules.ApplyHeaders(outReq.Header, reqRules)
	}

	// Only needed when the client leg isn't tee-captured (HTTP/2): tee
	// the body as it streams through so the reconstructed capture isn't
	// missing it.
	var reqBodyCap *cappedTee
	if clientTee == nil && outReq.Body != nil {
		reqBodyCap = newCappedTee(outReq.Body)
		outReq.Body = io.NopCloser(reqBodyCap)
	}

	started := time.Now()
	conn, negotiated, dialErr := dial(r.Context())
	if dialErr != nil {
		reqRaw, reqExact := captureRequest(r, clientTee, reqBodyCap)
		s.record(started, time.Since(started), scheme, hostname, r, reqRaw, reqExact, nil, false, 0, dialErr.Error())
		http.Error(w, dialErr.Error(), http.StatusBadGateway)
		return
	}
	defer conn.Close()
	// The dial itself is bounded (net.Dialer.Timeout / HandshakeContext);
	// without this, a server that accepts the connection and then never
	// writes or never finishes writing would hang the request forever -
	// there's no other timeout covering the write-request/read-response
	// phase. Bounds the whole exchange, so a legitimately slow multi-
	// minute transfer would also get cut off; a fixed default is enough
	// for now, not worth a config surface yet.
	conn.SetDeadline(time.Now().Add(upstreamTimeout))

	var resp *http.Response
	var upstreamTee *teeConn
	var err error
	if negotiated == http2.NextProtoTLS {
		resp, err = roundTripH2(conn, outReq)
	} else {
		resp, upstreamTee, err = roundTripH1(conn, outReq)
	}
	duration := time.Since(started)

	reqRaw, reqExact := captureRequest(r, clientTee, reqBodyCap)

	if err != nil {
		s.record(started, duration, scheme, hostname, r, reqRaw, reqExact, nil, false, 0, err.Error())
		http.Error(w, err.Error(), http.StatusBadGateway)
		return
	}
	defer resp.Body.Close()

	var respBodyCap *cappedTee
	if upstreamTee == nil {
		respBodyCap = newCappedTee(resp.Body)
		resp.Body = io.NopCloser(respBodyCap)
	}

	// Same split as the request side: response_raw keeps reflecting what
	// the origin server actually sent (captured below, from upstreamTee
	// or respBodyCap, both already wired to resp.Body independent of
	// resp.Header), while the client actually receives the rule-modified
	// headers.
	if respRules, err := s.enabledRules("response"); err != nil {
		log.Printf("load response rules: %v", err)
	} else if len(respRules) > 0 {
		resp.Header = rules.ApplyHeaders(resp.Header, respRules)
	}

	stripHopByHop(resp.Header)
	for k, vv := range resp.Header {
		for _, v := range vv {
			w.Header().Add(k, v)
		}
	}
	w.WriteHeader(resp.StatusCode)
	io.Copy(w, resp.Body)

	var respRaw []byte
	var respExact bool
	if upstreamTee != nil {
		respRaw, respExact = upstreamTee.Take(), true
	} else {
		respRaw, respExact = captureResponse(resp, respBodyCap)
	}

	s.record(started, duration, scheme, hostname, r, reqRaw, reqExact, respRaw, respExact, resp.StatusCode, "")
}

// enabledRules fetches the current enabled match-and-replace rules for
// scope ("request" or "response") fresh from the store on every call -
// simple and always current, and cheap enough (a local, in-process
// SQLite query) not to bother caching for how this is actually used.
func (s *Server) enabledRules(scope string) ([]rules.Rule, error) {
	if s.store == nil {
		return nil, nil
	}
	return s.store.EnabledRules(scope)
}

// record stores one history entry and notifies OnEntry.
func (s *Server) record(started time.Time, duration time.Duration, scheme, host string, r *http.Request,
	reqRaw []byte, reqExact bool, respRaw []byte, respExact bool, status int, errMsg string) {
	if s.store == nil {
		return
	}

	e := &store.Entry{
		StartedAt:     started,
		Duration:      duration,
		Method:        r.Method,
		Scheme:        scheme,
		Host:          host,
		Path:          r.URL.Path,
		StatusCode:    status,
		RequestRaw:    reqRaw,
		ResponseRaw:   respRaw,
		RequestExact:  reqExact,
		ResponseExact: respExact,
		Error:         errMsg,
	}
	id, err := s.store.Insert(e)
	if err != nil {
		log.Printf("store history entry: %v", err)
		return
	}

	if s.OnEntry != nil {
		s.OnEntry(store.Summary{
			ID:         id,
			StartedAt:  e.StartedAt,
			Duration:   e.Duration,
			Method:     e.Method,
			Scheme:     e.Scheme,
			Host:       e.Host,
			Path:       e.Path,
			StatusCode: e.StatusCode,
			ReqSize:    len(reqRaw),
			RespSize:   len(respRaw),
			Error:      errMsg,
			Source:     "proxy",
		})
	}
}

// singleConnListener adapts one already-accepted net.Conn into a
// net.Listener so http.Serve can drive it, returning io.EOF from the
// second Accept once the connection closes.
type singleConnListener struct {
	ch   chan net.Conn
	addr net.Addr
}

// newSingleConnListener wraps c for one Accept, teeConn on the outside
// so a *teeConn is what ConnContext sees (see withClientTee) - wrapping
// it the other way around lets closeSignalConn's concrete type mask the
// teeConn from that type assertion, silently disabling capture.
func newSingleConnListener(c net.Conn) *singleConnListener {
	ch := make(chan net.Conn, 1)
	signaled := &closeSignalConn{Conn: c, onClose: sync.OnceFunc(func() { close(ch) })}
	ch <- newTeeConn(signaled)
	return &singleConnListener{ch: ch, addr: c.LocalAddr()}
}

func (l *singleConnListener) Accept() (net.Conn, error) {
	c, ok := <-l.ch
	if !ok {
		return nil, io.EOF
	}
	return c, nil
}

func (l *singleConnListener) Close() error   { return nil }
func (l *singleConnListener) Addr() net.Addr { return l.addr }

type closeSignalConn struct {
	net.Conn
	onClose func()
}

func (c *closeSignalConn) Close() error {
	err := c.Conn.Close()
	c.onClose()
	return err
}

// handleHTTP forwards a plain (non-CONNECT) proxy request, copies the
// response back, and records it to history.
func (s *Server) handleHTTP(w http.ResponseWriter, r *http.Request) {
	if !r.URL.IsAbs() {
		http.Error(w, "mitmux: request must use absolute-form URI (configure as a proxy, not a target)", http.StatusBadRequest)
		return
	}
	host := r.URL.Host
	dial := func(ctx context.Context) (net.Conn, string, error) {
		return dialUpstreamPlain(ctx, host)
	}
	s.forward(dial, r.URL.Scheme, r.URL.Host, w, r)
}

func stripHopByHop(h http.Header) {
	for _, k := range hopByHopHeaders {
		h.Del(k)
	}
}