llmgw/test/e2e/signup_test.go

package e2e

import (
	"bytes"
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/cookiejar"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"
	"time"

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

// browser is a client that keeps cookies and does not follow redirects, so a
// test can see them.
type browser struct {
	t    *testing.T
	base string
	c    *http.Client
}

func newBrowser(t *testing.T, base string) *browser {
	t.Helper()
	jar, err := cookiejar.New(nil)
	if err != nil {
		t.Fatal(err)
	}
	return &browser{t: t, base: base, c: &http.Client{
		Jar: jar,
		CheckRedirect: func(*http.Request, []*http.Request) error {
			return http.ErrUseLastResponse
		},
	}}
}

func (b *browser) get(path string) (int, string) {
	b.t.Helper()
	req, _ := http.NewRequestWithContext(context.Background(), http.MethodGet, b.base+path, nil)
	resp, err := b.c.Do(req)
	if err != nil {
		b.t.Fatal(err)
	}
	defer resp.Body.Close()
	body, _ := io.ReadAll(resp.Body)
	return resp.StatusCode, string(body)
}

func (b *browser) post(path string, form url.Values) (int, string, string) {
	b.t.Helper()
	// The sign-in form carries a token its page set, as a browser's would.
	if path == "/login" && form.Get("csrf") == "" {
		_, page := b.get("/login")
		form = cloneValues(form)
		form.Set("csrf", csrf(b.t, page))
	}
	req, _ := http.NewRequestWithContext(context.Background(), http.MethodPost,
		b.base+path, strings.NewReader(form.Encode()))
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	resp, err := b.c.Do(req)
	if err != nil {
		b.t.Fatal(err)
	}
	defer resp.Body.Close()
	body, _ := io.ReadAll(resp.Body)
	return resp.StatusCode, string(body), resp.Header.Get("Location")
}

// csrf pulls the token the server put in the page, from a form field or the
// meta tag the JSON-posting pages use.
func csrf(t *testing.T, body string) string {
	t.Helper()
	for _, marker := range []string{`name="csrf" value="`, `name="csrf" content="`} {
		if i := strings.Index(body, marker); i >= 0 {
			rest := body[i+len(marker):]
			return rest[:strings.IndexByte(rest, '"')]
		}
	}
	t.Fatal("no csrf token in the page")
	return ""
}

func signupHarness(t *testing.T) (*harness, string) {
	t.Helper()
	path, _ := provision(t)
	h := start(t, 2, func(c *config.Server) {
		c.Auth = config.Auth{
			Enabled: true, Store: path, SignupOpen: true,
			CacheTTL: 50 * time.Millisecond,
		}
	})
	return h, path
}

const password = "a long enough passphrase"

func TestAnAccountIsUselessUntilItIsApproved(t *testing.T) {
	// The whole reason a public signup form is safe: it produces something that
	// cannot do anything.
	h, path := signupHarness(t)
	b := newBrowser(t, h.router.URL)

	code, body, _ := b.post("/signup", url.Values{
		"email": {"ayse@example.com"}, "name": {"Ayse"}, "password": {password},
	})
	if code != http.StatusOK || !strings.Contains(body, "waiting for approval") {
		t.Fatalf("signup returned %d: %s", code, body)
	}

	code, body, _ = b.post("/login", url.Values{
		"email": {"ayse@example.com"}, "password": {password},
	})
	if code != http.StatusUnauthorized || !strings.Contains(body, "waiting for approval") {
		t.Fatalf("a pending account signed in: %d", code)
	}

	// The profile page is not reachable either.
	if code, _ := b.get("/profile"); code != http.StatusSeeOther {
		t.Errorf("an unauthenticated visit to /profile returned %d", code)
	}

	// Approve, and it all works.
	withStore(t, path, func(s *store.DB) {
		users, err := s.Users(context.Background(), store.UserPending)
		if err != nil || len(users) != 1 {
			t.Fatalf("pending users: %v %v", users, err)
		}
		if err := s.ApproveUser(context.Background(), users[0].ID, "proj", "test", time.Now()); err != nil {
			t.Fatal(err)
		}
	})

	code, _, loc := b.post("/login", url.Values{
		"email": {"ayse@example.com"}, "password": {password},
	})
	if code != http.StatusSeeOther || loc != "/profile" {
		t.Fatalf("sign-in after approval returned %d to %q", code, loc)
	}
}

func TestAUserIssuesTheirOwnKeyAndItWorks(t *testing.T) {
	h, path := signupHarness(t)
	b := newBrowser(t, h.router.URL)
	signUpAndApprove(t, b, path, "ayse@example.com")

	_, profile := b.get("/profile")
	token := csrf(t, profile)

	code, body, _ := b.post("/profile/keys", url.Values{
		"csrf": {token}, "name": {"laptop"},
	})
	if code != http.StatusOK {
		t.Fatalf("creating a key returned %d", code)
	}
	key := extractKey(t, body)

	// The key the person made for themselves is a real key.
	if code, msg := h.call(t, key, testModel); code != http.StatusOK {
		t.Fatalf("a self-issued key does not work: %d %s", code, msg)
	}

	// And it is listed afterwards, without the secret.
	_, profile = b.get("/profile")
	if !strings.Contains(profile, "laptop") {
		t.Error("the key is not listed")
	}
	if strings.Contains(profile, key) {
		t.Error("the profile page shows the key again; it is meant to be shown once")
	}
}

func TestAUserCannotRevokeSomeoneElsesKey(t *testing.T) {
	// The ownership check is in the statement rather than a prior read, so a
	// guessed id from another account changes nothing.
	h, path := signupHarness(t)

	ayse := newBrowser(t, h.router.URL)
	signUpAndApprove(t, ayse, path, "ayse@example.com")
	_, p1 := ayse.get("/profile")
	_, body, _ := ayse.post("/profile/keys", url.Values{"csrf": {csrf(t, p1)}, "name": {"ayse-key"}})
	ayseKey := extractKey(t, body)

	var ayseKeyID string
	withStore(t, path, func(s *store.DB) {
		users, _ := s.Users(context.Background(), store.UserActive)
		for _, u := range users {
			if u.Email == "ayse@example.com" {
				keys, _ := s.KeysForUser(context.Background(), u.ID)
				if len(keys) != 1 {
					t.Fatalf("ayse has %d keys", len(keys))
				}
				ayseKeyID = keys[0].ID
			}
		}
	})

	mehmet := newBrowser(t, h.router.URL)
	signUpAndApprove(t, mehmet, path, "mehmet@example.com")
	_, p2 := mehmet.get("/profile")

	if code, _, _ := mehmet.post("/profile/keys/"+ayseKeyID+"/revoke",
		url.Values{"csrf": {csrf(t, p2)}}); code != http.StatusSeeOther {
		t.Fatalf("status %d", code)
	}
	// It must still work.
	if code, msg := h.call(t, ayseKey, testModel); code != http.StatusOK {
		t.Errorf("someone else revoked this key: %d %s", code, msg)
	}
}

func TestAFormWithoutItsTokenIsRefused(t *testing.T) {
	// A cookie authenticates the browser, not the page that asked. Without this,
	// any site the person visits can post here and have the browser attach the
	// session for it.
	h, path := signupHarness(t)
	b := newBrowser(t, h.router.URL)
	signUpAndApprove(t, b, path, "ayse@example.com")

	for _, bad := range []url.Values{
		{"name": {"x"}},
		{"csrf": {""}, "name": {"x"}},
		{"csrf": {"forged"}, "name": {"x"}},
	} {
		if code, _, _ := b.post("/profile/keys", bad); code != http.StatusBadRequest {
			t.Errorf("a form with csrf=%q was accepted: %d", bad.Get("csrf"), code)
		}
	}
}

func TestSigningOutEndsTheSession(t *testing.T) {
	h, path := signupHarness(t)
	b := newBrowser(t, h.router.URL)
	signUpAndApprove(t, b, path, "ayse@example.com")

	_, profile := b.get("/profile")
	if _, _, _ = b.post("/logout", url.Values{"csrf": {csrf(t, profile)}}); false {
		t.Fatal("unreachable")
	}
	if code, _ := b.get("/profile"); code != http.StatusSeeOther {
		t.Errorf("the profile was still reachable after signing out: %d", code)
	}
}

func TestSuspensionClosesOpenSessions(t *testing.T) {
	h, path := signupHarness(t)
	b := newBrowser(t, h.router.URL)
	id := signUpAndApprove(t, b, path, "ayse@example.com")

	if code, _ := b.get("/profile"); code != http.StatusOK {
		t.Fatalf("profile before suspension: %d", code)
	}
	withStore(t, path, func(s *store.DB) {
		if err := s.SetUserState(context.Background(), id, store.UserSuspended); err != nil {
			t.Fatal(err)
		}
	})
	if code, _ := b.get("/profile"); code != http.StatusSeeOther {
		t.Errorf("a suspended account kept its session: %d", code)
	}
}

func TestSignupCanBeClosed(t *testing.T) {
	path, _ := provision(t)
	h := start(t, 2, func(c *config.Server) {
		c.Auth = config.Auth{Enabled: true, Store: path, SignupOpen: false}
	})
	b := newBrowser(t, h.router.URL)

	if code, body := b.get("/signup"); code == http.StatusOK && strings.Contains(body, "<form") {
		t.Errorf("the signup form is served although signup is closed: %d", code)
	}
	// The status the mux gives depends on which catch-all the path falls
	// through to, which is not the point. What matters is that no account
	// appears.
	if code, _, _ := b.post("/signup", url.Values{
		"email": {"a@example.com"}, "password": {password},
	}); code == http.StatusOK || code == http.StatusSeeOther {
		t.Errorf("signup accepted a post although it is closed: %d", code)
	}
	withStore(t, path, func(s *store.DB) {
		users, err := s.Users(context.Background(), "")
		if err != nil {
			t.Fatal(err)
		}
		if len(users) != 0 {
			t.Errorf("%d accounts exist although signup is closed", len(users))
		}
	})
}

// signUpAndApprove runs the operator's half and returns the new user's id.
func signUpAndApprove(t *testing.T, b *browser, path, email string) string {
	t.Helper()
	if code, _, _ := b.post("/signup", url.Values{
		"email": {email}, "name": {email}, "password": {password},
	}); code != http.StatusOK {
		t.Fatalf("signup returned %d", code)
	}

	var id string
	withStore(t, path, func(s *store.DB) {
		users, err := s.Users(context.Background(), store.UserPending)
		if err != nil {
			t.Fatal(err)
		}
		for _, u := range users {
			if u.Email == email {
				id = u.ID
			}
		}
		if id == "" {
			t.Fatalf("no pending account for %s", email)
		}
		if err := s.ApproveUser(context.Background(), id, "proj", "test", time.Now()); err != nil {
			t.Fatal(err)
		}
	})

	if code, _, _ := b.post("/login", url.Values{
		"email": {email}, "password": {password},
	}); code != http.StatusSeeOther {
		t.Fatalf("sign-in returned %d", code)
	}
	return id
}

func TestChangingAPasswordKeepsThisBrowserAndDropsTheOthers(t *testing.T) {
	// Both halves matter. Signing every session out is the security property;
	// keeping this one is what stops the person who just proved their password
	// from being thrown back to the sign-in page.
	h, path := signupHarness(t)

	laptop := newBrowser(t, h.router.URL)
	signUpAndApprove(t, laptop, path, "ayse@example.com")

	// A second browser, signed in as the same person.
	phone := newBrowser(t, h.router.URL)
	if code, _, _ := phone.post("/login", url.Values{
		"email": {"ayse@example.com"}, "password": {password},
	}); code != http.StatusSeeOther {
		t.Fatalf("second sign-in returned %d", code)
	}
	if code, _ := phone.get("/profile"); code != http.StatusOK {
		t.Fatalf("the second browser was not signed in: %d", code)
	}

	_, profile := laptop.get("/profile")
	const newPassword = "an entirely different passphrase"
	code, body, _ := laptop.post("/profile/password", url.Values{
		"csrf": {csrf(t, profile)}, "current": {password}, "new": {newPassword},
	})
	if code != http.StatusOK {
		t.Fatalf("changing the password returned %d", code)
	}
	if !strings.Contains(body, "password was changed") {
		t.Errorf("the page does not confirm the change")
	}

	// This browser is still signed in, and its forms still work — which they
	// only do if the new CSRF token reached the page that was just rendered.
	newToken := csrf(t, body)
	if code, _, _ := laptop.post("/profile/keys", url.Values{
		"csrf": {newToken}, "name": {"after-change"},
	}); code != http.StatusOK {
		t.Errorf("the form on the page after a password change was refused: %d", code)
	}

	// The other browser is not.
	if code, _ := phone.get("/profile"); code != http.StatusSeeOther {
		t.Errorf("the other browser kept its session: %d", code)
	}

	// And the old password is gone.
	fresh := newBrowser(t, h.router.URL)
	if code, _, _ := fresh.post("/login", url.Values{
		"email": {"ayse@example.com"}, "password": {password},
	}); code != http.StatusUnauthorized {
		t.Errorf("the old password still signs in: %d", code)
	}
	if code, _, _ := fresh.post("/login", url.Values{
		"email": {"ayse@example.com"}, "password": {newPassword},
	}); code != http.StatusSeeOther {
		t.Errorf("the new password does not sign in: %d", code)
	}
}

func TestAWrongCurrentPasswordChangesNothing(t *testing.T) {
	h, path := signupHarness(t)
	b := newBrowser(t, h.router.URL)
	signUpAndApprove(t, b, path, "ayse@example.com")

	_, profile := b.get("/profile")
	_, body, _ := b.post("/profile/password", url.Values{
		"csrf": {csrf(t, profile)}, "current": {"wrong"}, "new": {"something else entirely"},
	})
	if !strings.Contains(body, "not your current password") {
		t.Error("the page does not say why it was refused")
	}
	// Still signed in, still on the old password.
	if code, _ := b.get("/profile"); code != http.StatusOK {
		t.Errorf("a failed change signed the person out: %d", code)
	}
}

func TestTheClusterPageIsOnlyForAdmins(t *testing.T) {
	// Serving the dashboard over the web needs an actual answer to "who may see
	// this", and the answer must not be "anyone with an account".
	path, _ := provision(t)
	h := start(t, 2, func(c *config.Server) {
		c.Auth = config.Auth{Enabled: true, Store: path, SignupOpen: true, CacheTTL: 50 * time.Millisecond}
	})
	h.rt.SetClusterView(fakeCluster{}, nil)
	// Rebuild the handler now that there is a cluster to show.
	srv := httptest.NewServer(h.rt.Handler())
	defer srv.Close()

	// A stranger is sent to sign in, which reveals nothing: /login is public.
	anon := newBrowser(t, srv.URL)
	if code, _ := anon.get("/admin"); code != http.StatusSeeOther {
		t.Errorf("an anonymous visit returned %d", code)
	}

	// An ordinary account gets 404, not 403. 403 would confirm the page exists
	// and that admin accounts are a thing here.
	user := newBrowser(t, srv.URL)
	id := signUpAndApproveAt(t, user, srv.URL, path, "ayse@example.com")
	code, body := user.get("/admin")
	if code != http.StatusNotFound {
		t.Errorf("a non-admin got %d, want 404", code)
	}
	if strings.Contains(strings.ToLower(body), "forbidden") ||
		strings.Contains(strings.ToLower(body), "admin") {
		t.Error("the refusal admits the page exists")
	}
	// And the profile page does not advertise a door they cannot open.
	if _, profile := user.get("/profile"); strings.Contains(profile, `href="/admin"`) {
		t.Error("a non-admin is shown a link to the cluster")
	}

	withStore(t, path, func(s *store.DB) {
		if err := s.SetAdmin(context.Background(), id, true); err != nil {
			t.Fatal(err)
		}
	})

	code, body = user.get("/admin")
	if code != http.StatusOK {
		t.Fatalf("an admin got %d", code)
	}
	if !strings.Contains(body, "llmgw-node01") {
		t.Error("the page does not show the cluster")
	}
	if _, profile := user.get("/profile"); !strings.Contains(profile, `href="/admin"`) {
		t.Error("an admin is not shown the link")
	}
}

// signUpAndApproveAt is signUpAndApprove against an explicitly given base URL,
// for tests that rebuild the handler.
func signUpAndApproveAt(t *testing.T, b *browser, base, path, email string) string {
	t.Helper()
	_ = base
	return signUpAndApprove(t, b, path, email)
}

// fakeCluster is enough of a cluster for the page to render.
type fakeCluster struct{}

func (fakeCluster) Nodes() []adminapi.NodeView {
	return []adminapi.NodeView{{ID: "llmgw-node01", State: "ready", GPUTotal: 8, GPUFree: 8}}
}
func (fakeCluster) Models() []adminapi.ModelView           { return nil }
func (fakeCluster) Replicas(string) []adminapi.ReplicaView { return nil }
func (fakeCluster) Summary() adminapi.Summary              { return adminapi.Summary{} }
func (fakeCluster) Explain(string) adminapi.ExplainView    { return adminapi.ExplainView{} }

func TestTheChatPageTalksToTheModelWithoutTheBrowserHoldingAKey(t *testing.T) {
	// The point of the page: someone tries a model without setting up a client,
	// and the browser never sees a credential. The request still goes down the
	// ordinary path, so every limit that applies to an API key applies here.
	path, _ := provision(t)
	h := start(t, 2, func(c *config.Server) {
		c.Auth = config.Auth{Enabled: true, Store: path, SignupOpen: true, CacheTTL: 50 * time.Millisecond}
	})
	b := newBrowser(t, h.router.URL)
	signUpAndApprove(t, b, path, "ayse@example.com")

	code, page := b.get("/chat")
	if code != http.StatusOK {
		t.Fatalf("/chat returned %d", code)
	}
	if !strings.Contains(page, testModel) {
		t.Errorf("the model picker does not offer %s", testModel)
	}
	// Whatever else is on the page, a key is not.
	if strings.Contains(page, store.KeyPrefix) {
		t.Error("an API key was rendered into the chat page")
	}

	token := csrf(t, page)
	body, _ := json.Marshal(map[string]any{
		"model":    testModel,
		"messages": []map[string]string{{"role": "user", "content": "hello"}},
	})
	req, _ := http.NewRequestWithContext(context.Background(), http.MethodPost,
		h.router.URL+"/chat/send", bytes.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("X-CSRF-Token", token)
	resp, err := b.c.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		out, _ := io.ReadAll(resp.Body)
		t.Fatalf("chat send returned %d: %s", resp.StatusCode, out)
	}
	out, _ := io.ReadAll(resp.Body)
	if !strings.Contains(string(out), "data:") {
		t.Errorf("the reply did not stream: %.200s", out)
	}
}

func TestTheChatPageIsNotOpenToStrangers(t *testing.T) {
	path, _ := provision(t)
	h := start(t, 2, func(c *config.Server) {
		c.Auth = config.Auth{Enabled: true, Store: path, SignupOpen: true}
	})
	anon := newBrowser(t, h.router.URL)

	if code, _ := anon.get("/chat"); code != http.StatusSeeOther {
		t.Errorf("an anonymous visit to /chat returned %d", code)
	}
	req, _ := http.NewRequestWithContext(context.Background(), http.MethodPost,
		h.router.URL+"/chat/send", strings.NewReader(`{"model":"m","messages":[{}]}`))
	req.Header.Set("Content-Type", "application/json")
	resp, err := anon.c.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusUnauthorized {
		t.Errorf("an anonymous send returned %d", resp.StatusCode)
	}
}

func TestChatObeysTheModelAllowlist(t *testing.T) {
	// If the page could reach a model the key path refuses, it would be a hole
	// in the limits everything else enforces.
	path, _ := provision(t)
	h := start(t, 2, func(c *config.Server) {
		c.Auth = config.Auth{Enabled: true, Store: path, SignupOpen: true, CacheTTL: 50 * time.Millisecond}
	})
	b := newBrowser(t, h.router.URL)
	signUpAndApprove(t, b, path, "ayse@example.com")

	withStore(t, path, func(s *store.DB) {
		if err := s.PutLimits(context.Background(), store.Limits{
			Scope: store.ScopeProject, ScopeID: "proj",
			AllowedModels: []string{"something-else"},
		}); err != nil {
			t.Fatal(err)
		}
	})

	_, page := b.get("/chat")
	if strings.Contains(page, `value="`+testModel+`"`) {
		t.Error("the picker offers a model the allowlist forbids")
	}

	body, _ := json.Marshal(map[string]any{
		"model": testModel, "messages": []map[string]string{{"role": "user", "content": "hi"}},
	})
	req, _ := http.NewRequestWithContext(context.Background(), http.MethodPost,
		h.router.URL+"/chat/send", bytes.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("X-CSRF-Token", csrf(t, page))
	resp, err := b.c.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusForbidden {
		t.Errorf("asking for a forbidden model returned %d", resp.StatusCode)
	}
}

func TestChatSendNeedsItsToken(t *testing.T) {
	// A JSON post carries the token in a header, which a cross-site form cannot
	// set at all.
	path, _ := provision(t)
	h := start(t, 2, func(c *config.Server) {
		c.Auth = config.Auth{Enabled: true, Store: path, SignupOpen: true}
	})
	b := newBrowser(t, h.router.URL)
	signUpAndApprove(t, b, path, "ayse@example.com")

	for _, token := range []string{"", "forged"} {
		req, _ := http.NewRequestWithContext(context.Background(), http.MethodPost,
			h.router.URL+"/chat/send", strings.NewReader(`{"model":"m","messages":[{}]}`))
		req.Header.Set("Content-Type", "application/json")
		if token != "" {
			req.Header.Set("X-CSRF-Token", token)
		}
		resp, err := b.c.Do(req)
		if err != nil {
			t.Fatal(err)
		}
		resp.Body.Close()
		if resp.StatusCode != http.StatusBadRequest {
			t.Errorf("a send with token %q returned %d", token, resp.StatusCode)
		}
	}
}

func cloneValues(v url.Values) url.Values {
	out := url.Values{}
	for k, vs := range v {
		out[k] = append([]string{}, vs...)
	}
	return out
}