llmgw/internal/auth/auth_test.go

package auth

import (
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"sync"
	"sync/atomic"
	"testing"
	"time"

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

// fixture builds a store with one org/team/project and one key.
func fixture(t *testing.T) (*store.DB, string, string) {
	t.Helper()
	ctx := context.Background()
	s, err := store.OpenSQLite(ctx, filepath.Join(t.TempDir(), "t.db"))
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = s.Close() })

	must(t, s.PutOrg(ctx, store.Org{ID: "org1", Name: "Acme"}))
	must(t, s.PutTeam(ctx, store.Team{ID: "team1", OrgID: "org1", Name: "R"}))
	must(t, s.PutProject(ctx, store.Project{ID: "proj1", TeamID: "team1", Name: "P"}))

	secret, k, err := store.NewKey("key1", "laptop", "proj1")
	if err != nil {
		t.Fatal(err)
	}
	k.CreatedAt = time.Now()
	k.Band = "gold"
	must(t, s.CreateKey(ctx, k))
	return s, secret, k.ID
}

func must(t *testing.T, err error) {
	t.Helper()
	if err != nil {
		t.Fatal(err)
	}
}

func TestAValidKeyResolvesItsWholeChain(t *testing.T) {
	s, secret, keyID := fixture(t)
	a := New(s, DefaultConfig())

	p, err := a.Authenticate(context.Background(), secret)
	if err != nil {
		t.Fatal(err)
	}
	if p.KeyID != keyID || p.ProjectID != "proj1" || p.TeamID != "team1" || p.OrgID != "org1" {
		t.Errorf("principal %+v", p)
	}
	if p.Band != "gold" {
		t.Errorf("band %q", p.Band)
	}
}

func TestRejections(t *testing.T) {
	s, secret, keyID := fixture(t)
	ctx := context.Background()
	a := New(s, DefaultConfig())

	if _, err := a.Authenticate(ctx, ""); !errors.Is(err, ErrMissing) {
		t.Errorf("no credential: %v", err)
	}
	// A bearer token from some other system must be rejected on shape alone.
	if _, err := a.Authenticate(ctx, "ghp_0123456789abcdefghij"); !errors.Is(err, ErrInvalid) {
		t.Errorf("foreign token: %v", err)
	}
	if _, err := a.Authenticate(ctx, alter(secret)); !errors.Is(err, ErrInvalid) {
		t.Errorf("wrong key: %v", err)
	}

	must(t, s.RevokeKey(ctx, keyID, time.Now()))
	a.Invalidate()
	if _, err := a.Authenticate(ctx, secret); !errors.Is(err, ErrRevoked) {
		t.Errorf("revoked key: %v", err)
	}
}

// countingStore counts how often the database is actually read.
type countingStore struct {
	store.Store
	lookups atomic.Int64
	fail    atomic.Bool
}

func (c *countingStore) KeyByHash(ctx context.Context, hash []byte) (store.Key, error) {
	c.lookups.Add(1)
	if c.fail.Load() {
		return store.Key{}, errors.New("database is down")
	}
	return c.Store.KeyByHash(ctx, hash)
}

// alter changes a key by one character, guaranteeing the result differs.
//
// Replacing the last character with a fixed one does not: keys are base62, so
// one in sixty-two already ends in that character and the "wrong" key is the
// right one. That test passed locally for a day and failed in CI, which is
// exactly how long a one-in-sixty-two flake takes to look like a real bug.
func alter(secret string) string {
	last := secret[len(secret)-1]
	replacement := byte('a')
	if last == replacement {
		replacement = 'b'
	}
	return secret[:len(secret)-1] + string(replacement)
}

func TestTheRequestPathDoesNotReadTheDatabase(t *testing.T) {
	// The rule the whole package is shaped around. Without the cache this is a
	// database read per request, on every request.
	s, secret, _ := fixture(t)
	c := &countingStore{Store: s}
	a := New(c, DefaultConfig())

	for range 100 {
		if _, err := a.Authenticate(context.Background(), secret); err != nil {
			t.Fatal(err)
		}
	}
	if n := c.lookups.Load(); n != 1 {
		t.Errorf("100 requests caused %d database reads", n)
	}
}

func TestARejectedKeyIsNotADatabaseReadPerRequest(t *testing.T) {
	// A client looping with a bad key would otherwise be a free way to load the
	// database from outside.
	s, _, _ := fixture(t)
	c := &countingStore{Store: s}
	a := New(c, DefaultConfig())

	bad := store.KeyPrefix + "0000000000000000000000000000000000000000000"
	for range 50 {
		if _, err := a.Authenticate(context.Background(), bad); !errors.Is(err, ErrInvalid) {
			t.Fatal(err)
		}
	}
	if n := c.lookups.Load(); n != 1 {
		t.Errorf("50 bad requests caused %d database reads", n)
	}
}

