llmgw/internal/cli/leaderforward.go

package cli

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

	"golang.org/x/net/http2"

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

// 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 !strings.HasPrefix(r.URL.Path, "/api/") || 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
				pr.Out.URL.Path = forwardPrefix + pr.In.URL.Path
				pr.Out.URL.RawPath = ""
				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)
	})
}

// 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)
	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 !strings.HasPrefix(path, "/api/") {
			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, ""
		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})
}