llmgw/internal/cli/leaderforward_test.go

package cli

import (
	"encoding/json"
	"io"
	"log/slog"
	"net"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"strings"
	"sync/atomic"
	"testing"
	"time"

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

type leading struct{ v atomic.Bool }

func (l *leading) IsLeader() bool { return l.v.Load() }

// forwardCluster is two control planes of one cluster on one store: a
// leader whose agent port answers forwarded admin requests, and a standby.
type forwardCluster struct {
	st              *store.DB
	dir             string
	leaderLeads     *leading
	standby, leader *leaderForward
	leaderURL       string
}

func newForwardCluster(t *testing.T) *forwardCluster {
	t.Helper()
	dir := t.TempDir()
	ca, err := tlsconf.NewCA("test CA")
	if err != nil {
		t.Fatal(err)
	}
	if err := ca.Write(dir); err != nil {
		t.Fatal(err)
	}
	for _, name := range []string{"ctl-a", "ctl-b"} {
		leaf, err := ca.IssuePeer(name, []string{"127.0.0.1"})
		if err != nil {
			t.Fatal(err)
		}
		if err := leaf.Write(dir, name); err != nil {
			t.Fatal(err)
		}
	}
	node, err := ca.IssueClient("node1")
	if err != nil {
		t.Fatal(err)
	}
	if err := node.Write(dir, "node1"); err != nil {
		t.Fatal(err)
	}
	cfgFor := func(name string) tlsconf.Config {
		return tlsconf.Config{CA: filepath.Join(dir, "ca.pem"), Cert: filepath.Join(dir, name+".pem"), Key: filepath.Join(dir, name+"-key.pem")}
	}

	st, err := store.Open(t.Context(), filepath.Join(dir, "s.db"))
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { st.Close() })

	c := &forwardCluster{st: st, dir: dir, leaderLeads: &leading{}}
	c.leaderLeads.v.Store(true)
	log := slog.New(slog.DiscardHandler)
	c.leader = &leaderForward{st: st, lead: c.leaderLeads, self: "ctl-a", tls: cfgFor("ctl-a"), log: log}
	standbyLeads := &leading{}
	c.standby = &leaderForward{st: st, lead: standbyLeads, self: "ctl-b", tls: cfgFor("ctl-b"), log: log}

	// The leader's admin API: says whose it is, and who it was asked by.
	admin := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_ = json.NewEncoder(w).Encode(map[string]string{"path": r.URL.Path, "by": r.Header.Get(hdrForwardedBy),
			"auth": r.Header.Get("Authorization"), "actor": r.Header.Get("X-Llmgw-Actor")})
	})
	mux := http.NewServeMux()
	mux.Handle(forwardPrefix+"/", c.leader.leaderSide(admin))
	srv := httptest.NewUnstartedServer(mux)
	tc, err := cfgFor("ctl-a").ServerConfig()
	if err != nil {
		t.Fatal(err)
	}
	srv.TLS = tc
	srv.EnableHTTP2 = true
	srv.StartTLS()
	t.Cleanup(srv.Close)
	c.leaderURL = srv.URL

	for _, cp := range []store.ControlPlane{{URL: srv.URL, Name: "ctl-a"}, {URL: "https://127.0.0.1:1", Name: "ctl-b"}} {
		if err := st.PutControlPlane(t.Context(), cp); err != nil {
			t.Fatal(err)
		}
	}
	if _, err := st.AcquireLeadership(t.Context(), control.ControlRole, "ctl-a-0123456789abcdef", time.Minute); err != nil {
		t.Fatal(err)
	}
	return c
}

// ask sends one request through the standby's admin API.
func (c *forwardCluster) ask(t *testing.T, path string) (*http.Response, map[string]string) {
	t.Helper()
	local := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		_ = json.NewEncoder(w).Encode(map[string]string{"path": "local"})
	})
	srv := httptest.NewServer(c.standby.standby(local))
	defer srv.Close()
	req, _ := http.NewRequestWithContext(t.Context(), http.MethodGet, srv.URL+path, nil)
	req.Header.Set("Authorization", "Bearer the-standbys-own-token")
	req.Header.Set("X-Llmgw-Actor", "root@ctl-b")
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	b, _ := io.ReadAll(resp.Body)
	var got map[string]string
	_ = json.Unmarshal(b, &got)
	return resp, got
}

