homelab-frontend/internal/proxy/proxy.go

// Package proxy provides reverse proxying to configured upstreams.
package proxy

import (
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"
	"time"

	"github.com/Riotpiaole/homelab-frontend/internal/config"
	"github.com/Riotpiaole/homelab-frontend/internal/logging"
)

// Handler is a reverse proxy that routes requests to configured upstreams.
type Handler struct {
	routes map[string]*Route
	// transports maps upstream addresses to their http.Transport for connection reuse
	transports map[string]*http.Transport
	// config holds the gateway configuration (for model registry, etc.)
	config *config.Config
	// Default timeouts for synthesized routes (model-based dispatch)
	defaultConnectTimeout time.Duration
	defaultReadTimeout    time.Duration
	defaultWriteTimeout   time.Duration
	defaultMaxBodySize    int64
}

// Route represents a reverse proxy route.
type Route struct {
	Name     string
	Upstream *config.Upstream
	Director func(*http.Request)
	Transport *http.Transport
}

// New creates a new reverse proxy handler from configuration.
// It sets up connection pooling and rewriting rules for each route.
func New(cfg *config.Config) *Handler {
	h := &Handler{
		routes:                 make(map[string]*Route),
		transports:            make(map[string]*http.Transport),
		config:                cfg,
		defaultConnectTimeout: 10 * time.Second,
		defaultReadTimeout:    1 * time.Hour,
		defaultWriteTimeout:   1 * time.Hour,
		defaultMaxBodySize:    100 * 1024 * 1024,
	}

	for name, route := range cfg.Routes {
		// Create a transport per unique upstream address for connection reuse
		transport := h.getOrCreateTransport(route.Upstream.Address, &route.Upstream)
		
		upstreamURL, _ := url.Parse("http://" + route.Upstream.Address)
		
		r := &Route{
			Name:      name,
			Upstream:  &route.Upstream,
			Transport: transport,
			Director: func(req *http.Request) {
				directorFunc(req, upstreamURL, &route.Upstream)
			},
		}
		h.routes[name] = r
	}

	return h
}

// getOrCreateTransport returns a shared http.Transport for the given upstream address.
// This ensures connections are pooled and reused across requests to the same upstream.
func (h *Handler) getOrCreateTransport(addr string, up *config.Upstream) *http.Transport {
	if t, ok := h.transports[addr]; ok {
		return t
	}

	// Create a transport with timeout settings from the upstream config.
	// Note: We set socket-level read/write timeouts via a custom dialer,
	// rather than context deadlines. Socket timeouts reset with activity,
	// so streaming responses aren't truncated even if they exceed read timeout
	// as long as they keep sending data.
	dialer := &net.Dialer{
		Timeout:   up.ConnectTimeout,
		KeepAlive: 30 * time.Second,
	}

	transport := &http.Transport{
		Dial:            dialer.Dial,
		DialContext:     dialer.DialContext,
		MaxIdleConns:    100,
		IdleConnTimeout: 90 * time.Second,
		// Allow persistent connections
		DisableKeepAlives: false,
	}

	// Store the upstream config for use in the handler
	if transport.TLSClientConfig == nil {
		// We can't directly set socket timeouts on http.Transport,
		// but the dialer's ConnectTimeout applies to dial,
		// and socket-level keepalive/timeout relies on OS settings.
		// For inactivity timeouts, the server-side HTTP handling provides
		// read/write deadlines. Client-side, we rely on TCP keepalive.
	}

	h.transports[addr] = transport
	return transport
}

