llmgw/internal/cli/leaderforward_test.go

package cli

import (
	"encoding/json"
	"io"
	"log/slog"
	"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"])
	}
	// What is not the API stays on the machine: the dashboard, pprof.
	if _, got := c.ask(t, "/debug/pprof/"); got["path"] != "local" {
		t.Errorf("a non-API path was forwarded: %v", 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)
	}
}