// 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 with the real server first and mirrored to the client so
// HTTP/2 connections stay HTTP/2 end to end rather than being downgraded.
package proxy

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

	"golang.org/x/net/http2"

	"mitmux/internal/ca"
)

// 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

	ca        *ca.CA
	transport *http.Transport
	server    *http.Server
}

// New creates a proxy Server bound to addr (e.g. "127.0.0.1:8080"),
// signing intercepted TLS connections with root.
func New(addr string, root *ca.CA) *Server {
	s := &Server{
		Addr: addr,
		ca:   root,
		transport: &http.Transport{
			Proxy: nil,
			DialContext: (&net.Dialer{
				Timeout: 10 * time.Second,
			}).DialContext,
			ForceAttemptHTTP2:     false,
			MaxIdleConns:          100,
			IdleConnTimeout:       90 * time.Second,
			TLSHandshakeTimeout:   10 * time.Second,
			ExpectContinueTimeout: 1 * time.Second,
		},
	}
	s.server = &http.Server{
		Addr:    addr,
		Handler: http.HandlerFunc(s.handle),
	}
	return s
}

// ListenAndServe starts the proxy and blocks until it stops.
func (s *Server) ListenAndServe() error {
	log.Printf("proxy listening on %s", s.Addr)
	return s.server.ListenAndServe()
}

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

// 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-only-preferring server doesn't fail to connect -
// http.Transport (via http2.ConfigureTransport) bridges the two sides
// independently per request.
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
	}

	tr := &http.Transport{
		DialTLSContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
			return dialUpstreamTLS(ctx, hostPort, hostname)
		},
	}
	if err := http2.ConfigureTransport(tr); err != nil {
		log.Printf("configure h2 transport for %s: %v", hostname, err)
	}
	defer tr.CloseIdleConnections()

	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		s.forward(tr, "https", hostname, w, r)
	})

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

	err = http.Serve(newSingleConnListener(clientTLS), handler)
	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) (*tls.Conn, error) {
	dialer := &net.Dialer{Timeout: 10 * time.Second}
	raw, err := dialer.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, nil
}

// forward sends r upstream via rt and copies the response back to w,
// rewriting r's URL from origin-form (as read off the terminated TLS
// connection) to absolute-form for the round trip.
func (s *Server) forward(rt http.RoundTripper, scheme, hostname string, w http.ResponseWriter, r *http.Request) {
	outReq := r.Clone(r.Context())
	outReq.URL.Scheme = scheme
	outReq.URL.Host = hostname
	outReq.RequestURI = ""
	stripHopByHop(outReq.Header)

	resp, err := rt.RoundTrip(outReq)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadGateway)
		return
	}
	defer resp.Body.Close()

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

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

func newSingleConnListener(c net.Conn) *singleConnListener {
	ch := make(chan net.Conn, 1)
	ch <- &closeSignalConn{Conn: c, onClose: sync.OnceFunc(func() { close(ch) })}
	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 and copies the
// response back unmodified.
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
	}

	outReq := r.Clone(r.Context())
	outReq.RequestURI = ""
	stripHopByHop(outReq.Header)

	resp, err := s.transport.RoundTrip(outReq)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadGateway)
		return
	}
	defer resp.Body.Close()

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

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