llmgw/internal/web/claim_test.go

package web

import (
	"context"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
	"time"

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

// fakeClaims is one pending claim.
type fakeClaims struct {
	secret   string
	consumed atomic.Bool
	peeks    atomic.Int64
}

func (f *fakeClaims) PeekClaim(context.Context, string) (store.Claim, error) {
	f.peeks.Add(1)
	if f.consumed.Load() {
		return store.Claim{}, store.ErrClaimed
	}
	return store.Claim{KeyID: "key-1", ExpiresAt: time.Now().Add(time.Hour)}, nil
}

func (f *fakeClaims) ConsumeClaim(context.Context, string) (string, error) {
	if f.consumed.Swap(true) {
		return "", store.ErrClaimed
	}
	return f.secret, nil
}

func do(t *testing.T, h http.Handler, method, target string) (int, string, http.Header) {
	t.Helper()
	// The handler reads the token from a path value, so it has to go through a
	// mux that sets one.
	mux := http.NewServeMux()
	mux.Handle("GET /claim/{token}", h)
	mux.Handle("POST /claim/{token}", h)

	r := httptest.NewRequest(method, target, nil)
	w := httptest.NewRecorder()
	mux.ServeHTTP(w, r)
	return w.Code, w.Body.String(), w.Header()
}

func TestOpeningTheLinkDoesNotBurnIt(t *testing.T) {
	// The reason GET and POST differ. Chat clients, mail scanners and link
	// previewers fetch a URL the moment it is posted; a GET that consumed the
	// claim would mean the first thing to unfurl the link destroys the key
	// before anyone reads it.
	f := &fakeClaims{secret: "sk-llmgw-thesecret"}
	h := Claim(f, "v1", nil)

	for range 5 {
		code, body, _ := do(t, h, http.MethodGet, "http://x/claim/tok")
		if code != http.StatusOK {
			t.Fatalf("status %d", code)
		}
		if strings.Contains(body, "thesecret") {
			t.Fatal("a GET revealed the key")
		}
	}
	if f.consumed.Load() {
		t.Fatal("a GET consumed the claim")
	}
}

func TestRevealingShowsTheKeyOnceAndThenTheLinkIsGone(t *testing.T) {
	f := &fakeClaims{secret: "sk-llmgw-thesecret"}
	h := Claim(f, "v1", nil)

	code, body, _ := do(t, h, http.MethodPost, "http://x/claim/tok")
	if code != http.StatusOK {
		t.Fatalf("status %d", code)
	}
	if !strings.Contains(body, "sk-llmgw-thesecret") {
		t.Fatal("the key was not shown")
	}

	code, body, _ = do(t, h, http.MethodPost, "http://x/claim/tok")
	if code != http.StatusGone {
		t.Errorf("a second reveal returned %d", code)
	}
	if strings.Contains(body, "thesecret") {
		t.Fatal("the key was shown twice")
	}
}

func TestTheKeyPageIsNotCached(t *testing.T) {
	// A key rendered into a page must not sit in a proxy, a browser cache or a
	// search index.
	f := &fakeClaims{secret: "sk-llmgw-thesecret"}
	_, _, hdr := do(t, Claim(f, "v1", nil), http.MethodPost, "http://x/claim/tok")

	if cc := hdr.Get("Cache-Control"); !strings.Contains(cc, "no-store") {
		t.Errorf("Cache-Control is %q", cc)
	}
	if hdr.Get("X-Robots-Tag") == "" {
		t.Error("the page is not marked noindex")
	}
	if hdr.Get("Referrer-Policy") != "no-referrer" {
		t.Error("the URL — which is the credential — would leak in a referrer")
	}
}

func TestAUsedLinkAndAnUnknownOneReadTheSame(t *testing.T) {
	// Distinguishing them confirms that a token was once valid.
	used := &fakeClaims{secret: "s"}
	used.consumed.Store(true)
	unknown := &fakeClaims{secret: "s"}
	unknown.consumed.Store(true)

	_, a, _ := do(t, Claim(used, "v1", nil), http.MethodGet, "http://x/claim/real")
	_, b, _ := do(t, Claim(unknown, "v1", nil), http.MethodGet, "http://x/claim/never")
	// Strip the timestamp, which differs by construction.
	if strip(a) != strip(b) {
		t.Error("a used link and an unknown one render differently")
	}
}

func strip(s string) string {
	if i := strings.Index(s, "<footer>"); i >= 0 {
		return s[:i]
	}
	return s
}

func TestTheExamplesUseTheAddressTheUserReached(t *testing.T) {
	trustTestProxy(t)
	f := &fakeClaims{secret: "sk-llmgw-thesecret"}
	mux := http.NewServeMux()
	mux.Handle("POST /claim/{token}", Claim(f, "v1", nil))

	r := httptest.NewRequest(http.MethodPost, "http://10.0.0.140:8080/claim/tok", nil)
	r.Header.Set("X-Forwarded-Proto", "https")
	r.Header.Set("X-Forwarded-Host", "llmgw.mynodes.xyz")
	w := httptest.NewRecorder()
	mux.ServeHTTP(w, r)

	body := w.Body.String()
	if !strings.Contains(body, "https://llmgw.mynodes.xyz/v1/chat/completions") {
		t.Error("the example names an address the user cannot use")
	}
}