jarl/internal/server/totp_test.go

package server

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

	"jarl/internal/api"
	"jarl/internal/auth"
	"jarl/internal/totp"
)

// browser is a web UI session: cookies plus the CSRF token.
type browser struct {
	t    *testing.T
	e    *apiEnv
	c    *http.Client
	csrf string
}

func newBrowser(e *apiEnv) *browser {
	jar, _ := cookiejar.New(nil)
	return &browser{t: e.t, e: e, c: &http.Client{Jar: jar}}
}

// call sends a request; header is an optional X-Jarl-TOTP code.
func (b *browser) call(method, path string, body any, code string) (map[string]any, int) {
	b.t.Helper()
	var rd io.Reader
	if body != nil {
		j, _ := json.Marshal(body)
		rd = bytes.NewReader(j)
	}
	r, _ := http.NewRequest(method, b.e.ts.URL+api.Prefix+path, rd)
	r.Header.Set("Content-Type", "application/json")
	if b.csrf != "" {
		r.Header.Set(csrfHeader, b.csrf)
	}
	if code != "" {
		r.Header.Set(TOTPHeader, code)
	}
	resp, err := b.c.Do(r)
	if err != nil {
		b.t.Fatal(err)
	}
	defer resp.Body.Close()
	var m map[string]any
	json.NewDecoder(resp.Body).Decode(&m)
	if c, ok := m["csrf_token"].(string); ok && c != "" {
		b.csrf = c
	}
	return m, resp.StatusCode
}

func (b *browser) must(method, path string, body any, code string, want int) map[string]any {
	b.t.Helper()
	m, status := b.call(method, path, body, code)
	if status != want {
		b.t.Fatalf("%s %s = %d, want %d: %v", method, path, status, want, m)
	}
	return m
}

// sessionHash is the stored hash of the browser's session cookie.
func (b *browser) sessionHash() string {
	u, _ := url.Parse(b.e.ts.URL)
	for _, c := range b.c.Jar.Cookies(u) {
		if c.Name == sessionCookie {
			return auth.HashSecret(c.Value)
		}
	}
	b.t.Fatal("no session cookie")
	return ""
}

// codeAt is secret's code for a step. Tests count steps from a base
// taken at their start, so crossing a step boundary mid-test changes
// nothing (the server accepts one step either side).
func codeAt(secret string, step int64) string {
	c, _ := totp.Code(secret, step)
	return c
}

func bodyErrCode(m map[string]any) string {
	if e, ok := m["error"].(map[string]any); ok {
		s, _ := e["code"].(string)
		return s
	}
	return ""
}

