mcias/internal/ui/ui_test.go

package ui

import (
	"crypto/ed25519"
	"crypto/rand"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"
	"time"

	"git.wntrmute.dev/kyle/mcias/internal/config"
	"git.wntrmute.dev/kyle/mcias/internal/db"
	"git.wntrmute.dev/kyle/mcias/internal/model"
	"git.wntrmute.dev/kyle/mcias/internal/token"
)

const testIssuer = "https://auth.example.com"

// newTestUIServer creates a UIServer backed by an in-memory DB.
func newTestUIServer(t *testing.T) *UIServer {
	t.Helper()

	pub, priv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatalf("generate key: %v", err)
	}

	database, err := db.Open(":memory:")
	if err != nil {
		t.Fatalf("open db: %v", err)
	}
	if err := db.Migrate(database); err != nil {
		t.Fatalf("migrate db: %v", err)
	}
	t.Cleanup(func() { _ = database.Close() })

	masterKey := make([]byte, 32)
	if _, err := rand.Read(masterKey); err != nil {
		t.Fatalf("generate master key: %v", err)
	}

	cfg := config.NewTestConfig(testIssuer)
	logger := slog.New(slog.NewTextHandler(io.Discard, nil))

	uiSrv, err := New(database, cfg, priv, pub, masterKey, logger)
	if err != nil {
		t.Fatalf("new UIServer: %v", err)
	}
	return uiSrv
}

// newTestMux creates a UIServer and returns the http.Handler used in production
// (a ServeMux with all UI routes registered, wrapped with securityHeaders).
func newTestMux(t *testing.T) http.Handler {
	t.Helper()
	uiSrv := newTestUIServer(t)
	mux := http.NewServeMux()
	uiSrv.Register(mux)
	return mux
}

// assertSecurityHeaders verifies all mandatory defensive headers are present in
// resp with acceptable values. The label is used in failure messages to identify
// which endpoint the test was checking.
func assertSecurityHeaders(t *testing.T, h http.Header, label string) {
	t.Helper()

	checks := []struct {
		header  string
		wantSub string
	}{
		{"Content-Security-Policy", "default-src 'self'"},
		{"X-Content-Type-Options", "nosniff"},
		{"X-Frame-Options", "DENY"},
		{"Strict-Transport-Security", "max-age="},
		{"Referrer-Policy", "no-referrer"},
	}
	for _, c := range checks {
		val := h.Get(c.header)
		if val == "" {
			t.Errorf("[%s] missing security header %s", label, c.header)
			continue
		}
		if c.wantSub != "" && !strings.Contains(val, c.wantSub) {
			t.Errorf("[%s] %s = %q, want substring %q", label, c.header, val, c.wantSub)
		}
	}
}

// TestSecurityHeadersOnLoginPage verifies headers are present on the public login page.
func TestSecurityHeadersOnLoginPage(t *testing.T) {
	mux := newTestMux(t)

	req := httptest.NewRequest(http.MethodGet, "/login", nil)
	rr := httptest.NewRecorder()
	mux.ServeHTTP(rr, req)

	assertSecurityHeaders(t, rr.Result().Header, "GET /login")
}

// TestSecurityHeadersOnUnauthenticatedDashboard verifies headers are present even
// when the response is a redirect to login (no session cookie supplied).
func TestSecurityHeadersOnUnauthenticatedDashboard(t *testing.T) {
	mux := newTestMux(t)

	req := httptest.NewRequest(http.MethodGet, "/dashboard", nil)
	rr := httptest.NewRecorder()
	mux.ServeHTTP(rr, req)

	assertSecurityHeaders(t, rr.Result().Header, "GET /dashboard (no session)")
}

// TestSecurityHeadersOnRootRedirect verifies headers on the "/" → "/login" redirect.
func TestSecurityHeadersOnRootRedirect(t *testing.T) {
	mux := newTestMux(t)

	req := httptest.NewRequest(http.MethodGet, "/", nil)
	rr := httptest.NewRecorder()
	mux.ServeHTTP(rr, req)

	assertSecurityHeaders(t, rr.Result().Header, "GET /")
}