func TestAStandbysAdminAPIIsTheLeaders(t *testing.T) {
	c := newForwardCluster(t)
	resp, got := c.ask(t, "/api/v1/summary")
	if resp.StatusCode != http.StatusOK || got["path"] != "/api/v1/summary" {
		t.Fatalf("%d %v", resp.StatusCode, got)
	}
	if resp.Header.Get(hdrAnsweredBy) != "ctl-a" || got["by"] != "ctl-b" {
		t.Errorf("answered by %q, asked by %q", resp.Header.Get(hdrAnsweredBy), got["by"])
	}
	if got["auth"] != "" {
		t.Error("the standby's token was passed on to the leader")
	}
	if got["actor"] != "root@ctl-b" {
		t.Errorf("who asked was lost: %q", got["actor"])
	}
	// The dashboard's pages are the leader's too.
	if _, got := c.ask(t, "/nodes/cuda01"); got["path"] != "/nodes/cuda01" {
		t.Errorf("a dashboard page was answered here: %v", got)
	}
	// What is about this process stays on it: pprof, its metrics.
	for _, p := range []string{"/debug/pprof/", "/metrics"} {
		if _, got := c.ask(t, p); got["path"] != "local" {
			t.Errorf("%s was forwarded: %v", p, got)
		}
	}
}

func TestTheLeaderAnswersOnlyControlPlanesAndOnlyWhileItLeads(t *testing.T) {
	c := newForwardCluster(t)
	call := func(cert string) int {
		cfg := tlsconf.Config{CA: filepath.Join(c.dir, "ca.pem"), Cert: filepath.Join(c.dir, cert+".pem"), Key: filepath.Join(c.dir, cert+"-key.pem")}
		tc, err := cfg.ClientConfig("127.0.0.1")
		if err != nil {
			t.Fatal(err)
		}
		tc.NextProtos = nil
		client := &http.Client{Transport: &http.Transport{TLSClientConfig: tc}}
		resp, err := client.Get(c.leaderURL + forwardPrefix + "/api/v1/summary")
		if err != nil {
			t.Fatal(err)
		}
		resp.Body.Close()
		return resp.StatusCode
	}
	if code := call("ctl-b"); code != http.StatusOK {
		t.Errorf("a control plane was refused: %d", code)
	}
	// A node's certificate is from the same authority, and is no way in.
	if code := call("node1"); code != http.StatusForbidden {
		t.Errorf("a node reached the admin API: %d", code)
	}
	c.leaderLeads.v.Store(false)
	if code := call("ctl-b"); code != http.StatusServiceUnavailable {
		t.Errorf("a control plane that no longer leads answered: %d", code)
	}
}

func TestAStandbyWithNoLeaderSaysSo(t *testing.T) {
	c := newForwardCluster(t)
	if err := c.st.ReleaseLeadership(t.Context(), control.ControlRole, "ctl-a-0123456789abcdef"); err != nil {
		t.Fatal(err)
	}
	resp, got := c.ask(t, "/api/v1/summary")
	if resp.StatusCode != http.StatusServiceUnavailable || !strings.Contains(got["error"], "standing by") {
		t.Errorf("%d %v", resp.StatusCode, got)
	}
}

func TestAStandbysPortalIsTheLeaders(t *testing.T) {
	c := newForwardCluster(t)
	// The leader's client port, on the address the store lists for it.
	var got http.Header
	var gotHost string
	leader := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		got, gotHost = r.Header.Clone(), r.Host
		_, _ = w.Write([]byte("leader " + r.URL.Path))
	}))
	defer leader.Close()
	_, port, _ := net.SplitHostPort(strings.TrimPrefix(leader.URL, "http://"))

	local := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _, _ = w.Write([]byte("standby " + r.URL.Path)) })
	srv := httptest.NewServer(c.standby.public(local, port))
	defer srv.Close()
	get := func(path string, hdr map[string]string) string {
		req, _ := http.NewRequestWithContext(t.Context(), http.MethodGet, srv.URL+path, nil)
		req.Host = "llmgw.example"
		for k, v := range hdr {
			req.Header.Set(k, v)
		}
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			t.Fatal(err)
		}
		defer resp.Body.Close()
		b, _ := io.ReadAll(resp.Body)
		return string(b)
	}

	if body := get("/admin/replicas", map[string]string{"X-Forwarded-For": "203.0.113.7", "X-Forwarded-Proto": "https"}); body != "leader /admin/replicas" {
		t.Fatalf("the portal was answered by %q", body)
	}
	if gotHost != "llmgw.example" || got.Get("X-Forwarded-Proto") != "https" ||
		!strings.HasPrefix(got.Get("X-Forwarded-For"), "203.0.113.7, ") || got.Get(hdrForwardedBy) != "ctl-b" {
		t.Errorf("host %q, headers %v", gotHost, got)
	}
	// What a load balancer asks is about this machine.
	for _, p := range []string{"/ready", "/health", "/leader"} {
		if body := get(p, nil); body != "standby "+p {
			t.Errorf("%s answered by %q", p, body)
		}
	}
	// Forwarded once only.
	if body := get("/v1/models", map[string]string{hdrForwardedBy: "ctl-c"}); body != "standby /v1/models" {
		t.Errorf("a forwarded request was forwarded again: %q", body)
	}
}