// directorFunc modifies the request to be sent to the upstream.
// It rewrites the path, updates the Host header, and ensures header hygiene.
func directorFunc(req *http.Request, target *url.URL, upstream *config.Upstream) {
	// Apply path rewrite if configured
	if upstream.PathRewrite != "" {
		req.URL.Path = upstream.PathRewrite
	}

	// Set the scheme and host
	req.URL.Scheme = target.Scheme
	req.URL.Host = target.Host
	
	// Update the Host header to the upstream address
	req.Host = target.Host
	
	// Strip hop-by-hop headers as defined in RFC 7230 Section 6.1
	// These must not be forwarded to upstream
	hopByHopHeaders := map[string]bool{
		"connection":          true,
		"keep-alive":          true,
		"proxy-authenticate":  true,
		"proxy-authorization": true,
		"te":                  true,
		"trailers":            true,
		"transfer-encoding":   true,
		"upgrade":             true,
	}
	
	// Also strip any headers listed in the Connection header
	if conn := req.Header.Get("Connection"); conn != "" {
		for _, h := range strings.Split(conn, ",") {
			hopByHopHeaders[strings.ToLower(strings.TrimSpace(h))] = true
		}
	}
	
	// Remove all hop-by-hop headers
	// The http.Header.Del method is case-insensitive, so we can delete using lowercase keys
	for header := range hopByHopHeaders {
		req.Header.Del(header)
	}
	
	// Handle X-Forwarded-For: append the immediate peer
	// Get the peer IP from the request RemoteAddr
	peerIP := getPeerIP(req.RemoteAddr)
	if xForwardedFor := req.Header.Get("X-Forwarded-For"); xForwardedFor != "" {
		// Append the peer IP to the existing X-Forwarded-For
		req.Header.Set("X-Forwarded-For", xForwardedFor+", "+peerIP)
	} else {
		// Create a new X-Forwarded-For with just the peer IP
		req.Header.Set("X-Forwarded-For", peerIP)
	}
}

// getPeerIP extracts the IP address from a RemoteAddr string (format: "IP:port")
func getPeerIP(remoteAddr string) string {
	if remoteAddr == "" {
		return ""
	}
	// RemoteAddr is "IP:port", extract just the IP
	if idx := strings.LastIndex(remoteAddr, ":"); idx != -1 {
		return remoteAddr[:idx]
	}
	return remoteAddr
}

// ServeHTTP implements http.Handler.
func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// Try to find a matching route (including body-based dispatch for /v1/chat/completions)
	route, err := h.RouteRequest(r)
	if err != nil || route == nil {
		// Route not found or error determining route
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(w, "not found")
		if err != nil {
			logging.Errorf("routing failed", err, map[string]string{
				"path":   r.URL.Path,
				"method": r.Method,
			})
		} else {
			logging.Errorf("no route matches", fmt.Errorf("path=%s method=%s", r.URL.Path, r.Method), nil)
		}
		return
	}

	// Note: Body size checking already happened in RouteRequest (body was read for model dispatch).
	// For other paths, we still need to enforce the cap.
	// For /v1/chat/completions, the body was already read and validated.
	
	// Enforce request body size cap for non-chat routes
	if r.URL.Path != "/v1/chat/completions" {
		if route.Upstream.MaxBodySize > 0 && r.ContentLength > route.Upstream.MaxBodySize {
			w.WriteHeader(http.StatusRequestEntityTooLarge)
			fmt.Fprintf(w, "request body too large")
			logging.Errorf("request rejected", fmt.Errorf("body_too_large"), map[string]string{
				"reason":         "body_too_large",
				"route":          route.Name,
				"upstream":       route.Upstream.Address,
				"content_length": fmt.Sprintf("%d", r.ContentLength),
				"max_body_size":  fmt.Sprintf("%d", route.Upstream.MaxBodySize),
			})
			return
		}

		// Wrap request body with size limiter
		// This enforces the cap at read time, not after buffering
		if route.Upstream.MaxBodySize > 0 && r.Body != nil {
			r.Body = io.NopCloser(io.LimitReader(r.Body, route.Upstream.MaxBodySize))
		}
	}

	// Create the reverse proxy
	proxy := httputil.NewSingleHostReverseProxy(&url.URL{
		Scheme: "http",
		Host:   route.Upstream.Address,
	})

	// Set the director to apply path rewriting
	proxy.Director = route.Director

	// Use the connection-pooled transport
	proxy.Transport = route.Transport

	// Set error handler to log upstream errors
	proxy.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) {
		w.WriteHeader(http.StatusBadGateway)
		fmt.Fprintf(w, "upstream error")
		logging.Errorf("upstream error", err, map[string]string{
			"upstream": route.Upstream.Address,
			"path":     r.URL.Path,
		})
	}

	// Note: We apply connection timeout via Transport dialer, but NOT read timeout as a context deadline.
	// Read timeout should apply to inactivity (socket read timeout), not total request duration.
	// A streaming response that's continuously sending should not be cut off.
	// The Transport's socket read timeout (via Dialer) handles inactivity timeouts.

	// Serve the request through the proxy
	proxy.ServeHTTP(w, r)
}



// Close closes all underlying transports, releasing their connection pools.
func (h *Handler) Close() error {
	for _, transport := range h.transports {
		transport.CloseIdleConnections()
	}
	return nil
}