mailgoserver/internal/webui/mfa_enforcement_test.go

package webui

import (
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"

	"mailgoserver/internal/db"
)

// TestAdminMFAEnforcementForcesSetupThenReleases confirms enforce_admin_mfa blocks
// every other admin page — redirecting to /account, which has the TOTP/passkey
// enrollment forms — until the admin actually sets up a second factor, after which
// normal access resumes.
func TestAdminMFAEnforcementForcesSetupThenReleases(t *testing.T) {
	app := newTestApp(t)
	app.Cfg.Section("Auth").Key("enforce_admin_mfa").SetValue("true")
	mux := app.Mux()

	hash, err := db.HashPassword("no-mfa-yet-password-1!")
	if err != nil {
		t.Fatal(err)
	}
	userID, err := app.DB.CreateAdminUser("no-mfa-admin", hash, false)
	if err != nil {
		t.Fatal(err)
	}
	token, err := app.DB.CreateSession(userID, true, sessionTTL)
	if err != nil {
		t.Fatal(err)
	}
	cookie := &http.Cookie{Name: sessionCookieName, Value: token}

	// Blocked from an ordinary page, redirected to /account.
	req := httptest.NewRequest(http.MethodGet, Prefix+"/domains", nil)
	req.AddCookie(cookie)
	rec := httptest.NewRecorder()
	mux.ServeHTTP(rec, req)
	if rec.Code != http.StatusFound || rec.Header().Get("Location") != Prefix+"/account" {
		t.Fatalf("expected redirect to /account, got %d Location=%q", rec.Code, rec.Header().Get("Location"))
	}

	// /account itself must be reachable (that's where MFA setup happens).
	req = httptest.NewRequest(http.MethodGet, Prefix+"/account", nil)
	req.AddCookie(cookie)
	rec = httptest.NewRecorder()
	mux.ServeHTTP(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("/account: status=%d body=%s", rec.Code, rec.Body.String())
	}
	if !strings.Contains(rec.Body.String(), "requires two-factor authentication") {
		t.Error("expected the MFA-required banner on /account")
	}

	// Once TOTP is enabled, other pages become reachable again.
	if err := app.DB.SetAdminTOTPSecret(userID, "JBSWY3DPEHPK3PXP", true); err != nil {
		t.Fatal(err)
	}
	req = httptest.NewRequest(http.MethodGet, Prefix+"/domains", nil)
	req.AddCookie(cookie)
	rec = httptest.NewRecorder()
	mux.ServeHTTP(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("expected /domains reachable after enabling MFA, got %d", rec.Code)
	}
}

// TestAdminMFAEnforcementOffByDefault confirms nothing changes for existing installs
// unless the admin explicitly turns enforcement on.
func TestAdminMFAEnforcementOffByDefault(t *testing.T) {
	app := newTestApp(t)
	mux := app.Mux()
	cookie := loginSession(t, app)

	req := httptest.NewRequest(http.MethodGet, Prefix+"/domains", nil)
	req.AddCookie(cookie)
	rec := httptest.NewRecorder()
	mux.ServeHTTP(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("expected /domains reachable with enforcement off, got %d", rec.Code)
	}
}

// TestMailboxMFAEnforcementBlocksLogin confirms enforce_mailbox_mfa blocks the
// self-service webmail login outright (no session is ever created) for a mailbox with
// no MFA configured, and that a mailbox-level or domain-level exemption lets the login
// through instead — the bootstrap path for a mailbox to set up its own MFA under
// enforcement. App-password creation/use is deliberately untouched by any of this;
// see mailboxNeedsMFASetup's doc comment.
func TestMailboxMFAEnforcementBlocksLogin(t *testing.T) {
	app := newTestApp(t)
	app.Cfg.Section("Auth").Key("enforce_mailbox_mfa").SetValue("true")
	mux := app.Mux()

	domainID, err := app.DB.CreateDomain("mfatest.example")
	if err != nil {
		t.Fatal(err)
	}
	mhash, err := db.HashPassword("mailbox-owner-password-1!")
	if err != nil {
		t.Fatal(err)
	}
	dek := make([]byte, 32)
	mboxID, err := app.DB.CreateMailbox("owner@mfatest.example", mhash, domainID, 1<<30, dek, dek)
	if err != nil {
		t.Fatal(err)
	}

	tryLogin := func() (status int, sessionCookieSet bool) {
		form := url.Values{"email": {"owner@mfatest.example"}, "password": {"mailbox-owner-password-1!"}}
		req := httptest.NewRequest(http.MethodPost, MailboxPrefix+"/login", strings.NewReader(form.Encode()))
		req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
		rec := httptest.NewRecorder()
		mux.ServeHTTP(rec, req)
		for _, c := range rec.Result().Cookies() {
			if c.Name == mailboxSessionCookieName && c.Value != "" {
				sessionCookieSet = true
			}
		}
		return rec.Code, sessionCookieSet
	}

	status, gotSession := tryLogin()
	if status != http.StatusOK || gotSession {
		t.Fatalf("expected login rejected with no session, got status=%d session=%v", status, gotSession)
	}

	// Mailbox-level exemption lets the login through.
	if err := app.DB.SetMailboxMFAExempt(mboxID, true); err != nil {
		t.Fatal(err)
	}
	status, gotSession = tryLogin()
	if status != http.StatusFound || !gotSession {
		t.Fatalf("expected login to succeed once mailbox-exempt, got status=%d session=%v", status, gotSession)
	}

	// Un-exempt the mailbox but exempt its domain instead — still overrides.
	if err := app.DB.SetMailboxMFAExempt(mboxID, false); err != nil {
		t.Fatal(err)
	}
	if err := app.DB.SetDomainMFAExempt(domainID, true); err != nil {
		t.Fatal(err)
	}
	status, gotSession = tryLogin()
	if status != http.StatusFound || !gotSession {
		t.Fatalf("expected login to succeed once domain-exempt, got status=%d session=%v", status, gotSession)
	}

	// Un-exempt everything, but set up TOTP MFA on the mailbox directly — login
	// succeeds (goes to the pending-MFA step) without needing any exemption at all.
	if err := app.DB.SetDomainMFAExempt(domainID, false); err != nil {
		t.Fatal(err)
	}
	if err := app.DB.SetMailboxTOTPSecret(mboxID, "JBSWY3DPEHPK3PXP", true); err != nil {
		t.Fatal(err)
	}
	form := url.Values{"email": {"owner@mfatest.example"}, "password": {"mailbox-owner-password-1!"}}
	req := httptest.NewRequest(http.MethodPost, MailboxPrefix+"/login", strings.NewReader(form.Encode()))
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	rec := httptest.NewRecorder()
	mux.ServeHTTP(rec, req)
	if rec.Code != http.StatusFound || rec.Header().Get("Location") != MailboxPrefix+"/login/mfa" {
		t.Fatalf("expected redirect to MFA step once TOTP is configured, got %d Location=%q", rec.Code, rec.Header().Get("Location"))
	}
}