mailgoserver/internal/webui/csrf_test.go

package webui

import (
	"io"
	"net/http"
	"net/http/cookiejar"
	"net/http/httptest"
	"net/url"
	"regexp"
	"strings"
	"testing"
)

var csrfTokenInPage = regexp.MustCompile(`window\.__csrfToken\s*=\s*"([0-9a-f]+)"`)

// TestCSRFProtectionAppliesAcrossAdminAndWebmail is a live-HTTP test (real
// httptest.NewServer wrapped exactly like main.go composes it —
// SecurityHeaders(app.CSRFProtect(mux)) — not just httptest.NewRecorder against the
// bare mux) confirming: a forged/missing CSRF token on a state-changing POST is
// rejected for BOTH an admin route and a webmail route, a real page-driven
// submission (token scraped from the actual rendered page, exactly as the injected
// csrf_script.html partial would hand it to a real form) succeeds, and every
// response carries the new security headers.
func TestCSRFProtectionAppliesAcrossAdminAndWebmail(t *testing.T) {
	app := newTestApp(t)
	srv := httptest.NewServer(SecurityHeaders(app.CSRFProtect(app.Mux())))
	defer srv.Close()

	// Security headers present on a plain unauthenticated GET too.
	headResp, err := http.Get(srv.URL + Prefix + "/login")
	if err != nil {
		t.Fatal(err)
	}
	headResp.Body.Close()
	if headResp.Header.Get("X-Frame-Options") != "SAMEORIGIN" {
		t.Fatalf("expected X-Frame-Options on every response, got headers: %v", headResp.Header)
	}
	if headResp.Header.Get("Content-Security-Policy") == "" {
		t.Fatal("expected a Content-Security-Policy header")
	}

	adminCookie := loginSession(t, app)
	domains, _ := app.DB.ListDomains()
	mailboxID := createTestMailboxWithPassword(t, app, "csrf-mailbox@example.com", domains[0].ID, "csrf-password-1!")
	mailboxCookie := webmailLoginSession(t, app, mailboxID)

	jarClient := func(cookie *http.Cookie) *http.Client {
		jar, _ := cookiejar.New(nil)
		u, _ := url.Parse(srv.URL)
		jar.SetCookies(u, []*http.Cookie{cookie})
		return &http.Client{Jar: jar}
	}

	// Regression check: the compose popup (webmail_compose_widget.html) loads
	// /webmail/mail/compose in a same-origin <iframe> — X-Frame-Options: DENY or
	// frame-ancestors 'none' would silently break that popup (this exact regression
	// shipped once already), so explicitly confirm the compose route itself allows
	// same-origin framing.
	composeResp, err := jarClient(mailboxCookie).Get(srv.URL + MailboxPrefix + "/mail/compose")
	if err != nil {
		t.Fatal(err)
	}
	composeResp.Body.Close()
	if fo := composeResp.Header.Get("X-Frame-Options"); fo == "DENY" {
		t.Fatalf("compose route sets X-Frame-Options: DENY — this breaks the compose popup's own same-origin iframe")
	}
	if csp := composeResp.Header.Get("Content-Security-Policy"); strings.Contains(csp, "frame-ancestors 'none'") {
		t.Fatalf("compose route's CSP sets frame-ancestors 'none' — this breaks the compose popup's own same-origin iframe, got: %s", csp)
	}

	scrapeCSRFToken := func(client *http.Client, path string) string {
		t.Helper()
		resp, err := client.Get(srv.URL + path)
		if err != nil {
			t.Fatal(err)
		}
		defer resp.Body.Close()
		body, _ := io.ReadAll(resp.Body)
		m := csrfTokenInPage.FindSubmatch(body)
		if m == nil {
			t.Fatalf("no CSRF token found in rendered page %s: %s", path, body)
		}
		return string(m[1])
	}

	// --- Admin route: use the always-present, state-changing "logout" POST (changing
	// an account setting would need extra setup like enabling TOTP first).
	adminClient := jarClient(adminCookie)
	adminToken := scrapeCSRFToken(adminClient, Prefix+"/")

	forgedResp, err := adminClient.PostForm(srv.URL+Prefix+"/logout", url.Values{"csrf_token": {"forged-not-real"}})
	if err != nil {
		t.Fatal(err)
	}
	forgedResp.Body.Close()
	if forgedResp.StatusCode != http.StatusForbidden {
		t.Fatalf("admin route: expected 403 for a forged CSRF token, got %d", forgedResp.StatusCode)
	}

	missingResp, err := adminClient.PostForm(srv.URL+Prefix+"/logout", url.Values{})
	if err != nil {
		t.Fatal(err)
	}
	missingResp.Body.Close()
	if missingResp.StatusCode != http.StatusForbidden {
		t.Fatalf("admin route: expected 403 for a missing CSRF token, got %d", missingResp.StatusCode)
	}

	realResp, err := adminClient.PostForm(srv.URL+Prefix+"/logout", url.Values{"csrf_token": {adminToken}})
	if err != nil {
		t.Fatal(err)
	}
	realResp.Body.Close()
	if realResp.StatusCode != http.StatusOK && realResp.StatusCode != http.StatusFound {
		t.Fatalf("admin route: expected the real-token logout to succeed, got %d", realResp.StatusCode)
	}

	// --- Webmail route: use the webmail logout POST the same way.
	mailboxClient := jarClient(mailboxCookie)
	mailboxToken := scrapeCSRFToken(mailboxClient, MailboxPrefix+"/mail/INBOX")

	forgedMail, err := mailboxClient.PostForm(srv.URL+MailboxPrefix+"/logout", url.Values{"csrf_token": {"forged-not-real"}})
	if err != nil {
		t.Fatal(err)
	}
	forgedMail.Body.Close()
	if forgedMail.StatusCode != http.StatusForbidden {
		t.Fatalf("webmail route: expected 403 for a forged CSRF token, got %d", forgedMail.StatusCode)
	}

	realMail, err := mailboxClient.PostForm(srv.URL+MailboxPrefix+"/logout", url.Values{"csrf_token": {mailboxToken}})
	if err != nil {
		t.Fatal(err)
	}
	realMail.Body.Close()
	if realMail.StatusCode != http.StatusOK && realMail.StatusCode != http.StatusFound {
		t.Fatalf("webmail route: expected the real-token logout to succeed, got %d", realMail.StatusCode)
	}

	// Also confirm the header-based path works (what the fetch() wrapper uses) —
	// re-login first since the account above just logged itself out.
	mailboxCookie2 := webmailLoginSession(t, app, mailboxID)
	mailboxClient2 := jarClient(mailboxCookie2)
	mailboxToken2 := scrapeCSRFToken(mailboxClient2, MailboxPrefix+"/mail/INBOX")
	req, _ := http.NewRequest(http.MethodPost, srv.URL+MailboxPrefix+"/logout", strings.NewReader(""))
	req.Header.Set("X-CSRF-Token", mailboxToken2)
	headerResp, err := mailboxClient2.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	headerResp.Body.Close()
	if headerResp.StatusCode != http.StatusOK && headerResp.StatusCode != http.StatusFound {
		t.Fatalf("webmail route: expected the header-tokened logout to succeed, got %d", headerResp.StatusCode)
	}
}