llmgw/test/e2e/auth_test.go

package e2e

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"strconv"
	"strings"
	"testing"
	"time"

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

// provision builds a store with a tenancy tree and returns its path plus a
// function that mints keys in it.
func provision(t *testing.T) (string, func(t *testing.T, name string, tune ...func(*store.Key)) string) {
	t.Helper()
	path := filepath.Join(t.TempDir(), "auth.db")
	ctx := context.Background()

	s, err := store.OpenSQLite(ctx, path)
	if err != nil {
		t.Fatal(err)
	}
	for _, err := range []error{
		s.PutOrg(ctx, store.Org{ID: "org", Name: "Org"}),
		s.PutTeam(ctx, store.Team{ID: "team", OrgID: "org", Name: "Team"}),
		s.PutProject(ctx, store.Project{ID: "proj", TeamID: "team", Name: "Proj"}),
	} {
		if err != nil {
			t.Fatal(err)
		}
	}
	_ = s.Close()

	mint := func(t *testing.T, name string, tune ...func(*store.Key)) string {
		t.Helper()
		s, err := store.OpenSQLite(context.Background(), path)
		if err != nil {
			t.Fatal(err)
		}
		defer s.Close()

		secret, k, err := store.NewKey("key-"+name, name, "proj")
		if err != nil {
			t.Fatal(err)
		}
		k.CreatedAt = time.Now()
		for _, f := range tune {
			f(&k)
		}
		if err := s.CreateKey(context.Background(), k); err != nil {
			t.Fatal(err)
		}
		return secret
	}
	return path, mint
}

// withStore sets limits on an existing store.
func withStore(t *testing.T, path string, fn func(*store.SQLite)) {
	t.Helper()
	s, err := store.OpenSQLite(context.Background(), path)
	if err != nil {
		t.Fatal(err)
	}
	defer s.Close()
	fn(s)
}

func (h *harness) call(t *testing.T, key, model string) (int, string) {
	t.Helper()
	body, _ := json.Marshal(map[string]any{
		"model":      model,
		"messages":   []map[string]string{{"role": "user", "content": "hi"}},
		"max_tokens": 4,
	})
	req, err := http.NewRequestWithContext(context.Background(), http.MethodPost,
		h.router.URL+"/v1/chat/completions", bytes.NewReader(body))
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Set("Content-Type", "application/json")
	if key != "" {
		req.Header.Set("Authorization", "Bearer "+key)
	}
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()

	var e struct {
		Error struct{ Message string } `json:"error"`
	}
	_ = json.NewDecoder(resp.Body).Decode(&e)
	return resp.StatusCode, e.Error.Message
}

func TestTheGatewayIsClosedWithoutAKey(t *testing.T) {
	// The whole point of phase 5's first half: before this, anything that could
	// reach the port could use every model in the catalog.
	path, mint := provision(t)
	key := mint(t, "good")

	h := start(t, 2, func(c *config.Server) {
		c.Auth = config.Auth{Enabled: true, Store: path, CacheTTL: 50 * time.Millisecond}
	})

	if code, _ := h.call(t, "", testModel); code != http.StatusUnauthorized {
		t.Errorf("no key got %d", code)
	}
	if code, _ := h.call(t, "sk-llmgw-"+"0123456789012345678901234567890123456789012", testModel); code != http.StatusUnauthorized {
		t.Errorf("an unknown key got %d", code)
	}
	if code, msg := h.call(t, key, testModel); code != http.StatusOK {
		t.Errorf("a valid key got %d: %s", code, msg)
	}
}

func TestRevocationTakesEffectWithinTheCacheTTL(t *testing.T) {
	// The cache is what keeps the database off the request path, so the price
	// is that a revocation is not instant. It has to be bounded and it has to
	// actually happen.
	path, mint := provision(t)
	key := mint(t, "doomed")

	h := start(t, 2, func(c *config.Server) {
		c.Auth = config.Auth{Enabled: true, Store: path, CacheTTL: 100 * time.Millisecond}
	})
	if code, _ := h.call(t, key, testModel); code != http.StatusOK {
		t.Fatal("the key did not work to begin with")
	}

	withStore(t, path, func(s *store.SQLite) {
		if err := s.RevokeKey(context.Background(), "key-doomed", time.Now()); err != nil {
			t.Fatal(err)
		}
	})

	deadline := time.Now().Add(3 * time.Second)
	for time.Now().Before(deadline) {
		if code, _ := h.call(t, key, testModel); code == http.StatusUnauthorized {
			return
		}
		time.Sleep(20 * time.Millisecond)
	}
	t.Fatal("a revoked key still works well past the cache TTL")
}