// TestSecurityHeadersOnStaticAsset verifies headers are present on static file responses.
func TestSecurityHeadersOnStaticAsset(t *testing.T) {
	mux := newTestMux(t)

	req := httptest.NewRequest(http.MethodGet, "/static/style.css", nil)
	rr := httptest.NewRecorder()
	mux.ServeHTTP(rr, req)

	// 200 or 404 — either way the securityHeaders wrapper must fire.
	assertSecurityHeaders(t, rr.Result().Header, "GET /static/style.css")
}

// TestCSPDirectives verifies the Content-Security-Policy includes same-origin
// directives for scripts and styles.
func TestCSPDirectives(t *testing.T) {
	mux := newTestMux(t)

	req := httptest.NewRequest(http.MethodGet, "/login", nil)
	rr := httptest.NewRecorder()
	mux.ServeHTTP(rr, req)

	csp := rr.Header().Get("Content-Security-Policy")
	for _, directive := range []string{
		"default-src 'self'",
		"script-src 'self'",
		"style-src 'self'",
	} {
		if !strings.Contains(csp, directive) {
			t.Errorf("CSP missing directive %q; full value: %q", directive, csp)
		}
	}
}

// TestHSTSMinAge verifies HSTS max-age is at least two years (63072000 seconds).
func TestHSTSMinAge(t *testing.T) {
	mux := newTestMux(t)

	req := httptest.NewRequest(http.MethodGet, "/login", nil)
	rr := httptest.NewRecorder()
	mux.ServeHTTP(rr, req)

	hsts := rr.Header().Get("Strict-Transport-Security")
	if !strings.Contains(hsts, "max-age=63072000") {
		t.Errorf("HSTS = %q, want max-age=63072000 (2 years)", hsts)
	}
}

// TestSecurityHeadersMiddlewareUnit tests the securityHeaders middleware in
// isolation, independent of routing, to guard against future refactoring.
func TestSecurityHeadersMiddlewareUnit(t *testing.T) {
	reached := false
	inner := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		reached = true
		w.WriteHeader(http.StatusOK)
	})

	handler := securityHeaders(inner)
	req := httptest.NewRequest(http.MethodGet, "/test", nil)
	rr := httptest.NewRecorder()
	handler.ServeHTTP(rr, req)

	if !reached {
		t.Error("inner handler was not reached")
	}
	assertSecurityHeaders(t, rr.Result().Header, "unit test")
}

// TestTOTPNonceIssuedAndConsumed verifies that issueTOTPNonce produces a
// non-empty nonce and consumeTOTPNonce returns the correct account ID exactly
// once (single-use).
func TestTOTPNonceIssuedAndConsumed(t *testing.T) {
	u := newTestUIServer(t)

	const accountID int64 = 42
	nonce, err := u.issueTOTPNonce(accountID)
	if err != nil {
		t.Fatalf("issueTOTPNonce: %v", err)
	}
	if nonce == "" {
		t.Fatal("expected non-empty nonce")
	}

	// First consumption must succeed.
	got, ok := u.consumeTOTPNonce(nonce)
	if !ok {
		t.Fatal("consumeTOTPNonce: expected ok=true on first use")
	}
	if got != accountID {
		t.Errorf("accountID = %d, want %d", got, accountID)
	}

	// Second consumption must fail (single-use).
	_, ok2 := u.consumeTOTPNonce(nonce)
	if ok2 {
		t.Error("consumeTOTPNonce: expected ok=false on second use (single-use guarantee violated)")
	}
}

// TestTOTPNonceUnknownRejected verifies that a never-issued nonce is rejected.
func TestTOTPNonceUnknownRejected(t *testing.T) {
	u := newTestUIServer(t)
	_, ok := u.consumeTOTPNonce("not-a-real-nonce")
	if ok {
		t.Error("consumeTOTPNonce: expected ok=false for unknown nonce")
	}
}

