llmgw/internal/cli/leaderforward.go

package cli

import (
	"context"
	"crypto/x509"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"slices"
	"strings"
	"time"

	"golang.org/x/net/http2"

	"gitea.mynodes.xyz/ilker/llmgw/internal/clientip"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
	"gitea.mynodes.xyz/ilker/llmgw/internal/tlsconf"
)

// forwardPrefix is where the leader answers a standby's admin API, on the
// agent port: the one port control planes already reach each other on, with
// the cluster's certificates both ways. The admin port itself listens on
// loopback, and should — it serves pprof.
const forwardPrefix = "/_admin"

// Headers a forwarded request and its answer carry: which standby asked, and
// which control plane answered.
const (
	hdrForwardedBy = "X-Llmgw-Forwarded-By"
	hdrAnsweredBy  = "X-Llmgw-Answered-By"
)

// leaderForward sends a standby's admin API to the leader.
//
// Everything the console and ctl show is the leader's: nodes, replicas and
// jobs live in its memory, put there by the agent streams only it holds. A
// standby answering from its own would show an empty cluster, and accept
// changes it cannot carry out. So a standby answers nothing under /api/
// itself; the leader does, and the console looks the same on either machine.
type leaderForward struct {
	st   *store.DB
	lead interface{ IsLeader() bool }
	self string
	tls  tlsconf.Config
	log  *slog.Logger
	// routes, with router.active_active, is the leader's replicas this
	// standby routes to; nil when every request goes to the leader.
	routes *standbyRoutes
}

// standby wraps the admin handler: /api/ goes to the leader while this
// process is not it.
func (f *leaderForward) standby(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if adminStaysHere(r.URL.Path) || f.lead.IsLeader() {
			next.ServeHTTP(w, r)
			return
		}
		target, name, err := f.leader(r.Context())
		if err != nil {
			writeForwardError(w, http.StatusServiceUnavailable,
				"this control plane is standing by, and cannot reach the leader: "+err.Error())
			return
		}
		tc, err := f.tls.ClientConfig(target.Hostname())
		if err != nil {
			writeForwardError(w, http.StatusInternalServerError, err.Error())
			return
		}
		p := &httputil.ReverseProxy{
			Rewrite: func(pr *httputil.ProxyRequest) {
				pr.Out.URL.Scheme, pr.Out.URL.Host = "https", target.Host
				// Escaped as it came: a name with a slash in it, Demo/hr
				// as Demo%2Fhr, is one segment, not two.
				pr.Out.URL.Path = forwardPrefix + pr.In.URL.Path
				pr.Out.URL.RawPath = forwardPrefix + pr.In.URL.EscapedPath()
				pr.Out.Host = target.Host
				// The leader's own token is not this one's, and the
				// certificate is what it believes.
				pr.Out.Header.Del("Authorization")
				pr.Out.Header.Set(hdrForwardedBy, f.self)
			},
			Transport: &http2.Transport{TLSClientConfig: tc},
			// Streams, such as a job's progress, pass through as they come.
			FlushInterval: -1,
			ErrorHandler: func(w http.ResponseWriter, _ *http.Request, err error) {
				writeForwardError(w, http.StatusBadGateway, fmt.Sprintf("the leader, %s, did not answer: %v", name, err))
			},
		}
		p.ServeHTTP(w, r)
	})
}

// adminStaysHere are the admin paths about this process rather than the
// cluster: its profile, its own metrics, its own health. Never forwarded, and
// never answered for another — pprof dumps the heap, keys and all.
func adminStaysHere(path string) bool {
	return strings.HasPrefix(path, "/debug/") || path == "/metrics" || path == "/health"
}

// publicStaysHere are the client port's paths that must describe this
// process: a load balancer and keepalived ask them which machine to send
// traffic to, and a standby saying the leader is ready would draw it.
func publicStaysHere(path string) bool {
	return path == "/health" || path == "/ready" || path == "/leader" || path == "/metrics"
}

// public wraps the client-facing handler: on a standby, everything but its
// own health goes to the leader's client port — the portal, its admin pages
// and the API, which on a standby would show an empty cluster and serve
// nothing. Plain HTTP, as the port is; the leader trusts a control plane's
// X-Forwarded-For, so a client is still known by its own address.
func (f *leaderForward) public(next http.Handler, port string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Once only: a request another control plane forwarded is answered
		// here, leading or not, rather than sent on while the lease moves.
		if publicStaysHere(r.URL.Path) || f.lead.IsLeader() || r.Header.Get(hdrForwardedBy) != "" ||
			f.answersHere(r.URL.Path) {
			next.ServeHTTP(w, r)
			return
		}
		target, _, err := f.leader(r.Context())
		if err != nil || target.Hostname() == "" {
			next.ServeHTTP(w, r)
			return
		}
		host := net.JoinHostPort(target.Hostname(), port)
		p := &httputil.ReverseProxy{
			Rewrite: func(pr *httputil.ProxyRequest) {
				pr.Out.URL.Scheme, pr.Out.URL.Host = "http", host
				// The name the client used, for the portal's links and cookies.
				pr.Out.Host = pr.In.Host
				peer, _, _ := net.SplitHostPort(pr.In.RemoteAddr)
				xff := pr.In.Header.Get("X-Forwarded-For")
				if xff != "" {
					xff += ", "
				}
				pr.Out.Header.Set("X-Forwarded-For", xff+peer)
				for _, h := range []string{"X-Forwarded-Proto", "X-Forwarded-Host"} {
					if v := pr.In.Header.Get(h); v != "" {
						pr.Out.Header.Set(h, v)
					}
				}
				pr.Out.Header.Set(hdrForwardedBy, f.self)
			},
			FlushInterval: -1,
			ErrorHandler: func(w http.ResponseWriter, _ *http.Request, err error) {
				writeForwardError(w, http.StatusBadGateway, "this control plane is standing by, and the leader did not answer: "+err.Error())
			},
		}
		p.ServeHTTP(w, r)
	})
}

