Vellum/internal/auth/auth_test.go

package auth

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"strconv"
	"strings"
	"testing"
	"time"

	"github.com/mithral/vellum/internal/core"
	"github.com/mithral/vellum/migrations"
)

func testDB(t *testing.T) *sql.DB {
	t.Helper()
	db, err := core.OpenDB(filepath.Join(t.TempDir(), "auth.db"))
	if err != nil {
		t.Fatalf("OpenDB: %v", err)
	}
	t.Cleanup(func() { _ = db.Close() })

	log := slog.New(slog.NewTextHandler(io.Discard, nil))
	if err := core.Migrate(context.Background(), db, migrations.FS, log); err != nil {
		t.Fatalf("Migrate: %v", err)
	}
	return db
}

func mustCreateUser(t *testing.T, db *sql.DB, username string, admin bool) User {
	t.Helper()
	u, err := CreateUser(context.Background(), db, username, "", "password123", admin)
	if err != nil {
		t.Fatalf("CreateUser(%s): %v", username, err)
	}
	return u
}

func mustCreateLibrary(t *testing.T, db *sql.DB, name, kind, root string) int64 {
	t.Helper()
	var id int64
	err := db.QueryRowContext(context.Background(),
		`INSERT INTO libraries (name, kind, root_path) VALUES (?, ?, ?) RETURNING id`,
		name, kind, root).Scan(&id)
	if err != nil {
		t.Fatalf("insert library: %v", err)
	}
	return id
}

// ── users ──────────────────────────────────────────────────────────────

func TestCreateUserAndAuthenticate(t *testing.T) {
	db := testDB(t)
	ctx := context.Background()

	u, err := CreateUser(ctx, db, "alice", "Alice", "hunter2hunter2", false)
	if err != nil {
		t.Fatalf("CreateUser: %v", err)
	}
	if u.DisplayName != "Alice" {
		t.Errorf("DisplayName = %q", u.DisplayName)
	}

	got, err := Authenticate(ctx, db, "alice", "hunter2hunter2")
	if err != nil {
		t.Fatalf("Authenticate: %v", err)
	}
	if got.ID != u.ID {
		t.Errorf("authenticated id = %d, want %d", got.ID, u.ID)
	}

	// Usernames are case-insensitive.
	if _, err := Authenticate(ctx, db, "ALICE", "hunter2hunter2"); err != nil {
		t.Errorf("case-insensitive login failed: %v", err)
	}
}

func TestAuthenticateRejectsBadCredentials(t *testing.T) {
	db := testDB(t)
	ctx := context.Background()
	mustCreateUser(t, db, "bob", false)

	for _, tc := range []struct{ user, pass string }{
		{"bob", "wrong"},
		{"nobody", "password123"},
	} {
		_, err := Authenticate(ctx, db, tc.user, tc.pass)
		if !errors.Is(err, ErrInvalidCredentials) {
			t.Errorf("Authenticate(%q,%q) error = %v, want ErrInvalidCredentials", tc.user, tc.pass, err)
		}
	}
}

func TestCreateUserRejectsDuplicateAndBadUsername(t *testing.T) {
	db := testDB(t)
	ctx := context.Background()
	mustCreateUser(t, db, "carol", false)

	if _, err := CreateUser(ctx, db, "CAROL", "", "password123", false); !errors.Is(err, ErrUserExists) {
		t.Errorf("duplicate error = %v, want ErrUserExists", err)
	}
	if _, err := CreateUser(ctx, db, "x", "", "password123", false); !errors.Is(err, ErrInvalidUsername) {
		t.Errorf("short username error = %v, want ErrInvalidUsername", err)
	}
	if _, err := CreateUser(ctx, db, "has space", "", "password123", false); !errors.Is(err, ErrInvalidUsername) {
		t.Errorf("username with space error = %v, want ErrInvalidUsername", err)
	}
}