// TestTOTPNonceExpired verifies that an expired nonce is rejected even if
// the token exists in the map.
func TestTOTPNonceExpired(t *testing.T) {
	u := newTestUIServer(t)

	const accountID int64 = 99
	nonce, err := u.issueTOTPNonce(accountID)
	if err != nil {
		t.Fatalf("issueTOTPNonce: %v", err)
	}

	// Back-date the stored entry so it appears expired.
	v, loaded := u.pendingLogins.Load(nonce)
	if !loaded {
		t.Fatal("nonce not found in pendingLogins immediately after issuance")
	}
	pl, castOK := v.(*pendingLogin)
	if !castOK {
		t.Fatal("pendingLogins value is not *pendingLogin")
	}
	pl.expiresAt = time.Now().Add(-time.Second)

	_, ok := u.consumeTOTPNonce(nonce)
	if ok {
		t.Error("consumeTOTPNonce: expected ok=false for expired nonce")
	}
}

// TestLoginPostPasswordNotInTOTPForm verifies that after step 1, the TOTP
// step form body does not contain the user's password.
func TestLoginPostPasswordNotInTOTPForm(t *testing.T) {
	u := newTestUIServer(t)

	// Create an account with a known password and TOTP required flag.
	// We use the auth package to hash and the db to store directly.
	acct, err := u.db.CreateAccount("totpuser", model.AccountTypeHuman, "$argon2id$v=19$m=65536,t=3,p=4$dGVzdHNhbHQ$dGVzdGhhc2g")
	if err != nil {
		t.Fatalf("CreateAccount: %v", err)
	}
	// Enable TOTP required flag directly (use a stub secret so the account is
	// consistent; the step-1→step-2 nonce test only covers step 1 here).
	if err := u.db.StorePendingTOTP(acct.ID, []byte("enc"), []byte("nonce")); err != nil {
		t.Fatalf("StorePendingTOTP: %v", err)
	}
	if err := u.db.SetTOTP(acct.ID, []byte("enc"), []byte("nonce")); err != nil {
		t.Fatalf("SetTOTP: %v", err)
	}

	// POST step 1 with wrong password (will fail auth but verify form shape doesn't matter).
	// Instead, test the nonce store directly: issueTOTPNonce must be called once
	// per password-verified login attempt, and the form must carry Nonce not Password.
	nonce, err := u.issueTOTPNonce(acct.ID)
	if err != nil {
		t.Fatalf("issueTOTPNonce: %v", err)
	}

	// Simulate what the template renders: the LoginData for the TOTP step.
	data := LoginData{Nonce: nonce}
	if data.Nonce == "" {
		t.Error("LoginData.Nonce is empty after issueTOTPNonce")
	}
	// Password field must be empty — it is no longer part of LoginData.
	// (This is a compile-time structural guarantee; the field was removed.)
	// The nonce must be non-empty and different on each issuance.
	nonce2, _ := u.issueTOTPNonce(acct.ID)
	if nonce == nonce2 {
		t.Error("two consecutive nonces are identical (randomness failure)")
	}
}

// ---- PG credentials UI tests ----

// issueAdminSession creates a system account with the "admin" role, issues a
// JWT for it, tracks the token in the DB, and returns the raw token string
// along with the account UUID for constructing request paths.
func issueAdminSession(t *testing.T, u *UIServer) (tokenStr, accountUUID string, accountID int64) {
	t.Helper()
	acct, err := u.db.CreateAccount("pgtest-admin", model.AccountTypeHuman, "")
	if err != nil {
		t.Fatalf("CreateAccount: %v", err)
	}
	if err := u.db.SetRoles(acct.ID, []string{"admin"}, nil); err != nil {
		t.Fatalf("SetRoles: %v", err)
	}
	tok, claims, err := token.IssueToken(u.privKey, testIssuer, acct.UUID, []string{"admin"}, time.Hour)
	if err != nil {
		t.Fatalf("IssueToken: %v", err)
	}
	if err := u.db.TrackToken(claims.JTI, acct.ID, claims.IssuedAt, claims.ExpiresAt); err != nil {
		t.Fatalf("TrackToken: %v", err)
	}
	return tok, acct.UUID, acct.ID
}