func TestTOTPLogin(t *testing.T) {
	base := totp.Step(time.Now())
	e := newAPIEnv(t)
	const pw = "correct horse battery"
	b := newBrowser(e)
	b.must("POST", "/auth/login", map[string]string{"username": "admin", "password": pw}, "", 200)
	if st := b.must("GET", "/me/totp", nil, "", 200); st["enabled"] != false {
		t.Fatalf("status = %v", st)
	}
	// Setting up needs the password; the secret only counts once a code
	// confirms it.
	if m, status := b.call("POST", "/me/totp/setup", map[string]string{"password": "wrong"}, ""); status != 403 {
		t.Errorf("setup with a wrong password = %d %v", status, m)
	}
	setup := b.must("POST", "/me/totp/setup", map[string]string{"password": pw}, "", 200)
	secret := setup["secret"].(string)
	if setup["qr"] == "" || setup["account"] != "admin" {
		t.Errorf("setup = %v", setup)
	}
	if m, status := b.call("POST", "/me/totp/enable", map[string]string{"code": "000000"}, ""); status != 400 || bodyErrCode(m) != "totp_invalid" {
		t.Errorf("enable with a wrong code = %d %v", status, m)
	}
	// The codes of three steps are accepted (the clock may be off by one);
	// each works once. Enable with the previous step's, log in with the
	// current one's, and step up with the next one's.
	en := b.must("POST", "/me/totp/enable", map[string]string{"code": codeAt(secret, base-1)}, "", 200)
	recovery := en["recovery_codes"].([]any)
	if len(recovery) != totp.RecoveryCodes {
		t.Fatalf("recovery codes = %v", en)
	}
	// Turning it on kept this session and recorded the proved code: a
	// token can be made right away.
	b.must("POST", "/me/tokens", map[string]any{"name": "laptop"}, "", 201)

	// The password alone is not a session any more.
	b2 := newBrowser(e)
	ch := b2.must("POST", "/auth/login", map[string]string{"username": "admin", "password": pw}, "", 200)
	if ch["totp_required"] != true || ch["challenge"] == "" {
		t.Fatalf("login = %v", ch)
	}
	if _, status := b2.call("GET", "/auth/whoami", nil, ""); status != 401 {
		t.Errorf("whoami after the password only = %d", status)
	}
	// The code used to turn it on can't be used again.
	if m, status := b2.call("POST", "/auth/login/totp", map[string]string{"challenge": ch["challenge"].(string), "code": codeAt(secret, base-1)}, ""); status != 401 || bodyErrCode(m) != "totp_invalid" {
		t.Errorf("reused code = %d %v", status, m)
	}
	who := b2.must("POST", "/auth/login/totp", map[string]string{"challenge": ch["challenge"].(string), "code": codeAt(secret, base)}, "", 200)
	if who["totp"].(map[string]any)["enabled"] != true {
		t.Errorf("whoami = %v", who)
	}
	if _, status := b2.call("POST", "/auth/login/totp", map[string]string{"challenge": ch["challenge"].(string), "code": codeAt(secret, base)}, ""); status != 401 {
		t.Errorf("challenge used twice = %d", status)
	}

	// Step-up: once the proved code is older than the window, making a
	// token needs a fresh one.
	old := time.Now().Add(-stepUpWindow - time.Minute)
	e.st.SetSessionTOTP(context.Background(), b2.sessionHash(), old)
	if m, status := b2.call("POST", "/me/tokens", map[string]any{"name": "ci"}, ""); status != 403 || bodyErrCode(m) != "totp_step_up" {
		t.Errorf("token without a recent code = %d %v", status, m)
	}
	if m, status := b2.call("POST", "/me/tokens", map[string]any{"name": "ci"}, "123456"); status != 403 || bodyErrCode(m) != "totp_invalid" {
		t.Errorf("token with a wrong code = %d %v", status, m)
	}
	b2.must("POST", "/me/tokens", map[string]any{"name": "ci"}, codeAt(secret, base+1), 201)
	// That proved a code again: the next action needs none.
	b2.must("POST", "/me/tokens", map[string]any{"name": "ci2"}, "", 201)
	// API tokens never need a code.
	e.do("POST", "/me/tokens", e.token, map[string]any{"name": "script"}, 201)

	// A recovery code logs in once.
	rc := recovery[0].(string)
	for i, want := range []int{200, 401} {
		b3 := newBrowser(e)
		ch := b3.must("POST", "/auth/login", map[string]string{"username": "admin", "password": pw}, "", 200)
		if _, status := b3.call("POST", "/auth/login/totp", map[string]string{"challenge": ch["challenge"].(string), "code": "  " + rc + " "}, ""); status != want {
			t.Errorf("recovery code use %d = %d, want %d", i+1, status, want)
		}
	}
	if st := b2.must("GET", "/me/totp", nil, "", 200); st["recovery_codes_left"] != float64(totp.RecoveryCodes-1) {
		t.Errorf("recovery codes left = %v", st["recovery_codes_left"])
	}
	evs := e.list("/audit?action=user.totp", e.token)
	if len(evs) == 0 {
		t.Error("no user.totp audit events")
	}
}