func TestAnAllowlistCoversEveryNameOfAModel(t *testing.T) {
	// The decision is per model, not per string. Matching only what the client
	// sent would make every alias a way around the allowlist; matching only the
	// canonical name would silently ignore an allowlist written in the names
	// clients actually use. Both directions have to work.
	path, mint := provision(t)
	byCanonical := mint(t, "canonical")
	byAlias := mint(t, "alias")
	other := mint(t, "other")

	withStore(t, path, func(s *store.SQLite) {
		for _, l := range []store.Limits{
			{Scope: store.ScopeKey, ScopeID: "key-canonical", AllowedModels: []string{testModel}},
			{Scope: store.ScopeKey, ScopeID: "key-alias", AllowedModels: []string{"alias-for-test-model"}},
			{Scope: store.ScopeKey, ScopeID: "key-other", AllowedModels: []string{"something-else"}},
		} {
			if err := s.PutLimits(context.Background(), l); err != nil {
				t.Fatal(err)
			}
		}
	})

	h := start(t, 2, func(c *config.Server) {
		c.Auth = config.Auth{Enabled: true, Store: path, CacheTTL: 50 * time.Millisecond}
		c.Models = []catalog.Model{{Name: testModel, Aliases: []string{"alias-for-test-model"}}}
	})

	cases := []struct {
		name, key, model string
		want             int
	}{
		{"canonical allowlist, canonical request", byCanonical, testModel, http.StatusOK},
		{"canonical allowlist, alias request", byCanonical, "alias-for-test-model", http.StatusOK},
		{"alias allowlist, alias request", byAlias, "alias-for-test-model", http.StatusOK},
		{"alias allowlist, canonical request", byAlias, testModel, http.StatusOK},
		{"unrelated allowlist, canonical request", other, testModel, http.StatusForbidden},
		{"unrelated allowlist, alias request", other, "alias-for-test-model", http.StatusForbidden},
	}
	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			if code, msg := h.call(t, c.key, c.model); code != c.want {
				t.Errorf("got %d (%s), want %d", code, msg, c.want)
			}
		})
	}
}

func TestTheKeyDecidesTheBandNotTheClient(t *testing.T) {
	// Until authentication there was nothing to consult, so the band came from
	// a header and any client could promote itself into the top priority band.
	path, mint := provision(t)
	bronze := mint(t, "bronze", func(k *store.Key) { k.Band = "bronze" })

	h := start(t, 2, func(c *config.Server) {
		c.Auth = config.Auth{Enabled: true, Store: path, CacheTTL: 50 * time.Millisecond}
		c.Router.Admission.Enabled = true
		c.Router.Admission.BandHeader = "x-llmgw-band"
		c.Router.Admission.DefaultBand = "bronze"
		c.Router.Admission.Bands = []config.Band{
			{Name: "platinum", Priority: 400, SaturationCeiling: 2},
			{Name: "bronze", Priority: 100, SaturationCeiling: 0.3},
		}
	})

	body, _ := json.Marshal(map[string]any{
		"model": testModel, "max_tokens": 4,
		"messages": []map[string]string{{"role": "user", "content": "hi"}},
	})
	req, _ := http.NewRequestWithContext(context.Background(), http.MethodPost,
		h.router.URL+"/v1/chat/completions", bytes.NewReader(body))
	req.Header.Set("Authorization", "Bearer "+bronze)
	// The client claims the top band. The key says otherwise.
	req.Header.Set("x-llmgw-band", "platinum")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("status %d", resp.StatusCode)
	}

	// The band a request was admitted in shows up in the admission metrics.
	if got := admitted(t, h, "platinum"); got != 0 {
		t.Errorf("a bronze key was admitted into platinum %d times", got)
	}
	if got := admitted(t, h, "bronze"); got == 0 {
		t.Error("the request was not admitted in the key's own band")
	}
}

// admitted counts requests the router admitted in a band, read from its own
// metrics rather than inferred.
func admitted(t *testing.T, h *harness, band string) int {
	t.Helper()
	mux, err := h.rt.AdminHandler()
	if err != nil {
		t.Fatal(err)
	}
	srv := httptest.NewServer(mux)
	defer srv.Close()

	req, _ := http.NewRequestWithContext(context.Background(), http.MethodGet, srv.URL+"/metrics", nil)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	body, _ := io.ReadAll(resp.Body)

	want := fmt.Sprintf("llmgw_admission_wait_seconds_count{band=\"%s\"", band)
	for _, line := range strings.Split(string(body), "\n") {
		if strings.HasPrefix(line, want) {
			fields := strings.Fields(line)
			n, err := strconv.Atoi(fields[len(fields)-1])
			if err != nil {
				t.Fatalf("parsing %q: %v", line, err)
			}
			return n
		}
	}
	return 0
}

func TestHealthStaysOpen(t *testing.T) {
	// A load balancer has no key, and a health check that needs one is a health
	// check that reports the gateway down.
	path, _ := provision(t)
	h := start(t, 1, func(c *config.Server) {
		c.Auth = config.Auth{Enabled: true, Store: path}
	})

	req, _ := http.NewRequestWithContext(context.Background(), http.MethodGet, h.router.URL+"/health", nil)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		t.Errorf("health needs a key: %d", resp.StatusCode)
	}
}

func TestWithoutAuthTheGatewayStillServes(t *testing.T) {
	// Turning authentication on is a deployment choice. A gateway on a trusted
	// network with no keys issued has to keep working.
	h := start(t, 2)
	if code, msg := h.call(t, "", testModel); code != http.StatusOK {
		t.Errorf("an unauthenticated gateway refused a request: %d %s", code, msg)
	}
}