// Package proxy is the mitmux proxy engine: a forward HTTP proxy that,
// at this build stage, passes traffic through unmodified. CONNECT
// requests (HTTPS) are tunneled raw rather than intercepted - TLS
// interception is a later build step.
package proxy

import (
	"context"
	"io"
	"log"
	"net"
	"net/http"
	"time"
)

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

	transport *http.Transport
	server    *http.Server
}

// New creates a proxy Server bound to addr (e.g. "127.0.0.1:8080").
func New(addr string) *Server {
	s := &Server{
		Addr: addr,
		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 tunnels a CONNECT request raw, byte for byte, without
// terminating TLS. This is the passthrough behavior for HTTPS traffic
// until TLS interception is implemented.
func (s *Server) handleConnect(w http.ResponseWriter, r *http.Request) {
	dst, err := net.DialTimeout("tcp", r.Host, 10*time.Second)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadGateway)
		return
	}
	defer dst.Close()

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

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

	done := make(chan struct{}, 2)
	go func() {
		io.Copy(dst, src)
		done <- struct{}{}
	}()
	go func() {
		io.Copy(src, dst)
		done <- struct{}{}
	}()
	<-done
}

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