// answersHere is whether a standby answers this client request itself: the
// API under /v1, with router.active_active, while it has the leader's
// replicas to send it to. Not the portal, its admin pages or /mcp, whose
// llmgw:admin tools read the cluster from memory only the leader has; and
// not while the leader has gone unheard, when it is passed on as before.
func (f *leaderForward) answersHere(path string) bool {
	return f.routes != nil && strings.HasPrefix(path, "/v1/") && f.routes.serving()
}

// trustControlPlanes keeps the client port trusting the other control
// planes' X-Forwarded-For, besides router.trusted_proxies: a standby passes
// requests on, and the client is the address it was given, not the standby.
// Read again every half minute, for one added since.
func (f *leaderForward) trustControlPlanes(ctx context.Context, configured []string) {
	for {
		list := slices.Clone(configured)
		if cps, err := f.st.ControlPlanes(ctx); err == nil {
			for _, cp := range cps {
				if u, err := url.Parse(cp.URL); err == nil && u.Hostname() != "" {
					list = append(list, u.Hostname())
				}
			}
		}
		clientip.SetTrusted(list)
		select {
		case <-ctx.Done():
			return
		case <-time.After(30 * time.Second):
		}
	}
}

// leader is where the lease's holder is reached, and its name.
func (f *leaderForward) leader(ctx context.Context) (*url.URL, string, error) {
	ctx, cancel := context.WithTimeout(ctx, 3*time.Second)
	defer cancel()
	lease, err := f.st.Leadership(ctx, control.ControlRole)
	if err != nil || lease.Owner == "" || time.Now().After(lease.ExpiresAt) {
		return nil, "", errors.New("no control plane holds the lease right now")
	}
	name := control.HolderName(lease.Owner)
	if name == f.self {
		// This machine's previous process, whose lease has not run out.
		return nil, name, errors.New("the lease is still held by this machine's previous run; it runs out within seconds")
	}
	cps, err := f.st.ControlPlanes(ctx)
	if err != nil {
		return nil, name, err
	}
	for _, cp := range cps {
		if cp.Name == name {
			u, err := url.Parse(cp.URL)
			return u, name, err
		}
	}
	return nil, name, fmt.Errorf("%s leads, and is not among the control planes the store lists", name)
}

// leaderSide answers forwarded requests, on the agent port: only from another
// control plane of this cluster, only while this one leads, and only the API.
func (f *leaderForward) leaderSide(admin http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		path := strings.TrimPrefix(r.URL.Path, forwardPrefix)
		if adminStaysHere(path) {
			http.NotFound(w, r)
			return
		}
		if err := f.fromControlPlane(r); err != nil {
			writeForwardError(w, http.StatusForbidden, err.Error())
			return
		}
		if !f.lead.IsLeader() {
			// Not passed on again: two standbys must not bounce a request
			// between them while the lease moves.
			writeForwardError(w, http.StatusServiceUnavailable, "this control plane no longer leads; ask again")
			return
		}
		r2 := r.Clone(r.Context())
		r2.URL.Path, r2.URL.RawPath = path, strings.TrimPrefix(r.URL.EscapedPath(), forwardPrefix)
		w.Header().Set(hdrAnsweredBy, f.self)
		admin.ServeHTTP(w, r2)
	})
}

// fromControlPlane checks the client certificate names a control plane the
// store lists. A node's certificate is from the same authority, and a node
// must not reach the admin API by presenting it.
func (f *leaderForward) fromControlPlane(r *http.Request) error {
	if r.TLS == nil || len(r.TLS.VerifiedChains) == 0 {
		return errors.New("a control plane's certificate is required")
	}
	cert := r.TLS.VerifiedChains[0][0]
	if !slices.Contains(cert.ExtKeyUsage, x509.ExtKeyUsageServerAuth) {
		return errors.New("not a control plane's certificate")
	}
	cps, err := f.st.ControlPlanes(r.Context())
	if err != nil {
		return err
	}
	for _, cp := range cps {
		if cp.Name == cert.Subject.CommonName {
			return nil
		}
	}
	return fmt.Errorf("%s is not a control plane of this cluster", cert.Subject.CommonName)
}

func writeForwardError(w http.ResponseWriter, code int, msg string) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(code)
	_ = json.NewEncoder(w).Encode(map[string]string{"error": msg})
}