func TestCachedDecisionsExpire(t *testing.T) {
	// The cache bounds how long a revocation can take to take effect, so it has
	// to actually expire.
	s, secret, keyID := fixture(t)
	ctx := context.Background()

	now := time.Now()
	a := New(s, Config{TTL: time.Second, NegativeTTL: time.Second})
	a.now = func() time.Time { return now }

	if _, err := a.Authenticate(ctx, secret); err != nil {
		t.Fatal(err)
	}
	must(t, s.RevokeKey(ctx, keyID, now))

	// Still inside the TTL: the stale answer is the point of the cache.
	if _, err := a.Authenticate(ctx, secret); err != nil {
		t.Errorf("the cache expired early: %v", err)
	}

	now = now.Add(2 * time.Second)
	if _, err := a.Authenticate(ctx, secret); !errors.Is(err, ErrRevoked) {
		t.Errorf("a revoked key still works after the TTL: %v", err)
	}
}

func TestAStoreFailureIsNotAnAuthenticationAnswer(t *testing.T) {
	// Two things must not happen: the failure must not be cached as a
	// rejection, and it must not read to the caller as a bad key.
	s, secret, _ := fixture(t)
	c := &countingStore{Store: s}
	a := New(c, DefaultConfig())

	c.fail.Store(true)
	_, err := a.Authenticate(context.Background(), secret)
	if err == nil {
		t.Fatal("a broken database authenticated a request")
	}
	if isRejection(err) {
		t.Errorf("a database failure was reported as a credential problem: %v", err)
	}

	c.fail.Store(false)
	if _, err := a.Authenticate(context.Background(), secret); err != nil {
		t.Errorf("the failure was cached: %v", err)
	}
}

func TestExpiryIsNotCached(t *testing.T) {
	ctx := context.Background()
	s, _, _ := fixture(t)

	secret, k, _ := store.NewKey("key2", "temp", "proj1")
	k.CreatedAt = time.Now()
	k.ExpiresAt = time.Now().Add(time.Hour)
	must(t, s.CreateKey(ctx, k))

	now := time.Now()
	a := New(s, Config{TTL: time.Minute})
	a.now = func() time.Time { return now }

	if _, err := a.Authenticate(ctx, secret); err != nil {
		t.Fatal(err)
	}
	// Past the key's expiry but inside the cache TTL. The cached decision is
	// allowed to be stale by at most the TTL, which is the documented trade.
	now = now.Add(2 * time.Hour)
	if _, err := a.Authenticate(ctx, secret); !errors.Is(err, ErrExpired) {
		t.Errorf("an expired key authenticated: %v", err)
	}
}

func TestConcurrentAuthentication(t *testing.T) {
	s, secret, _ := fixture(t)
	a := New(s, DefaultConfig())

	var wg sync.WaitGroup
	errs := make(chan error, 64)
	for range 64 {
		wg.Add(1)
		go func() {
			defer wg.Done()
			if _, err := a.Authenticate(context.Background(), secret); err != nil {
				errs <- err
			}
		}()
	}
	wg.Wait()
	close(errs)
	for err := range errs {
		t.Fatal(err)
	}
}

func TestMiddlewareGuardsEveryRouteButTheExemptOnes(t *testing.T) {
	s, secret, _ := fixture(t)
	a := New(s, DefaultConfig())

	var reached atomic.Int64
	inner := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		reached.Add(1)
		if _, ok := FromContext(r.Context()); !ok && r.URL.Path != "/health" {
			t.Error("an authenticated request carried no principal")
		}
		w.WriteHeader(http.StatusOK)
	})
	srv := httptest.NewServer(Middleware(a, nil, "/health")(inner))
	defer srv.Close()

	get := func(path, token string) int {
		req, _ := http.NewRequest(http.MethodGet, srv.URL+path, nil) //nolint:noctx // test
		if token != "" {
			req.Header.Set("Authorization", "Bearer "+token)
		}
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			t.Fatal(err)
		}
		defer resp.Body.Close()
		return resp.StatusCode
	}

	if code := get("/health", ""); code != http.StatusOK {
		t.Errorf("health needs no key, got %d", code)
	}
	if code := get("/v1/models", ""); code != http.StatusUnauthorized {
		t.Errorf("an unauthenticated request got %d", code)
	}
	if code := get("/v1/models", "sk-llmgw-bogus"); code != http.StatusUnauthorized {
		t.Errorf("a bad key got %d", code)
	}
	if code := get("/v1/models", secret); code != http.StatusOK {
		t.Errorf("a good key got %d", code)
	}
	if n := reached.Load(); n != 2 {
		t.Errorf("the handler was reached %d times; only health and the valid key should pass", n)
	}
}

