crossbar/internal/proxy/proxy.go

// Package proxy is the routing reverse proxy. It takes /{route}/v1/…, picks a host for the
// conversation from its ordered list using a lease table (or the health table alone), queues per
// (host, model), forwards the request streaming the answer back as it arrives, tees the response to
// read usage/timings, marks a host down when a forward fails, and records one accounting row per
// request.
package proxy

import (
	"bytes"
	"encoding/json"
	"errors"
	"io"
	"log/slog"
	"net/http"
	"strings"
	"time"

	"git.wntrmute.dev/kyle/crossbar/internal/config"
	"git.wntrmute.dev/kyle/crossbar/internal/fingerprint"
	"git.wntrmute.dev/kyle/crossbar/internal/health"
	"git.wntrmute.dev/kyle/crossbar/internal/lease"
	"git.wntrmute.dev/kyle/crossbar/internal/limiter"
	"git.wntrmute.dev/kyle/crossbar/internal/store"
)

const (
	MaxBody     = 16 << 20
	HostHeader  = "X-Crossbar-Host"
	LeaseHeader = "X-Crossbar-Lease" // "new" or "reused"
	RouteHeader = "X-Crossbar-Route" // client may name the route here instead of the path
)

// errBodyTooLarge is returned when a request body exceeds MaxBody during the model peek.
var errBodyTooLarge = errors.New("body too large")

// Health is what the proxy needs from the health table.
type Health interface {
	Get(name string) (health.Status, bool)
	MarkDown(name, reason string)
}

// Recorder is what the proxy needs to write an accounting row; *store.Store satisfies it.
type Recorder interface {
	RecordRequest(store.Request) error
}

// Handler forwards requests for a route to one of the route's healthy hosts, choosing by lease when
// one is configured and by health alone otherwise.
type Handler struct {
	cfg    *config.Config
	health Health
	leases *lease.Table
	lim    *limiter.Limiter
	rec    Recorder
	log    *slog.Logger
}

// New builds a Handler. A nil logger becomes slog.Default(). With a nil lease table it behaves like
// the v0 proxy: first healthy host, no queueing, no recording; nil limiter and recorder are likewise
// no-ops.
func New(cfg *config.Config, h Health, leases *lease.Table, lim *limiter.Limiter, rec Recorder, log *slog.Logger) *Handler {
	if log == nil {
		log = slog.Default()
	}
	return &Handler{cfg: cfg, health: h, leases: leases, lim: lim, rec: rec, log: log}
}

// SplitRoute takes the first path segment as the route. "/a/v1/x" -> ("a", "/v1/x", true); "/a" and
// "/a/" -> ("a", "/", true); "/", "//x", "", "noslash/v1" -> ("", "", false). The query string, if
// present, is kept in rest.
func SplitRoute(path string) (route, rest string, ok bool) {
	if path == "" || path[0] != '/' {
		return "", "", false
	}
	after := path[1:]
	slash := strings.IndexByte(after, '/')
	if slash == -1 {
		if after == "" {
			return "", "", false
		}
		return after, "/", true
	}
	if after[:slash] == "" {
		return "", "", false
	}
	return after[:slash], after[slash:], true
}

// Choose returns the first host in order that is healthy and lists model in Loaded; failing that,
// the first healthy host. ok is false if none. model may be "".
func Choose(hosts []string, model string, h Health) (string, bool) {
	for _, name := range hosts {
		s, ok := h.Get(name)
		if !ok || !s.Healthy || !hasModel(s.Loaded, model) {
			continue
		}
		return name, true
	}
	for _, name := range hosts {
		s, ok := h.Get(name)
		if ok && s.Healthy {
			return name, true
		}
	}
	return "", false
}

func hasModel(loaded []string, model string) bool {
	if model == "" {
		return false
	}
	for _, m := range loaded {
		if m == model {
			return true
		}
	}
	return false
}

// allowedPath reports whether rest may be proxied: under /v1/, or the two admin paths.
func allowedPath(rest string) bool {
	return strings.HasPrefix(rest, "/v1/") || rest == "/health" || rest == "/props"
}

