homelab-frontend/internal/proxy/proxy.go

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

import (
	"encoding/json"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"sort"
	"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
}

// Error types for model validation
type modelValidationError struct {
	Kind    string // "invalid_json", "missing_model", "unknown_model"
	Message string
	Model   string // only for unknown_model
}

func (e *modelValidationError) Error() string {
	return e.Message
}

// RFC 9457 Problem Details
type problemDetail struct {
	Type       string        `json:"type"`
	Title      string        `json:"title"`
	Status     int           `json:"status"`
	Detail     string        `json:"detail"`
	ValidModels []string     `json:"valid_models,omitempty"`
}

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

// writeProblemDetail writes an RFC 9457 problem detail response.
func writeProblemDetail(w http.ResponseWriter, status int, problemType, title, detail string, validModels []string) {
	w.Header().Set("Content-Type", "application/problem+json")
	w.WriteHeader(status)

	problem := problemDetail{
		Type:        problemType,
		Title:       title,
		Status:      status,
		Detail:      detail,
		ValidModels: validModels,
	}

	json.NewEncoder(w).Encode(problem)
}

// ServeHTTP implements http.Handler.
func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// Handle /v1/models endpoint (no routing needed, derived from config)
	if r.URL.Path == "/v1/models" && r.Method == "GET" {
		h.handleModelsEndpoint(w, r)
		return
	}

	// Try to find a matching route (including body-based dispatch for /v1/chat/completions)
	route, err := h.RouteRequest(r)

	// Check if this is a model validation error (from body-based dispatch)
	if validationErr, ok := err.(*modelValidationError); ok {
		// This is a client error, not a routing error
		var status int
		var problemType string
		var title string
		var detail string

		switch validationErr.Kind {
		case "invalid_json":
			status = http.StatusBadRequest
			problemType = "https://api.example.com/problems/invalid-request-body"
			title = "Invalid Request Body"
			detail = validationErr.Message
		case "missing_model", "empty_model", "null_model":
			status = http.StatusBadRequest
			problemType = "https://api.example.com/problems/missing-model"
			title = "Missing Model"
			detail = "The 'model' field is required and must be a non-empty string"
		case "unknown_model":
			status = http.StatusBadRequest
			problemType = "https://api.example.com/problems/unknown-model"
			title = "Unknown Model"
			detail = fmt.Sprintf("Model %q is not available. See valid_models for available options.", validationErr.Model)
		default:
			status = http.StatusBadRequest
			problemType = "https://api.example.com/problems/invalid-request"
			title = "Invalid Request"
			detail = validationErr.Message
		}

		// Get list of valid models (only for model-related errors)
		var validModels []string
		if validationErr.Kind == "unknown_model" || validationErr.Kind == "missing_model" || validationErr.Kind == "empty_model" || validationErr.Kind == "null_model" {
			validModels = h.getValidModels()
		}

		writeProblemDetail(w, status, problemType, title, detail, validModels)
		logging.Errorf("client error", validationErr, map[string]string{
			"path":   r.URL.Path,
			"method": r.Method,
			"reason": validationErr.Kind,
		})
		return
	}

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



// getValidModels returns a sorted list of all configured model names.
func (h *Handler) getValidModels() []string {
	var models []string
	for name := range h.config.Models {
		models = append(models, name)
	}
	sort.Strings(models)
	return models
}

// modelsListResponse represents the response for GET /v1/models
type modelsListResponse struct {
	Object string              `json:"object"`
	Data   []modelsListEntry   `json:"data"`
}

// modelsListEntry represents a single model in the list
type modelsListEntry struct {
	ID      string `json:"id"`
	Object  string `json:"object"`
	OwnedBy string `json:"owned_by"`
	Created int64  `json:"created"`
}

// handleModelsEndpoint serves GET /v1/models
// Returns a list of all configured models, derived from config not hardcoded
func (h *Handler) handleModelsEndpoint(w http.ResponseWriter, r *http.Request) {
	// Get all model names from config
	modelNames := h.getValidModels()

	// Build the response
	data := make([]modelsListEntry, len(modelNames))
	for i, name := range modelNames {
		data[i] = modelsListEntry{
			ID:      name,
			Object:  "model",
			OwnedBy: "api.riotpiao.com",
			Created: 1700000000, // Fixed timestamp; can be made configurable if needed
		}
	}

	response := modelsListResponse{
		Object: "list",
		Data:   data,
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)
	json.NewEncoder(w).Encode(response)
}

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