func TestARevokedKeyLooksExactlyLikeAnUnknownOne(t *testing.T) {
	// Telling a caller that a key exists but was revoked confirms the key was
	// real, which is information worth withholding.
	s, secret, keyID := fixture(t)
	must(t, s.RevokeKey(context.Background(), keyID, time.Now()))

	a := New(s, DefaultConfig())
	srv := httptest.NewServer(Middleware(a, nil)(http.HandlerFunc(
		func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusOK) })))
	defer srv.Close()

	body := func(token string) string {
		req, _ := http.NewRequest(http.MethodGet, srv.URL, nil) //nolint:noctx // test
		req.Header.Set("Authorization", "Bearer "+token)
		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 e.Error.Message
	}

	unknown := store.KeyPrefix + "0000000000000000000000000000000000000000000"
	if body(secret) != body(unknown) {
		t.Errorf("revoked says %q, unknown says %q", body(secret), body(unknown))
	}
}

func TestAStoreFailureDoesNotReadAsABadKey(t *testing.T) {
	s, secret, _ := fixture(t)
	c := &countingStore{Store: s}
	c.fail.Store(true)
	a := New(c, DefaultConfig())

	srv := httptest.NewServer(Middleware(a, nil)(http.HandlerFunc(
		func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusOK) })))
	defer srv.Close()

	req, _ := http.NewRequest(http.MethodGet, srv.URL, nil) //nolint:noctx // test
	req.Header.Set("Authorization", "Bearer "+secret)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	// 503 rather than 401: sending someone hunting for a key problem that does
	// not exist wastes the one thing an incident is short of.
	if resp.StatusCode != http.StatusServiceUnavailable {
		t.Errorf("status %d", resp.StatusCode)
	}
}

func TestBearerTokenAcceptsWhatClientsActuallySend(t *testing.T) {
	cases := map[string]string{
		"Bearer sk-llmgw-abc":  "sk-llmgw-abc",
		"sk-llmgw-abc":         "sk-llmgw-abc",
		"Bearer  sk-llmgw-abc": "sk-llmgw-abc",
		"":                     "",
	}
	for header, want := range cases {
		r := httptest.NewRequest(http.MethodGet, "/", nil)
		if header != "" {
			r.Header.Set("Authorization", header)
		}
		if got := BearerToken(r); got != want {
			t.Errorf("Authorization %q -> %q, want %q", header, got, want)
		}
	}

	// Some tools only manage to set this one.
	r := httptest.NewRequest(http.MethodGet, "/", nil)
	r.Header.Set("api-key", "sk-llmgw-xyz")
	if got := BearerToken(r); got != "sk-llmgw-xyz" {
		t.Errorf("api-key header -> %q", got)
	}
}

func TestLastUsedIsBatchedNotWrittenPerRequest(t *testing.T) {
	s, secret, keyID := fixture(t)
	ctx := context.Background()
	a := New(s, Config{TTL: time.Minute, TouchInterval: time.Hour})

	for range 20 {
		if _, err := a.Authenticate(ctx, secret); err != nil {
			t.Fatal(err)
		}
	}
	// Nothing written yet: the field is not read in real time.
	k, _ := s.KeyByHash(ctx, store.HashKey(secret))
	if !k.LastUsedAt.IsZero() {
		t.Error("last-used was written on the request path")
	}

	must(t, a.FlushTouches(ctx))
	k, _ = s.KeyByHash(ctx, store.HashKey(secret))
	if k.LastUsedAt.IsZero() {
		t.Error("last-used was never written")
	}
	if k.ID != keyID {
		t.Errorf("wrong key touched: %s", k.ID)
	}
}

func TestExemptingTheLandingPageDoesNotExemptEverything(t *testing.T) {
	// A near-miss worth a test. "/" ends in a slash, so a prefix rule would
	// match every path there is: authentication still installed, still logging,
	// and letting the entire API through.
	cases := []struct {
		path string
		want bool
	}{
		{"/", true},
		{"/_static/base.css", true},
		{"/health", true},
		{"/v1/models", false},
		{"/v1/chat/completions", false},
		{"/healthz", false},  // not a prefix rule
		{"/_staticx", false}, // prefix rules end at the slash
	}
	exempt := []string{"/health", "/", "/_static/"}
	for _, c := range cases {
		if got := isExempt(c.path, exempt); got != c.want {
			t.Errorf("isExempt(%q) = %v, want %v", c.path, got, c.want)
		}
	}
}

func TestAlterAlwaysChangesTheKey(t *testing.T) {
	// The helper the test above depends on. If it can return its input, the
	// test it serves silently stops testing anything.
	for range 500 {
		secret, _, err := store.NewKey("id", "n", "p")
		if err != nil {
			t.Fatal(err)
		}
		if alter(secret) == secret {
			t.Fatalf("alter returned its input for %q", secret)
		}
	}
	// Including the case that made the original flaky: a key ending in the
	// replacement character.
	if got := alter("sk-llmgw-aaaa"); got == "sk-llmgw-aaaa" {
		t.Error("alter is a no-op when the key already ends in the replacement")
	}
}