Maven/cmd/mavgpud/auth_test.go

package main

import (
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"
)

// ok is what the boundary is protecting: anything that reaches it has spent
// the card.
func ok(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusTeapot) }

func TestRequireTokenRefusesEveryWrongCredential(t *testing.T) {
	h := requireToken("s3cret", 1<<20, http.HandlerFunc(ok))
	cases := []struct {
		name string
		auth string
		want int
	}{
		{"no header", "", http.StatusUnauthorized},
		{"wrong token", "Bearer wrong", http.StatusUnauthorized},
		{"prefix of the token", "Bearer s3cre", http.StatusUnauthorized},
		{"token with no scheme", "s3cret", http.StatusUnauthorized},
		{"basic auth", "Basic czNjcmV0", http.StatusUnauthorized},
		{"right token", "Bearer s3cret", http.StatusTeapot},
		{"scheme is case-insensitive", "bearer s3cret", http.StatusTeapot},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			r := httptest.NewRequest(http.MethodGet, "/health", nil)
			if tc.auth != "" {
				r.Header.Set("Authorization", tc.auth)
			}
			w := httptest.NewRecorder()
			h.ServeHTTP(w, r)
			if w.Code != tc.want {
				t.Errorf("status %d, want %d", w.Code, tc.want)
			}
		})
	}
}

// The 401 must not say which part was wrong. A probe that can tell a malformed
// header from a wrong token learns the header shape for free.
func TestUnauthorizedSaysNothingUseful(t *testing.T) {
	h := requireToken("s3cret", 1<<20, http.HandlerFunc(ok))
	w := httptest.NewRecorder()
	h.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/health", nil))
	if got := strings.TrimSpace(w.Body.String()); got != "unauthorized" {
		t.Errorf("body %q, want %q", got, "unauthorized")
	}
	if got := w.Header().Get("WWW-Authenticate"); got != "Bearer" {
		t.Errorf("WWW-Authenticate %q, want Bearer", got)
	}
}

// The body cap applies to an authenticated request. An unauthenticated one
// never gets far enough to allocate anything.
func TestRequireTokenCapsTheBody(t *testing.T) {
	var read error
	h := requireToken("t", 8, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		buf := make([]byte, 64)
		for read == nil {
			if _, read = r.Body.Read(buf); read != nil {
				break
			}
		}
	}))
	r := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(strings.Repeat("x", 4096)))
	r.Header.Set("Authorization", "Bearer t")
	h.ServeHTTP(httptest.NewRecorder(), r)
	if read == nil || !strings.Contains(read.Error(), "too large") {
		t.Errorf("read error %v, want the body cap", read)
	}
}

func TestAllowlistRefusesWhatMavenDoesNotCall(t *testing.T) {
	h := allowlist(http.HandlerFunc(ok))
	cases := []struct {
		method, path string
		want         int
	}{
		{http.MethodPost, "/v1/chat/completions", http.StatusTeapot},
		{http.MethodGet, "/v1/models", http.StatusTeapot},
		// /slots returns the prompts of whoever else is using the card, and
		// its actions write files the request names.
		{http.MethodGet, "/slots", http.StatusNotFound},
		{http.MethodPost, "/slots/0?action=save", http.StatusNotFound},
		{http.MethodGet, "/", http.StatusNotFound},
		{http.MethodGet, "/v1/chat/completions", http.StatusMethodNotAllowed},
	}
	for _, tc := range cases {
		t.Run(tc.method+" "+tc.path, func(t *testing.T) {
			w := httptest.NewRecorder()
			h.ServeHTTP(w, httptest.NewRequest(tc.method, tc.path, nil))
			if w.Code != tc.want {
				t.Errorf("status %d, want %d", w.Code, tc.want)
			}
		})
	}
}

func TestLimitInflightCapsConcurrency(t *testing.T) {
	const cap = 2
	var mu sync.Mutex
	now, peak := 0, 0
	release := make(chan struct{})
	h := limitInflight(cap, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		mu.Lock()
		now++
		if now > peak {
			peak = now
		}
		mu.Unlock()
		<-release
		mu.Lock()
		now--
		mu.Unlock()
	}))

	var wg sync.WaitGroup
	for i := 0; i < 8; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			h.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil))
		}()
	}
	// Let the first wave arrive, then drain. The assertion is the peak, and a
	// peak that never reached the cap still cannot exceed it.
	close(release)
	wg.Wait()
	if peak > cap {
		t.Errorf("%d requests in flight at once, cap is %d", peak, cap)
	}
}

func TestLoopbackListen(t *testing.T) {
	cases := map[string]bool{
		":8080":              false, // the shipped default, and the whole LAN
		"0.0.0.0:8080":       false,
		"[::]:8080":          false,
		"192.168.1.105:8080": false,
		"127.0.0.1:8080":     true,
		"[::1]:8080":         true,
		"localhost:8080":     true,
	}
	for addr, want := range cases {
		if got := loopbackListen(addr); got != want {
			t.Errorf("loopbackListen(%q) = %v, want %v", addr, got, want)
		}
	}
}

func TestReadToken(t *testing.T) {
	dir := t.TempDir()
	good := filepath.Join(dir, "tok")
	if err := os.WriteFile(good, []byte("  abc123\n"), 0o600); err != nil {
		t.Fatal(err)
	}
	got, err := readToken(good)
	if err != nil || got != "abc123" {
		t.Errorf("readToken = %q, %v; want abc123", got, err)
	}

	blank := filepath.Join(dir, "blank")
	if err := os.WriteFile(blank, []byte("\n\n"), 0o600); err != nil {
		t.Fatal(err)
	}
	if _, err := readToken(blank); err == nil {
		t.Error("an empty token file is not a token")
	}
	if _, err := readToken(filepath.Join(dir, "absent")); err == nil {
		t.Error("a missing token file is not a token")
	}
}