// route resolves the route name and the upstream path (rest) from the request, honouring the
// optional route header. code is non-zero when the request must be answered; msg is the JSON error
// text for that code.
func (p *Handler) route(r *http.Request) (route, rest string, code int, msg string) {
	hdr := r.Header.Get(RouteHeader)
	if hdr != "" {
		rest := r.URL.Path
		// A path that also carries a (different) route name is a client mistake: the header is the
		// operator's intent, but the path disagrees.
		if rname, _, ok := SplitRoute(rest); ok {
			if _, known := p.cfg.Routes[rname]; known && rname != hdr {
				return "", "", http.StatusBadRequest, "conflicting route"
			}
		}
		if _, known := p.cfg.Routes[hdr]; !known {
			return "", "", http.StatusNotFound, "unknown route"
		}
		if !allowedPath(rest) {
			return "", "", http.StatusNotFound, "not found"
		}
		return hdr, rest, 0, ""
	}
	route, rest, ok := SplitRoute(r.URL.Path)
	if !ok {
		return "", "", http.StatusBadRequest, "missing route"
	}
	if _, known := p.cfg.Routes[route]; !known {
		return "", "", http.StatusNotFound, "unknown route"
	}
	if !allowedPath(rest) {
		return "", "", http.StatusNotFound, "not found"
	}
	return route, rest, 0, ""
}

// peekModel reads a non-GET/HEAD body up to MaxBody+1 bytes, restores it on the request, and returns
// the top-level "model" and the body itself (for fingerprinting). A non-JSON body or one without a
// model gives "". A body larger than MaxBody returns errBodyTooLarge.
func peekModel(r *http.Request) (string, []byte, error) {
	if r.Method == http.MethodGet || r.Method == http.MethodHead {
		return "", nil, nil
	}
	if r.Body == nil || r.Body == http.NoBody {
		return "", nil, nil
	}
	body, err := io.ReadAll(io.LimitReader(r.Body, MaxBody+1))
	if err != nil {
		return "", nil, err
	}
	if len(body) > MaxBody {
		return "", nil, errBodyTooLarge
	}
	r.Body = io.NopCloser(bytes.NewReader(body))
	r.ContentLength = int64(len(body))

	var req struct {
		Model string `json:"model"`
	}
	_ = json.Unmarshal(body, &req)
	return req.Model, body, nil
}

// ServeHTTP routes, fingerprints, leases a host, queues per (host, model), forwards with streaming,
// tees the response for usage/timings, and records one accounting row. Every error answer is JSON
// {"error":"…"}.
func (p *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	route, rest, code, msg := p.route(r)
	if code != 0 {
		p.writeError(w, code, msg)
		return
	}
	routeCfg := p.cfg.Routes[route]

	model, body, err := peekModel(r)
	if err != nil {
		p.writeError(w, http.StatusRequestEntityTooLarge, "body too large")
		return
	}
	if model == "" {
		model = routeCfg.DefaultModel
	}
	fp := fingerprint.Of(body)
	started := time.Now()

	// v0 compatibility path: no lease table, no limiter, no recording.
	if p.leases == nil {
		name, ok := Choose(routeCfg.Hosts, model, p.health)
		if !ok {
			p.writeError(w, http.StatusServiceUnavailable, "no healthy host")
			return
		}
		p.forward(w, r, route, name, "", rest, fp, model, started, 0)
		return
	}

	// Lease. The route's ordered host list is the candidate set.
	host, reused, err := p.leases.Acquire(lease.Key{Route: route, FP: fp, Model: model}, routeCfg.Hosts, time.Now())
	if err != nil {
		switch {
		case errors.Is(err, lease.ErrNoHost):
			p.writeError(w, http.StatusServiceUnavailable, "no healthy host")
		case errors.Is(err, lease.ErrPinnedDown):
			p.writeError(w, http.StatusServiceUnavailable, "pinned host down")
		default:
			p.writeError(w, http.StatusBadGateway, "upstream failed")
		}
		return
	}

	// Slot. A full queue is a 503; a context done while waiting means the client left.
	release, waited, err := p.lim.Acquire(r.Context(), host, model)
	if err != nil {
		if errors.Is(err, limiter.ErrQueueFull) {
			p.writeRecord(store.Request{
				Route:   route,
				FP:      fp,
				Model:   model,
				Host:    host,
				Started: started,
				TotalMs: time.Since(started).Milliseconds(),
				Status:  http.StatusServiceUnavailable,
				Err:     "queue full",
			})
			p.writeError(w, http.StatusServiceUnavailable, "queue full")
			return
		}
		p.log.Warn("request", "route", route, "host", host, "method", r.Method, "path", rest, "status", 499)
		return
	}
	defer release()

	p.forward(w, r, route, host, leaseState(reused), rest, fp, model, started, waited)
}