1 files changed, 200 insertions, 0 deletions
diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go
new file mode 100644
index 0000000..3e9845a
--- /dev/null
+++ b/internal/proxy/proxy.go
@@ -0,0 +1,200 @@
+package proxy
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"io"
+	"net"
+	"net/http"
+	"sync"
+	"time"
+
+	"github.com/shiv/internal/cert"
+	"github.com/shiv/internal/events"
+	internalhttp "github.com/shiv/internal/http"
+	"github.com/shiv/internal/logger"
+)
+
+type Proxy struct {
+	addr     string
+	certAuth *cert.CA
+	bus      *events.Bus
+	client   *http.Client
+
+	mu    sync.Mutex
+	srv   *http.Server
+	conns sync.Map // active hijacked net.Conn
+}
+
+const maxBodySize = 10 << 20
+
+func New(addr string, bus *events.Bus) (*Proxy, error) {
+	ca, err := cert.Load()
+	if err != nil {
+		return nil, fmt.Errorf("proxy: load CA: %w", err)
+	}
+	p := &Proxy{
+		addr:     addr,
+		certAuth: ca,
+		bus:      bus,
+		client: &http.Client{
+			Transport: &http.Transport{
+				ResponseHeaderTimeout: 30 * time.Second,
+			},
+			CheckRedirect: func(req *http.Request, via []*http.Request) error {
+				return http.ErrUseLastResponse
+			},
+		},
+	}
+	bus.Register(p)
+	return p, nil
+}
+
+func (p *Proxy) CA() *cert.CA {
+	return p.certAuth
+}
+
+func (p *Proxy) trackConn(c net.Conn) {
+	p.conns.Store(c, struct{}{})
+}
+
+func (p *Proxy) untrackConn(c net.Conn) {
+	p.conns.Delete(c)
+}
+
+func (p *Proxy) closeAllConns() {
+	p.conns.Range(func(k, _ any) bool {
+		k.(net.Conn).SetDeadline(time.Now())
+		return true
+	})
+}
+
+func (p *Proxy) Start(newAddr string) error {
+	ln, err := net.Listen("tcp", newAddr)
+	if err != nil {
+		return fmt.Errorf("proxy: listen on %s: %w", newAddr, err)
+	}
+
+	newSrv := &http.Server{
+		Handler:      p,
+		ReadTimeout:  60 * time.Second,
+		WriteTimeout: 60 * time.Second,
+		IdleTimeout:  120 * time.Second,
+	}
+
+	p.mu.Lock()
+	oldSrv := p.srv
+	p.addr = newAddr
+	p.srv = newSrv
+	p.mu.Unlock()
+
+	if oldSrv != nil {
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+		_ = oldSrv.Shutdown(ctx)
+	}
+	p.closeAllConns()
+
+	go func() {
+		if err := newSrv.Serve(ln); err != nil && err != http.ErrServerClosed {
+			logger.Error("proxy: Serve: %v", err)
+		}
+	}()
+
+	return nil
+}
+
+func (p *Proxy) Stop() {
+	p.mu.Lock()
+	srv := p.srv
+	p.mu.Unlock()
+
+	if srv != nil {
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+		_ = srv.Shutdown(ctx)
+
+		p.mu.Lock()
+		p.srv = nil
+		p.mu.Unlock()
+	}
+
+	p.closeAllConns()
+}
+
+// ObserveProxyRestart implements events.ProxyCommandObserver.
+func (p *Proxy) ObserveProxyRestart(e events.ProxyRestartEvent) {
+	_ = p.Start(e.Addr)
+}
+
+// ObserveProxyStop implements events.ProxyCommandObserver.
+func (p *Proxy) ObserveProxyStop(_ events.ProxyStopEvent) {
+	p.Stop()
+}
+
+func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	if r.Method == http.MethodConnect {
+		p.handleConnect(w, r)
+		return
+	}
+
+	reqBody, err := io.ReadAll(io.LimitReader(r.Body, maxBodySize))
+	if err != nil {
+		http.Error(w, "failed to read request body", http.StatusBadRequest)
+		return
+	}
+	r.Body = io.NopCloser(bytes.NewReader(reqBody))
+
+	result := p.bus.EmitRequest(events.RequestEvent{Request: r, Body: reqBody})
+	if result.Drop {
+		http.Error(w, "request dropped", http.StatusForbidden)
+		return
+	}
+	interceptedReq, reqBody := result.Request, result.Body
+	interceptedReq.Body = io.NopCloser(bytes.NewReader(reqBody))
+
+	start := time.Now()
+
+	resp, err := forward(interceptedReq, p.client)
+	if err != nil {
+		http.Error(w, "bad gateway", http.StatusBadGateway)
+		return
+	}
+	defer resp.Body.Close()
+
+	var respBuf bytes.Buffer
+	reader := io.TeeReader(io.LimitReader(resp.Body, maxBodySize), &respBuf)
+
+	for k, vals := range resp.Header {
+		for _, v := range vals {
+			w.Header().Add(k, v)
+		}
+	}
+	w.WriteHeader(resp.StatusCode)
+
+	_, _ = io.Copy(w, reader)
+
+	respBody := respBuf.Bytes()
+	elapsed := time.Since(start).Milliseconds()
+
+	logBody := internalhttp.Decompress(resp.Header, respBody)
+	if internalhttp.IsBinary(resp.Header) {
+		logBody = nil
+	}
+
+	p.bus.EmitResponse(events.ResponseEvent{
+		Timestamp:   start,
+		Host:        interceptedReq.Host,
+		Proto:       "HTTP/1.1",
+		Method:      interceptedReq.Method,
+		URL:         interceptedReq.URL.String(),
+		ReqHeaders:  interceptedReq.Header,
+		ReqBody:     reqBody,
+		StatusCode:  resp.StatusCode,
+		RespHeaders: resp.Header,
+		RespBody:    logBody,
+		DurationMs:  elapsed,
+		TLS:         false,
+	})
+}