vibeproxy/internal/proxy/handler.go

package proxy

import (
	"database/sql"
	"log/slog"
	"net"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/owen/vibeproxy/internal/access"
	"github.com/owen/vibeproxy/internal/auth"
	"github.com/owen/vibeproxy/internal/certificates"
	"github.com/owen/vibeproxy/internal/logging"
	"github.com/owen/vibeproxy/internal/metrics"
	"github.com/owen/vibeproxy/internal/models"
	"github.com/owen/vibeproxy/internal/runtime"
)

// maxRequestBodyBytes bounds request bodies to protect the proxy itself;
// it is generous enough for large uploads while still being a hard limit.
const maxRequestBodyBytes = 5 << 30 // 5 GiB

// Handler is the single entry point for both the HTTP and HTTPS listeners.
// It resolves the target proxy host, enforces access control and
// force-HTTPS/redirect rules, picks a healthy backend, and proxies the
// request — or serves an ACME HTTP-01 challenge.
type Handler struct {
	store   *runtime.Store
	certMgr *certificates.Manager
	reqLog  *logging.RequestLogger
	db      *sql.DB
	log     *slog.Logger
}

func NewHandler(store *runtime.Store, certMgr *certificates.Manager, reqLog *logging.RequestLogger, db *sql.DB, log *slog.Logger) *Handler {
	return &Handler{store: store, certMgr: certMgr, reqLog: reqLog, db: db, log: log}
}

func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if h.certMgr.ServeChallenge(w, r) {
		return
	}

	r.Body = http.MaxBytesReader(w, r.Body, maxRequestBodyBytes)

	snap := h.store.Current()
	rt := snap.LookupHost(r.Host)
	if rt == nil {
		http.Error(w, "404 - no proxy host configured for this domain", http.StatusNotFound)
		return
	}
	if !rt.Model.Enabled {
		http.Error(w, "404 - not found", http.StatusNotFound)
		return
	}

	clientIP := ClientIP(r, snap)

	if rt.AccessList != nil {
		if !h.checkAccess(w, r, rt.AccessList, clientIP) {
			return
		}
	}

	if rt.Model.ForceHTTPS && r.TLS == nil {
		target := "https://" + normalizeHostHeader(r.Host) + r.URL.RequestURI()
		http.Redirect(w, r, target, http.StatusMovedPermanently)
		return
	}

	if rt.Model.Type == models.ProxyHostTypeRedirect {
		h.serveRedirect(w, r, rt)
		return
	}

	h.serveProxy(w, r, rt, snap, clientIP)
}

func (h *Handler) checkAccess(w http.ResponseWriter, r *http.Request, list *access.List, clientIP string) bool {
	ip := parseIPLoose(clientIP)
	if ip == nil || !list.Allowed(ip) {
		http.Error(w, "403 - forbidden", http.StatusForbidden)
		return false
	}
	if list.RequiresAuth() {
		username, password, ok := r.BasicAuth()
		if !ok {
			w.Header().Set("WWW-Authenticate", `Basic realm="restricted"`)
			http.Error(w, "401 - authentication required", http.StatusUnauthorized)
			return false
		}
		hash, exists := list.Users[username]
		if !exists {
			w.Header().Set("WWW-Authenticate", `Basic realm="restricted"`)
			http.Error(w, "401 - authentication required", http.StatusUnauthorized)
			return false
		}
		valid, err := auth.VerifyPassword(password, hash)
		if err != nil || !valid {
			w.Header().Set("WWW-Authenticate", `Basic realm="restricted"`)
			http.Error(w, "401 - authentication required", http.StatusUnauthorized)
			return false
		}
	}
	return true
}

func (h *Handler) serveRedirect(w http.ResponseWriter, r *http.Request, rt *runtime.ProxyHostRuntime) {
	target := rt.Model.RedirectTargetURL
	if target == "" {
		http.Error(w, "502 - redirect host has no target configured", http.StatusBadGateway)
		return
	}
	target = strings.TrimRight(target, "/")
	if rt.Model.RedirectPreservePath {
		target += r.URL.Path
	}
	if rt.Model.RedirectPreserveQS && r.URL.RawQuery != "" {
		target += "?" + r.URL.RawQuery
	}
	code := rt.Model.RedirectCode
	switch code {
	case 301, 302, 307, 308:
	default:
		code = http.StatusMovedPermanently
	}
	http.Redirect(w, r, target, code)
}

func (h *Handler) serveProxy(w http.ResponseWriter, r *http.Request, rt *runtime.ProxyHostRuntime, snap *runtime.Snapshot, clientIP string) {
	if rt.Pool == nil {
		http.Error(w, "502 - proxy host has no backends", http.StatusBadGateway)
		return
	}
	backend, err := rt.Pool.Pick()
	if err != nil {
		h.log.Warn("no healthy backend", "proxy_host", rt.Model.Name, "err", err)
		metrics.HTTPRequestsTotal.WithLabelValues(rt.Model.Name, r.Method, "503").Inc()
		http.Error(w, "503 - no healthy backend available", http.StatusServiceUnavailable)
		return
	}
	backend.Acquire()
	defer backend.Release()

	metrics.ActiveConnections.WithLabelValues("http").Inc()
	defer metrics.ActiveConnections.WithLabelValues("http").Dec()

	rp := buildReverseProxy(rt, backend, snap, h.log)

	readTimeout := time.Duration(rt.Model.ReadTimeoutMS) * time.Millisecond
	if readTimeout <= 0 {
		readTimeout = 60 * time.Second
	}
	ctx, cancel := withConnectTimeout(r.Context(), readTimeout)
	defer cancel()
	r = r.WithContext(ctx)

	rec := &statusRecorder{ResponseWriter: w}
	start := time.Now()
	rp.ServeHTTP(rec, r)
	elapsed := time.Since(start)

	backendAddr := backend.Addr() + ":" + strconv.Itoa(backend.Model.Port)
	metrics.HTTPRequestsTotal.WithLabelValues(rt.Model.Name, r.Method, strconv.Itoa(rec.status)).Inc()
	metrics.HTTPRequestDuration.WithLabelValues(rt.Model.Name).Observe(elapsed.Seconds())
	metrics.BackendRequestDuration.WithLabelValues(rt.Model.Name, backendAddr).Observe(elapsed.Seconds())
	metrics.BytesTransferred.WithLabelValues("out", "http").Add(float64(rec.bytes))

	hostID := rt.Model.ID
	h.reqLog.Log(models.RequestLog{
		Timestamp:      start,
		ProxyHostID:    &hostID,
		ProxyHostName:  rt.Model.Name,
		ClientIP:       clientIP,
		Method:         r.Method,
		Host:           r.Host,
		Path:           r.URL.Path,
		Status:         rec.status,
		ResponseSize:   rec.bytes,
		Backend:        backendAddr,
		ResponseTimeMS: elapsed.Milliseconds(),
		UserAgent:      r.UserAgent(),
	})
}

func parseIPLoose(s string) net.IP {
	if h, _, err := net.SplitHostPort(s); err == nil {
		s = h
	}
	return net.ParseIP(s)
}