vibeproxy/internal/proxy/reverseproxy.go

package proxy

import (
	"context"
	"crypto/tls"
	"fmt"
	"log/slog"
	"net"
	"net/http"
	"net/http/httputil"
	"strings"
	"time"

	"github.com/owen/vibeproxy/internal/loadbalancer"
	"github.com/owen/vibeproxy/internal/metrics"
	"github.com/owen/vibeproxy/internal/models"
	"github.com/owen/vibeproxy/internal/runtime"
)

// sharedTransport is reused across every proxied request so backend
// connections are pooled (keep-alive) instead of dialed fresh each time.
// It supports HTTP/2 to upstreams that offer it via ALPN.
var sharedTransport = &http.Transport{
	Proxy: nil,
	DialContext: (&net.Dialer{
		Timeout:   10 * time.Second,
		KeepAlive: 30 * time.Second,
	}).DialContext,
	ForceAttemptHTTP2:     true,
	MaxIdleConns:          200,
	MaxIdleConnsPerHost:   50,
	IdleConnTimeout:       90 * time.Second,
	TLSHandshakeTimeout:   10 * time.Second,
	ExpectContinueTimeout: 1 * time.Second,
	TLSClientConfig:       &tls.Config{MinVersion: tls.VersionTLS12},
}

// buildReverseProxy constructs a one-shot ReverseProxy for a single
// request/backend pair. Constructing it per request is cheap (no
// goroutines or connections are created here); the expensive parts —
// backend TCP/TLS connections — are pooled by sharedTransport across
// calls.
func buildReverseProxy(rt *runtime.ProxyHostRuntime, backend *loadbalancer.Backend, snap *runtime.Snapshot, log *slog.Logger) *httputil.ReverseProxy {
	targetScheme := rt.Model.ForwardScheme
	if targetScheme == "" {
		targetScheme = "http"
	}
	targetAddr := net.JoinHostPort(backend.Model.Host, fmt.Sprintf("%d", backend.Model.Port))

	director := func(req *http.Request) {
		clientIP := ClientIP(req, snap)
		priorFor := req.Header.Get("X-Forwarded-For")

		req.URL.Scheme = targetScheme
		req.URL.Host = targetAddr
		if !rt.Model.PreserveHostHeader {
			req.Host = targetAddr
		}
		// else: leave req.Host as the original inbound Host header

		if priorFor == "" {
			req.Header.Set("X-Forwarded-For", clientIP)
		} else {
			req.Header.Set("X-Forwarded-For", priorFor+", "+clientIP)
		}
		proto := "http"
		if req.TLS != nil {
			proto = "https"
		}
		req.Header.Set("X-Forwarded-Proto", proto)
		if fh := req.Header.Get("X-Forwarded-Host"); fh == "" {
			req.Header.Set("X-Forwarded-Host", req.Host)
		}
		req.Header.Set("X-Real-IP", clientIP)

		applyCustomHeaders(req.Header, rt.Model.CustomHeaders, models.HeaderDirectionRequest)

		// WebSocket upgrade headers must pass through untouched; Go's
		// httputil.ReverseProxy handles the hijack/tunnel automatically
		// once Connection: Upgrade is present, which we never strip.
		if !rt.Model.WebsocketSupport {
			req.Header.Del("Upgrade")
			req.Header.Set("Connection", "close")
		}
	}

	modifyResponse := func(resp *http.Response) error {
		applyCustomHeaders(resp.Header, rt.Model.CustomHeaders, models.HeaderDirectionResponse)
		if rt.Model.HSTSEnabled && resp.Request.TLS != nil {
			resp.Header.Set("Strict-Transport-Security", "max-age=63072000; includeSubDomains")
		}
		return nil
	}

	errorHandler := func(w http.ResponseWriter, r *http.Request, err error) {
		log.Warn("backend request failed", "proxy_host", rt.Model.Name, "backend", targetAddr, "err", err)
		metrics.HTTPRequestsTotal.WithLabelValues(rt.Model.Name, r.Method, "502").Inc()
		http.Error(w, "Bad Gateway", http.StatusBadGateway)
	}

	return &httputil.ReverseProxy{
		Director:       director,
		ModifyResponse: modifyResponse,
		ErrorHandler:   errorHandler,
		Transport:      sharedTransport,
		FlushInterval:  200 * time.Millisecond, // needed for SSE/streaming responses
	}
}

func applyCustomHeaders(h http.Header, headers []models.CustomHeader, direction models.CustomHeaderDirection) {
	for _, c := range headers {
		if c.Direction != direction {
			continue
		}
		switch c.Action {
		case models.HeaderActionSet:
			h.Set(c.Name, c.Value)
		case models.HeaderActionRemove:
			h.Del(c.Name)
		}
	}
}

// withConnectTimeout returns a context that bounds only the dial phase of
// an outbound request via a deadline; read/write timeouts are handled by
// the surrounding http.Server / http.ResponseController.
func withConnectTimeout(ctx context.Context, d time.Duration) (context.Context, context.CancelFunc) {
	if d <= 0 {
		d = 5 * time.Second
	}
	deadline := time.Now().Add(d)
	return context.WithDeadline(ctx, deadline)
}

func normalizeHostHeader(host string) string {
	if h, _, err := net.SplitHostPort(host); err == nil {
		return strings.ToLower(h)
	}
	return strings.ToLower(host)
}