vibeproxy/internal/proxy/reverseproxy.go

package proxy

import (
	"context"
	"crypto/tls"
	"fmt"
	"io"
	"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) {
		// sharedTransport pools keep-alive connections to the backend. There's
		// an inherent race in any HTTP/1.1 keep-alive client: the backend can
		// decide to close an idle connection at (almost) the same moment we
		// pick it back out of the pool to send the next request, so the
		// request lands on a connection that's already gone. Go's transport
		// already retries this automatically in the common case, but not
		// always (e.g. lightweight backends like Wt's built-in wthttpd, which
		// can close more eagerly than Go expects - visible in LMS's own logs
		// as "handleReadBody(): while waiting for disconnect, received
		// unexpected data, closing"). For idempotent requests with no body,
		// retrying once here is safe and is standard practice for reverse
		// proxies (nginx's proxy_next_upstream, Traefik's retry middleware).
		if isRetryable(r) {
			log.Warn("backend request failed, retrying once", "proxy_host", rt.Model.Name, "backend", targetAddr, "err", err)
			retryReq := r.Clone(r.Context())
			if resp, rerr := sharedTransport.RoundTrip(retryReq); rerr == nil {
				defer resp.Body.Close()
				for k, vv := range resp.Header {
					for _, v := range vv {
						w.Header().Add(k, v)
					}
				}
				applyCustomHeaders(w.Header(), rt.Model.CustomHeaders, models.HeaderDirectionResponse)
				w.WriteHeader(resp.StatusCode)
				io.Copy(w, resp.Body)
				return
			}
		}

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

// isRetryable reports whether it's safe to silently retry a failed backend
// round trip. Only idempotent, bodyless requests qualify: retrying anything
// with a request body risks double-submitting it (the body reader may
// already be partially or fully consumed by the failed attempt).
func isRetryable(r *http.Request) bool {
	switch r.Method {
	case http.MethodGet, http.MethodHead, http.MethodOptions:
	default:
		return false
	}
	return r.Body == nil || r.Body == http.NoBody || r.ContentLength == 0
}

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