// authenticatedPUT builds a PUT request to path with the session cookie, a
// valid CSRF cookie + header, and the given form body.
func authenticatedPUT(t *testing.T, u *UIServer, path string, sessionToken string, form url.Values) *http.Request {
	t.Helper()
	body := strings.NewReader(form.Encode())
	req := httptest.NewRequest(http.MethodPut, path, body)
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	req.Header.Set("HX-Request", "true")

	req.AddCookie(&http.Cookie{Name: sessionCookieName, Value: sessionToken})

	csrfCookieVal, csrfHeaderVal, err := u.csrf.NewToken()
	if err != nil {
		t.Fatalf("csrf.NewToken: %v", err)
	}
	req.AddCookie(&http.Cookie{Name: csrfCookieName, Value: csrfCookieVal})
	req.Header.Set("X-CSRF-Token", csrfHeaderVal)
	return req
}

// authenticatedGET builds a GET request with the session cookie set.
func authenticatedGET(t *testing.T, sessionToken string, path string) *http.Request {
	t.Helper()
	req := httptest.NewRequest(http.MethodGet, path, nil)
	req.AddCookie(&http.Cookie{Name: sessionCookieName, Value: sessionToken})
	return req
}

// TestSetPGCredsRejectsHumanAccount verifies that the PUT /accounts/{id}/pgcreds
// endpoint returns 400 when the target account is a human (not system) account.
func TestSetPGCredsRejectsHumanAccount(t *testing.T) {
	u := newTestUIServer(t)
	mux := http.NewServeMux()
	u.Register(mux)

	sessionToken, _, _ := issueAdminSession(t, u)

	// Create a human account to target.
	humanAcct, err := u.db.CreateAccount("human-target", model.AccountTypeHuman, "")
	if err != nil {
		t.Fatalf("CreateAccount: %v", err)
	}

	form := url.Values{
		"host": {"db.example.com"}, "database": {"mydb"},
		"username": {"user"}, "password": {"s3cret"},
	}
	req := authenticatedPUT(t, u, fmt.Sprintf("/accounts/%s/pgcreds", humanAcct.UUID), sessionToken, form)
	rr := httptest.NewRecorder()
	mux.ServeHTTP(rr, req)

	if rr.Code != http.StatusBadRequest {
		t.Errorf("status = %d, want %d", rr.Code, http.StatusBadRequest)
	}
}

// TestSetPGCredsStoresAndDisplaysMetadata verifies that a valid PUT on a system
// account stores the credentials and returns host, database, and username in the
// response body (non-sensitive metadata). The password must not appear.
func TestSetPGCredsStoresAndDisplaysMetadata(t *testing.T) {
	u := newTestUIServer(t)
	mux := http.NewServeMux()
	u.Register(mux)

	sessionToken, _, _ := issueAdminSession(t, u)

	sysAcct, err := u.db.CreateAccount("sys-target", model.AccountTypeSystem, "")
	if err != nil {
		t.Fatalf("CreateAccount: %v", err)
	}

	const (
		testHost     = "pghost.internal"
		testDB       = "appdb"
		testUsername = "appuser"
		testPassword = "super-secret-pw"
	)

	form := url.Values{
		"host": {testHost}, "port": {"5433"}, "database": {testDB},
		"username": {testUsername}, "password": {testPassword},
	}
	req := authenticatedPUT(t, u, fmt.Sprintf("/accounts/%s/pgcreds", sysAcct.UUID), sessionToken, form)
	rr := httptest.NewRecorder()
	mux.ServeHTTP(rr, req)

	if rr.Code != http.StatusOK {
		t.Fatalf("status = %d, want 200; body: %s", rr.Code, rr.Body.String())
	}

	body := rr.Body.String()
	for _, want := range []string{testHost, testDB, testUsername} {
		if !strings.Contains(body, want) {
			t.Errorf("response body missing %q; got:\n%s", want, body)
		}
	}

	// Security: password must not appear anywhere in the response.
	if strings.Contains(body, testPassword) {
		t.Errorf("response body contains plaintext password — security violation")
	}
}