func TestSetPasswordRevokesSessions(t *testing.T) {
	db := testDB(t)
	ctx := context.Background()
	u := mustCreateUser(t, db, "dave", false)

	token, err := CreateSession(ctx, db, u.ID, time.Hour, "test", "127.0.0.1", "test")
	if err != nil {
		t.Fatalf("CreateSession: %v", err)
	}
	if _, err := SessionByToken(ctx, db, token); err != nil {
		t.Fatalf("session should exist: %v", err)
	}

	if err := SetPassword(ctx, db, u.ID, "brandnewpass"); err != nil {
		t.Fatalf("SetPassword: %v", err)
	}
	if _, err := SessionByToken(ctx, db, token); !errors.Is(err, ErrNoSession) {
		t.Errorf("session survived password change: %v", err)
	}
	if _, err := Authenticate(ctx, db, "dave", "brandnewpass"); err != nil {
		t.Errorf("new password does not work: %v", err)
	}
}

// ── sessions ───────────────────────────────────────────────────────────

func TestSessionExpiry(t *testing.T) {
	db := testDB(t)
	ctx := context.Background()
	u := mustCreateUser(t, db, "erin", false)

	token, err := CreateSession(ctx, db, u.ID, -time.Minute, "test", "", "")
	if err != nil {
		t.Fatalf("CreateSession: %v", err)
	}
	if _, err := SessionByToken(ctx, db, token); !errors.Is(err, ErrNoSession) {
		t.Errorf("expired session resolved: %v", err)
	}
}

func TestDeleteExpiredSessions(t *testing.T) {
	db := testDB(t)
	ctx := context.Background()
	u := mustCreateUser(t, db, "frank", false)

	if _, err := CreateSession(ctx, db, u.ID, -time.Hour, "", "", ""); err != nil {
		t.Fatal(err)
	}
	if _, err := CreateSession(ctx, db, u.ID, time.Hour, "", "", ""); err != nil {
		t.Fatal(err)
	}

	n, err := DeleteExpiredSessions(ctx, db)
	if err != nil {
		t.Fatalf("DeleteExpiredSessions: %v", err)
	}
	if n != 1 {
		t.Errorf("deleted %d sessions, want 1", n)
	}
}

func TestRawTokenIsNotStored(t *testing.T) {
	db := testDB(t)
	ctx := context.Background()
	u := mustCreateUser(t, db, "grace", false)

	token, err := CreateSession(ctx, db, u.ID, time.Hour, "", "", "")
	if err != nil {
		t.Fatal(err)
	}

	var count int
	if err := db.QueryRowContext(ctx,
		`SELECT count(*) FROM sessions WHERE token_hash = ?`, token).Scan(&count); err != nil {
		t.Fatal(err)
	}
	if count != 0 {
		t.Error("raw session token found in the database")
	}
}

// ── grants / visibility ────────────────────────────────────────────────

func TestLibraryVisibility(t *testing.T) {
	db := testDB(t)
	ctx := context.Background()

	comics := mustCreateLibrary(t, db, "Comics", "comics", "/media/comics")
	ebooks := mustCreateLibrary(t, db, "Ebooks", "ebooks", "/media/ebooks")

	kid := mustCreateUser(t, db, "kid", false)
	admin := mustCreateUser(t, db, "admin", true)

	if err := GrantLibrary(ctx, db, kid.ID, comics, true, false); err != nil {
		t.Fatalf("GrantLibrary: %v", err)
	}

	kidLibs, err := VisibleLibraryIDs(ctx, db, kid)
	if err != nil {
		t.Fatal(err)
	}
	if len(kidLibs) != 1 || kidLibs[0] != comics {
		t.Errorf("kid visible libraries = %v, want [%d]", kidLibs, comics)
	}

	if ok, err := CanAccessLibrary(ctx, db, kid, comics); err != nil || !ok {
		t.Errorf("kid should access comics: ok=%v err=%v", ok, err)
	}
	if ok, err := CanAccessLibrary(ctx, db, kid, ebooks); err != nil || ok {
		t.Errorf("kid should NOT access ebooks: ok=%v err=%v", ok, err)
	}

	adminLibs, err := VisibleLibraryIDs(ctx, db, admin)
	if err != nil {
		t.Fatal(err)
	}
	if len(adminLibs) != 2 {
		t.Errorf("admin visible libraries = %v, want both", adminLibs)
	}

	// Revoking removes access.
	if err := RevokeLibrary(ctx, db, kid.ID, comics); err != nil {
		t.Fatal(err)
	}
	if ok, _ := CanAccessLibrary(ctx, db, kid, comics); ok {
		t.Error("kid still has access after revocation")
	}
}

