// 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"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/url"
	"sync"
	"time"

	"golang.org/x/net/http2"

	"mitmux/internal/ca"
	"mitmux/internal/rules"
	"mitmux/internal/scope"
	"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

// clientHeaderTimeout bounds how long a client connection can sit
// sending request headers (or a TLS ClientHello, on the CONNECT-tunnel
// leg) before mitmux gives up on it - a slow-loris style connection
// that opens and then trickles bytes (or never sends a ClientHello at
// all) would otherwise hold a connection and its goroutine open
// indefinitely, with nothing else in the codebase bounding it. Narrow
// on purpose: this only covers the pre-body phase, not overall request
// duration - a legitimately slow multi-minute upload/download must
// still work, so this is deliberately not a blanket ReadTimeout/
// WriteTimeout on the whole connection.
const clientHeaderTimeout = 30 * time.Second

// clientIdleTimeout bounds how long a keep-alive client connection can
// sit idle between requests before mitmux closes it - cleans up
// abandoned idle connections without affecting any connection that's
// actively mid-transfer.
const clientIdleTimeout = 120 * 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. It can bind more than one address
// at once (Addrs) - all sharing the same handler, history store, CA and
// rules, so a client on any of them sees identical behavior; this is for
// cases like wanting a separate port per client/network segment, not
// for running logically different proxies in one process.
type Server struct {
	Addrs []string

	// UpstreamProxy, if set (host:port, no scheme), chains every
	// outbound connection through another HTTP CONNECT proxy instead of
	// dialing origins directly - e.g. routing mitmux's own traffic
	// through Burp, a corporate proxy, or a network-access proxy.
	// SOCKS5 upstreams aren't implemented (see PLAN.md).
	UpstreamProxy 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 addrs (e.g. ["127.0.0.1:8080"]),
// signing intercepted TLS connections with root and recording history to
// db. upstreamProxy chains outbound connections through another HTTP
// CONNECT proxy (host:port, no scheme) instead of dialing origins
// directly; empty disables chaining.
func New(addrs []string, root *ca.CA, db *store.Store, upstreamProxy string) *Server {
	s := &Server{Addrs: addrs, ca: root, store: db, UpstreamProxy: upstreamProxy}
	s.server = &http.Server{
		Handler:           http.HandlerFunc(s.handle),
		ConnContext:       withClientTee,
		ReadHeaderTimeout: clientHeaderTimeout,
		IdleTimeout:       clientIdleTimeout,
	}
	return s
}

// ListenAndServe binds every address in Addrs and blocks until one of
// them stops (including on Shutdown, which closes all of them - every
// Serve call below then returns http.ErrServerClosed). Addresses are all
// bound up front before any of them start serving, so a bad address
// (already in use, unparseable, ...) fails startup immediately rather
// than leaving the daemon partially listening.
func (s *Server) ListenAndServe() error {
	if len(s.Addrs) == 0 {
		return errors.New("no listen addresses configured")
	}
	lns := make([]net.Listener, len(s.Addrs))
	for i, addr := range s.Addrs {
		ln, err := net.Listen("tcp", addr)
		if err != nil {
			for _, opened := range lns[:i] {
				opened.Close()
			}
			return fmt.Errorf("listen on %s: %w", addr, err)
		}
		lns[i] = ln
	}

	errCh := make(chan error, len(lns))
	for i, ln := range lns {
		log.Printf("proxy listening on %s", s.Addrs[i])
		go func(ln net.Listener) {
			errCh <- s.server.Serve(&teeListener{Listener: ln})
		}(ln)
	}
	return <-errCh
}

// 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,
	})
	// Bounded the same way the upstream leg already is (see forward's
	// conn.SetDeadline): without this, a client that completes CONNECT
	// and then never sends a ClientHello at all holds the connection and
	// its goroutine open indefinitely. Cleared after a successful
	// handshake - the request/response phase that follows has no
	// business inheriting a short handshake-only deadline.
	client.SetDeadline(time.Now().Add(clientHeaderTimeout))
	if err := clientTLS.Handshake(); err != nil {
		log.Printf("mitm handshake with client for %s: %v", hostname, err)
		client.Close()
		return
	}
	client.SetDeadline(time.Time{})

	dial := func(ctx context.Context) (net.Conn, string, error) {
		return dialUpstreamTLS(ctx, hostPort, hostname, s.UpstreamProxy)
	}
	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,
		ReadHeaderTimeout: clientHeaderTimeout,
		IdleTimeout:       clientIdleTimeout,
	}
	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 (directly, or tunneled
// through upstreamProxy if set - see dialViaProxy), offering both
// HTTP/2 and HTTP/1.1 over ALPN and letting the server pick. Chaining
// through another proxy is transparent to everything from here on: once
// the CONNECT tunnel is up, TLS and the request/response on top of it
// look identical to a direct connection.
func dialUpstreamTLS(ctx context.Context, hostPort, sni, upstreamProxy string) (net.Conn, string, error) {
	raw, err := dialViaProxy(ctx, hostPort, upstreamProxy)
	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. Unlike the TLS/
// CONNECT path, chaining here means dialing the upstream proxy's own
// address directly and writing it an absolute-form request (what a
// proxy expects) rather than tunneling - see forward()'s viaProxyForm.
func dialUpstreamPlain(ctx context.Context, host, upstreamProxy string) (net.Conn, string, error) {
	if _, _, err := net.SplitHostPort(host); err != nil {
		host = net.JoinHostPort(host, "80")
	}
	target := host
	if upstreamProxy != "" {
		target = upstreamProxy
	}
	nd := &net.Dialer{Timeout: 10 * time.Second}
	conn, err := nd.DialContext(ctx, "tcp", target)
	return conn, "http/1.1", err
}

// dialViaProxy returns a raw TCP connection ready to speak TLS to
// hostPort - dialed directly if upstreamProxy is empty, or tunneled
// through upstreamProxy via an HTTP CONNECT request otherwise.
func dialViaProxy(ctx context.Context, hostPort, upstreamProxy string) (net.Conn, error) {
	nd := &net.Dialer{Timeout: 10 * time.Second}
	if upstreamProxy == "" {
		return nd.DialContext(ctx, "tcp", hostPort)
	}

	conn, err := nd.DialContext(ctx, "tcp", upstreamProxy)
	if err != nil {
		return nil, fmt.Errorf("dial upstream proxy %s: %w", upstreamProxy, err)
	}

	connectReq := &http.Request{
		Method: http.MethodConnect,
		URL:    &url.URL{Opaque: hostPort},
		Host:   hostPort,
		Header: make(http.Header),
	}
	if err := connectReq.Write(conn); err != nil {
		conn.Close()
		return nil, fmt.Errorf("write CONNECT to upstream proxy %s: %w", upstreamProxy, err)
	}

	br := bufio.NewReader(conn)
	resp, err := http.ReadResponse(br, connectReq)
	if err != nil {
		conn.Close()
		return nil, fmt.Errorf("read CONNECT response from upstream proxy %s: %w", upstreamProxy, err)
	}
	if resp.StatusCode != http.StatusOK {
		conn.Close()
		return nil, fmt.Errorf("upstream proxy %s refused CONNECT to %s: %s", upstreamProxy, hostPort, resp.Status)
	}
	if br.Buffered() > 0 {
		// The upstream proxy shouldn't send anything past the CONNECT
		// response before the tunnel starts, but if it did, those bytes
		// are sitting in br's buffer, not on conn - replay them first
		// rather than silently dropping the start of the TLS handshake.
		return &prefixedConn{Conn: conn, r: br}, nil
	}
	return conn, nil
}

// prefixedConn serves buffered bytes from r before falling through to
// reading directly off the underlying connection.
type prefixedConn struct {
	net.Conn
	r *bufio.Reader
}

func (c *prefixedConn) Read(p []byte) (int, error) {
	if c.r.Buffered() > 0 {
		return c.r.Read(p)
	}
	return c.Conn.Read(p)
}

// 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. proxyForm selects an absolute-form
// request line ("GET http://host/path HTTP/1.1") instead of origin-form
// - needed when conn is a connection to another proxy, which expects
// that form, rather than to the origin server itself.
func roundTripH1(conn net.Conn, outReq *http.Request, proxyForm bool) (*http.Response, *teeConn, error) {
	tee := newTeeConn(conn)
	var writeErr error
	if proxyForm {
		writeErr = outReq.WriteProxy(tee)
	} else {
		writeErr = outReq.Write(tee)
	}
	if writeErr != nil {
		return nil, nil, writeErr
	}
	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, reqTrunc := captureRequest(r, clientTee, reqBodyCap)
		s.record(started, time.Since(started), scheme, hostname, r, reqRaw, reqExact, reqTrunc, nil, false, 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 {
		// Only the plain-HTTP path needs absolute-form: a CONNECT tunnel
		// (chained through an upstream proxy or not) is transparent from
		// here on, so it always uses origin-form like a direct connection.
		proxyForm := scheme == "http" && s.UpstreamProxy != ""
		resp, upstreamTee, err = roundTripH1(conn, outReq, proxyForm)
	}
	duration := time.Since(started)

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

	if err != nil {
		s.record(started, duration, scheme, hostname, r, reqRaw, reqExact, reqTrunc, nil, false, 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, respTrunc bool
	if upstreamTee != nil {
		respRaw, respTrunc = upstreamTee.Take()
		respExact = !respTrunc
	} else {
		respRaw, respExact = captureResponse(resp, respBodyCap)
	}

	s.record(started, duration, scheme, hostname, r, reqRaw, reqExact, reqTrunc, respRaw, respExact, respTrunc, 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, reqTruncated bool, respRaw []byte, respExact, respTruncated bool, status int, errMsg string) {
	if s.store == nil {
		return
	}
	// Scope only filters what gets recorded here - the request has
	// already been forwarded and its response already written to the
	// client by the time record() runs (see forward()), so an
	// out-of-scope host still proxies completely normally, it just
	// doesn't clutter history. Repeat/Intrude (recordRaw, a different
	// function) deliberately don't go through this check: a user
	// explicitly resending or fuzzing a specific request wants to see
	// the result regardless of scope, which exists to cut passive-
	// capture noise, not to second-guess a deliberate action.
	if scopeRules, err := s.store.ListScopeRules(); err == nil && !scope.InScope(scopeRules, host) {
		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,
		RequestTruncated:  reqTruncated,
		ResponseTruncated: respTruncated,
		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.UpstreamProxy)
	}
	s.forward(dial, r.URL.Scheme, r.URL.Host, w, r)
}

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