func TestTOTPRequiredAndReset(t *testing.T) {
	base := totp.Step(time.Now())
	e := newAPIEnvConfig(t, Config{Version: "test", TOTPRequire: "admins", TOTPIssuer: "Jarl test"})
	const pw = "correct horse battery"
	// An administrator without TOTP gets a session that can only set it up.
	b := newBrowser(e)
	who := b.must("POST", "/auth/login", map[string]string{"username": "admin", "password": pw}, "", 200)
	if who["totp_setup_required"] != true {
		t.Fatalf("whoami = %v", who)
	}
	if m, status := b.call("GET", "/zones", nil, ""); status != 403 || bodyErrCode(m) != "totp_setup_required" {
		t.Errorf("zones before setup = %d %v", status, m)
	}
	// Tokens are not affected.
	e.do("GET", "/zones", e.token, nil, 200)
	setup := b.must("POST", "/me/totp/setup", map[string]string{"password": pw}, "", 200)
	if setup["issuer"] != "Jarl test" {
		t.Errorf("issuer = %v", setup["issuer"])
	}
	secret := setup["secret"].(string)
	b.must("POST", "/me/totp/enable", map[string]string{"code": codeAt(secret, base)}, "", 200)
	b.must("GET", "/zones", nil, "", 200)
	// It can't be turned off while required.
	if _, status := b.call("POST", "/me/totp/disable", map[string]string{"password": pw, "code": codeAt(secret, base+1)}, ""); status != 409 {
		t.Errorf("disable while required = %d", status)
	}
	// Users that aren't administrators may go without.
	e.do("POST", "/users", e.token, map[string]any{"username": "bob", "password": "bob-password-123", "role": "readonly"}, 201)
	bob := newBrowser(e)
	if who := bob.must("POST", "/auth/login", map[string]string{"username": "bob", "password": "bob-password-123"}, "", 200); who["totp_setup_required"] == true {
		t.Error("bob must set up TOTP")
	}
	bob.must("GET", "/zones", nil, "", 200)

	// The dashboard: another administrator without two-factor login is
	// listed, and the first "admin" account looks shared.
	e.do("POST", "/users", e.token, map[string]any{"username": "carol", "password": "carol-password-123", "admin": true}, 201)
	st := e.do("GET", "/status", e.token, nil, 200)["totp"].(map[string]any)
	if missing := st["missing"].([]any); len(missing) != 1 || missing[0] != "carol" || st["shared_admin"] != "admin" || st["policy"] != "admins" {
		t.Errorf("status totp = %v", st)
	}

	// An administrator resets someone's TOTP (with a token: no code).
	users := e.list("/users", e.token)
	var adminID string
	for _, u := range users {
		if u["username"] == "admin" {
			adminID = itoa(int64(u["id"].(float64)))
			if u["totp_required"] != true || u["totp_enabled_at"] == nil {
				t.Errorf("admin = %v", u)
			}
		}
	}
	e.do("DELETE", "/users/"+adminID+"/totp", e.token, nil, 200)
	if _, status := b.call("GET", "/zones", nil, ""); status != 401 {
		t.Errorf("session after reset = %d", status)
	}
	e.do("DELETE", "/users/"+adminID+"/totp", e.token, nil, 409)
	// Next login: password only, and setup again.
	b = newBrowser(e)
	if who := b.must("POST", "/auth/login", map[string]string{"username": "admin", "password": pw}, "", 200); who["totp_setup_required"] != true {
		t.Errorf("after reset = %v", who)
	}
}

func TestTOTPDisable(t *testing.T) {
	base := totp.Step(time.Now())
	e := newAPIEnv(t)
	const pw = "correct horse battery"
	b := newBrowser(e)
	b.must("POST", "/auth/login", map[string]string{"username": "admin", "password": pw}, "", 200)
	secret := b.must("POST", "/me/totp/setup", map[string]string{"password": pw}, "", 200)["secret"].(string)
	b.must("POST", "/me/totp/enable", map[string]string{"code": codeAt(secret, base)}, "", 200)
	if _, status := b.call("POST", "/me/totp/disable", map[string]string{"password": pw, "code": codeAt(secret, base)}, ""); status != 403 {
		t.Errorf("disable with a used code = %d", status)
	}
	b.must("POST", "/me/totp/disable", map[string]string{"password": pw, "code": codeAt(secret, base+1)}, "", 204)
	// The password alone logs in again.
	b2 := newBrowser(e)
	if who := b2.must("POST", "/auth/login", map[string]string{"username": "admin", "password": pw}, "", 200); who["csrf_token"] == "" {
		t.Errorf("login = %v", who)
	}
	// A token can't set up TOTP.
	if _, status := e.req("POST", "/me/totp/setup", e.token, map[string]string{"password": pw}); status != 403 {
		t.Errorf("setup with a token = %d", status)
	}
}