func TestGrantLibraryUpserts(t *testing.T) {
	db := testDB(t)
	ctx := context.Background()
	u := mustCreateUser(t, db, "henry", false)
	lib := mustCreateLibrary(t, db, "Music", "music", "/media/music")

	if err := GrantLibrary(ctx, db, u.ID, lib, true, false); err != nil {
		t.Fatal(err)
	}
	if err := GrantLibrary(ctx, db, u.ID, lib, false, true); err != nil {
		t.Fatal(err)
	}

	grants, err := LibraryGrants(ctx, db, u.ID)
	if err != nil {
		t.Fatal(err)
	}
	if len(grants) != 1 {
		t.Fatalf("grants = %v, want exactly 1", grants)
	}
	if grants[0].CanRead || !grants[0].CanDownload {
		t.Errorf("grant = %+v, want read=false download=true", grants[0])
	}
}

// ── HTTP flow ──────────────────────────────────────────────────────────

func newAuthServer(t *testing.T, db *sql.DB) http.Handler {
	t.Helper()
	log := slog.New(slog.NewTextHandler(io.Discard, nil))
	a := NewAuthenticator(db, "vellum_session", time.Hour, false, log)

	mux := http.NewServeMux()
	NewHandlers(a, log).Register(mux)

	return a.Middleware(mux)
}

func postJSON(t *testing.T, h http.Handler, path, body string, cookies []*http.Cookie) *httptest.ResponseRecorder {
	t.Helper()
	req := httptest.NewRequest(http.MethodPost, path, strings.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	for _, c := range cookies {
		req.AddCookie(c)
	}
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	return rec
}

func TestHTTPLoginMeLogout(t *testing.T) {
	db := testDB(t)
	h := newAuthServer(t, db)
	mustCreateUser(t, db, "irene", false)

	// Unauthenticated /me.
	req := httptest.NewRequest(http.MethodGet, "/api/v1/auth/me", nil)
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	if rec.Code != http.StatusUnauthorized {
		t.Fatalf("unauthenticated /me status = %d, want 401", rec.Code)
	}

	// Wrong password.
	rec = postJSON(t, h, "/api/v1/auth/login", `{"username":"irene","password":"nope"}`, nil)
	if rec.Code != http.StatusUnauthorized {
		t.Fatalf("bad login status = %d, want 401", rec.Code)
	}

	// Correct password.
	rec = postJSON(t, h, "/api/v1/auth/login", `{"username":"irene","password":"password123"}`, nil)
	if rec.Code != http.StatusOK {
		t.Fatalf("login status = %d, want 200; body=%s", rec.Code, rec.Body.String())
	}
	cookies := rec.Result().Cookies()
	if len(cookies) == 0 {
		t.Fatal("login did not set a cookie")
	}
	var sessCookie *http.Cookie
	for _, c := range cookies {
		if c.Name == "vellum_session" {
			sessCookie = c
		}
	}
	if sessCookie == nil {
		t.Fatal("no vellum_session cookie set")
	}
	if !sessCookie.HttpOnly {
		t.Error("session cookie is not HttpOnly")
	}

	// Authenticated /me.
	req = httptest.NewRequest(http.MethodGet, "/api/v1/auth/me", nil)
	req.AddCookie(sessCookie)
	rec = httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("/me status = %d, want 200", rec.Code)
	}
	var me struct {
		User struct {
			Username string `json:"username"`
		} `json:"user"`
		VisibleLibraries []int64 `json:"visible_libraries"`
	}
	if err := json.Unmarshal(rec.Body.Bytes(), &me); err != nil {
		t.Fatalf("decode /me: %v", err)
	}
	if me.User.Username != "irene" {
		t.Errorf("/me username = %q", me.User.Username)
	}

	// Logout clears the session server-side.
	rec = postJSON(t, h, "/api/v1/auth/logout", "", []*http.Cookie{sessCookie})
	if rec.Code != http.StatusOK {
		t.Fatalf("logout status = %d", rec.Code)
	}
	req = httptest.NewRequest(http.MethodGet, "/api/v1/auth/me", nil)
	req.AddCookie(sessCookie)
	rec = httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	if rec.Code != http.StatusUnauthorized {
		t.Errorf("session still valid after logout: %d", rec.Code)
	}
}

