From 0a477fd288d064d265395c88dbd0f261b723804e Mon Sep 17 00:00:00 2001 From: spiperac Date: Mon, 15 Jun 2026 06:42:13 +0200 Subject: Initial release --- internal/proxy/proxy.go | 200 ++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 200 insertions(+) create mode 100644 internal/proxy/proxy.go (limited to 'internal/proxy/proxy.go') 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, + }) +} -- cgit v1.2.3