// TestSetPGCredsPasswordNotEchoed explicitly confirms the plaintext password is
// absent from the response even when all other fields are valid.
func TestSetPGCredsPasswordNotEchoed(t *testing.T) {
	u := newTestUIServer(t)
	mux := http.NewServeMux()
	u.Register(mux)

	sessionToken, _, _ := issueAdminSession(t, u)

	sysAcct, err := u.db.CreateAccount("sys-echo-check", model.AccountTypeSystem, "")
	if err != nil {
		t.Fatalf("CreateAccount: %v", err)
	}

	const secretPassword = "must-not-appear-in-response"
	form := url.Values{
		"host": {"h"}, "database": {"d"}, "username": {"u"}, "password": {secretPassword},
	}
	req := authenticatedPUT(t, u, fmt.Sprintf("/accounts/%s/pgcreds", sysAcct.UUID), sessionToken, form)
	rr := httptest.NewRecorder()
	mux.ServeHTTP(rr, req)

	if rr.Code != http.StatusOK {
		t.Fatalf("status = %d, want 200", rr.Code)
	}
	if strings.Contains(rr.Body.String(), secretPassword) {
		t.Error("response body contains the plaintext password — security violation")
	}
}

// TestSetPGCredsRequiresPassword verifies that omitting the password returns an error.
func TestSetPGCredsRequiresPassword(t *testing.T) {
	u := newTestUIServer(t)
	mux := http.NewServeMux()
	u.Register(mux)

	sessionToken, _, _ := issueAdminSession(t, u)

	sysAcct, err := u.db.CreateAccount("sys-no-pw", model.AccountTypeSystem, "")
	if err != nil {
		t.Fatalf("CreateAccount: %v", err)
	}

	form := url.Values{
		"host": {"h"}, "database": {"d"}, "username": {"u"}, // password intentionally absent
	}
	req := authenticatedPUT(t, u, fmt.Sprintf("/accounts/%s/pgcreds", sysAcct.UUID), sessionToken, form)
	rr := httptest.NewRecorder()
	mux.ServeHTTP(rr, req)

	if rr.Code != http.StatusBadRequest {
		t.Errorf("status = %d, want %d (password required)", rr.Code, http.StatusBadRequest)
	}
}

// TestAccountDetailShowsPGCredsSection verifies that the account detail page for
// a system account includes the pgcreds section, and that a human account does not.
func TestAccountDetailShowsPGCredsSection(t *testing.T) {
	u := newTestUIServer(t)
	mux := http.NewServeMux()
	u.Register(mux)

	sessionToken, _, _ := issueAdminSession(t, u)

	sysAcct, err := u.db.CreateAccount("sys-detail", model.AccountTypeSystem, "")
	if err != nil {
		t.Fatalf("CreateAccount (system): %v", err)
	}
	humanAcct, err := u.db.CreateAccount("human-detail", model.AccountTypeHuman, "hash")
	if err != nil {
		t.Fatalf("CreateAccount (human): %v", err)
	}

	// System account detail must include the pgcreds section.
	sysReq := authenticatedGET(t, sessionToken, fmt.Sprintf("/accounts/%s", sysAcct.UUID))
	sysRR := httptest.NewRecorder()
	mux.ServeHTTP(sysRR, sysReq)

	if sysRR.Code != http.StatusOK {
		t.Fatalf("system account GET status = %d, want 200", sysRR.Code)
	}
	if !strings.Contains(sysRR.Body.String(), "pgcreds-section") {
		t.Error("system account detail page missing pgcreds-section")
	}

	// Human account detail must NOT include the pgcreds section.
	humanReq := authenticatedGET(t, sessionToken, fmt.Sprintf("/accounts/%s", humanAcct.UUID))
	humanRR := httptest.NewRecorder()
	mux.ServeHTTP(humanRR, humanReq)

	if humanRR.Code != http.StatusOK {
		t.Fatalf("human account GET status = %d, want 200", humanRR.Code)
	}
	if strings.Contains(humanRR.Body.String(), "pgcreds-section") {
		t.Error("human account detail page must not include pgcreds-section")
	}
}