func TestBearerTokenAlsoWorks(t *testing.T) {
	db := testDB(t)
	h := newAuthServer(t, db)
	u := mustCreateUser(t, db, "jill", false)

	token, err := CreateSession(context.Background(), db, u.ID, time.Hour, "", "", "")
	if err != nil {
		t.Fatal(err)
	}

	req := httptest.NewRequest(http.MethodGet, "/api/v1/auth/me", nil)
	req.Header.Set("Authorization", "Bearer "+token)
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("bearer /me status = %d, want 200", rec.Code)
	}
}

// sessionTouchCounter installs a trigger that counts every UPDATE to
// sessions.last_seen and returns a reader for the running total. Counting in
// SQLite rather than by diffing timestamps is exact even when two touches land
// in the same millisecond, which is common in a tight request loop.
func sessionTouchCounter(t *testing.T, db *sql.DB) func() int {
	t.Helper()
	ctx := context.Background()
	if _, err := db.ExecContext(ctx, `CREATE TABLE session_touch_count (n INTEGER NOT NULL)`); err != nil {
		t.Fatalf("create touch counter: %v", err)
	}
	if _, err := db.ExecContext(ctx, `INSERT INTO session_touch_count (n) VALUES (0)`); err != nil {
		t.Fatalf("seed touch counter: %v", err)
	}
	if _, err := db.ExecContext(ctx,
		`CREATE TRIGGER count_session_touch AFTER UPDATE OF last_seen ON sessions
		 BEGIN UPDATE session_touch_count SET n = n + 1; END`); err != nil {
		t.Fatalf("create touch trigger: %v", err)
	}
	return func() int {
		var n int
		if err := db.QueryRowContext(ctx, `SELECT n FROM session_touch_count`).Scan(&n); err != nil {
			t.Fatalf("read touch counter: %v", err)
		}
		return n
	}
}

func TestMiddlewareThrottlesSessionTouch(t *testing.T) {
	db := testDB(t)
	ctx := context.Background()
	log := slog.New(slog.NewTextHandler(io.Discard, nil))
	a := NewAuthenticator(db, "vellum_session", time.Hour, false, log)

	u := mustCreateUser(t, db, "kyle", false)
	token, err := CreateSession(ctx, db, u.ID, time.Hour, "", "", "")
	if err != nil {
		t.Fatalf("CreateSession: %v", err)
	}
	sess, err := SessionByToken(ctx, db, token)
	if err != nil {
		t.Fatalf("SessionByToken: %v", err)
	}

	touches := sessionTouchCounter(t, db)

	// A trivial terminal handler: the point is only that the request passes
	// through Middleware, which is where the touch decision lives.
	h := a.Middleware(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusOK)
	}))
	burst := func() {
		t.Helper()
		for i := 0; i < 10; i++ {
			req := httptest.NewRequest(http.MethodGet, "/api/v1/media/track", nil)
			req.Header.Set("Authorization", "Bearer "+token)
			rec := httptest.NewRecorder()
			h.ServeHTTP(rec, req)
			if rec.Code != http.StatusOK {
				t.Fatalf("burst request %d status = %d, want 200", i, rec.Code)
			}
		}
	}

	// A freshly created session has a recent last_seen, so a 10-request burst
	// must not take the write lock at all.
	burst()
	if got := touches(); got != 0 {
		t.Errorf("recent last_seen: %d touches after a 10-request burst, want 0", got)
	}

	// Backdate last_seen past touchEvery, then reset the counter. The first
	// request of the next burst sees a stale value and must touch; every later
	// request then sees a fresh one and must be throttled — so the whole burst
	// writes exactly once, not ten times.
	staleTime := time.Now().Add(-10 * time.Minute).UTC()
	if _, err := db.ExecContext(ctx, `UPDATE sessions SET last_seen = ? WHERE id = ?`,
		core.Stamp(staleTime), sess.ID); err != nil {
		t.Fatalf("backdate last_seen: %v", err)
	}
	if _, err := db.ExecContext(ctx, `UPDATE session_touch_count SET n = 0`); err != nil {
		t.Fatalf("reset touch counter: %v", err)
	}

	burst()
	if got := touches(); got != 1 {
		t.Errorf("stale last_seen: %d touches after a 10-request burst, want exactly 1", got)
	}

	// The single touch must have advanced last_seen past the stale value.
	after, err := SessionByToken(ctx, db, token)
	if err != nil {
		t.Fatalf("SessionByToken after burst: %v", err)
	}
	if !after.LastSeen.After(staleTime) {
		t.Errorf("last_seen = %v, want after %v", after.LastSeen, staleTime)
	}
}

func TestAdminRoutesRequireAdmin(t *testing.T) {
	db := testDB(t)
	h := newAuthServer(t, db)
	user := mustCreateUser(t, db, "mallory", false)

	token, err := CreateSession(context.Background(), db, user.ID, time.Hour, "", "", "")
	if err != nil {
		t.Fatal(err)
	}

	req := httptest.NewRequest(http.MethodGet, "/api/v1/users", nil)
	req.Header.Set("Authorization", "Bearer "+token)
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	if rec.Code != http.StatusForbidden {
		t.Errorf("non-admin /users status = %d, want 403", rec.Code)
	}
}

func TestAdminCreatesUserAndGrants(t *testing.T) {
	db := testDB(t)
	h := newAuthServer(t, db)
	admin := mustCreateUser(t, db, "root", true)
	lib := mustCreateLibrary(t, db, "Games", "games", "/media/games")

	token, err := CreateSession(context.Background(), db, admin.ID, time.Hour, "", "", "")
	if err != nil {
		t.Fatal(err)
	}
	authHeader := "Bearer " + token

	// Create a user.
	req := httptest.NewRequest(http.MethodPost, "/api/v1/users",
		strings.NewReader(`{"username":"newbie","password":"password123"}`))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", authHeader)
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	if rec.Code != http.StatusCreated {
		t.Fatalf("create user status = %d, body=%s", rec.Code, rec.Body.String())
	}

	// Grant a library to user 2 (newbie is the second user).
	req = httptest.NewRequest(http.MethodPut, "/api/v1/users/2/grants/"+itoa(lib),
		strings.NewReader(`{"can_read":true,"can_download":true}`))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", authHeader)
	rec = httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("grant status = %d, body=%s", rec.Code, rec.Body.String())
	}

	grants, err := LibraryGrants(context.Background(), db, 2)
	if err != nil {
		t.Fatal(err)
	}
	if len(grants) != 1 || !grants[0].CanRead || !grants[0].CanDownload {
		t.Errorf("grants = %+v", grants)
	}
}

func itoa(n int64) string { return strconv.FormatInt(n, 10) }

// ── login rate limiting ─────────────────────────────────────────────────

// AllowLogin/ResetLogin are what every password-checking surface calls, so the
// limiter can sit on the Authenticator and one budget covers all of them.
func TestAllowLoginRefusesAfterTenAndResetClears(t *testing.T) {
	log := slog.New(slog.NewTextHandler(io.Discard, nil))
	a := NewAuthenticator(nil, "", time.Hour, false, log)

	for i := 1; i <= 10; i++ {
		if !a.AllowLogin("10.0.0.1", "alice") {
			t.Fatalf("attempt %d refused, want it under the limit", i)
		}
	}
	if a.AllowLogin("10.0.0.1", "alice") {
		t.Error("11th attempt was allowed")
	}

	// The window is per ip and username, so nobody else is locked out.
	if !a.AllowLogin("10.0.0.2", "alice") {
		t.Error("another ip was refused")
	}
	if !a.AllowLogin("10.0.0.1", "bob") {
		t.Error("another username was refused")
	}

	a.ResetLogin("10.0.0.1", "alice")
	if !a.AllowLogin("10.0.0.1", "alice") {
		t.Error("attempt refused after a successful login cleared the window")
	}
}