pangolin/server/internal/admin/login_device_handler_test.go

package admin

import (
	"crypto/rand"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"
	"time"

	"github.com/alicebob/miniredis/v2"
	"github.com/redis/go-redis/v9"
	"github.com/wangjia/pangolin/server/internal/totp"
)

// newTrustEnv wires a full admin router with device-trust ENABLED and seeds
// one admin, returning the router and that admin's TOTP secret.
func newTrustEnv(t *testing.T) (http.Handler, string) {
	t.Helper()
	mr := miniredis.RunT(t)
	rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
	t.Cleanup(func() { rdb.Close() })

	key := make([]byte, 32)
	if _, err := rand.Read(key); err != nil {
		t.Fatal(err)
	}
	allow, _ := ParseCIDRs([]string{"127.0.0.0/8"})
	cfg := &Config{
		Listen: "127.0.0.1:9443", AllowCIDRs: allow, SecretKey: key,
		SessionTTL: 30 * time.Minute, LoginFailMax: 3, LoginLockDuration: time.Minute,
		CookieSecure: false, TrustedDeviceTTL: 30 * 24 * time.Hour,
	}
	store := newFakeStore()
	secret := newTestAdmin(t, store, key, "alice", "s3cret-pass")

	sessions := NewSessionStore(rdb, cfg.SessionTTL)
	sec := NewSecurityLog(store, nil)
	auth := NewAuthenticator(store, sessions, rdb, cfg, sec)
	svc := Services{Codes: &fakeCodes{}, Lifecycle: &recordingLifecycle{ready: true}, Provision: &recordingProvision{ready: true}}
	h, err := NewHandlers(cfg, store, sessions, auth, svc, sec, nil)
	if err != nil {
		t.Fatal(err)
	}
	return NewRouter(h, sessions, cfg, sec), secret
}

func cookieByName(resp *http.Response, name string) *http.Cookie {
	for _, c := range resp.Cookies() {
		if c.Name == name {
			return c
		}
	}
	return nil
}

func postLogin(t *testing.T, router http.Handler, form url.Values, cookies ...*http.Cookie) *http.Response {
	t.Helper()
	req := httptest.NewRequest("POST", "/login", strings.NewReader(form.Encode()))
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	req.RemoteAddr = "127.0.0.1:5000"
	for _, c := range cookies {
		req.AddCookie(c)
	}
	rr := httptest.NewRecorder()
	router.ServeHTTP(rr, req)
	return rr.Result()
}

// 端到端:勾选记住 → 登录成功 → 同时下发会话 cookie 与信任 cookie;
// 随后仅凭信任 cookie + 空 TOTP 再次登录成功(免二次验证)。
func TestLoginHandler_RememberThenSkipTOTP(t *testing.T) {
	router, secret := newTrustEnv(t)
	code, _ := totp.Code(secret, time.Now().UTC())

	resp := postLogin(t, router, url.Values{
		"username": {"alice"}, "password": {"s3cret-pass"},
		"totp": {code}, "remember": {"1"},
	})
	if resp.StatusCode != http.StatusFound {
		t.Fatalf("remember login: status = %d, want 302", resp.StatusCode)
	}
	trust := cookieByName(resp, TrustedDeviceCookieName)
	if trust == nil || trust.Value == "" {
		t.Fatal("remember login should set a non-empty trusted cookie")
	}
	if trust.MaxAge <= 0 {
		t.Errorf("trusted cookie MaxAge = %d, want > 0 (persistent)", trust.MaxAge)
	}
	sessCookie := cookieByName(resp, SessionCookieName)
	if sessCookie == nil || sessCookie.MaxAge <= int((30 * time.Minute).Seconds()) {
		t.Error("remember login should set a long-lived (persistent) session cookie")
	}

	// 第二次:只带信任 cookie,TOTP 留空 → 成功
	resp2 := postLogin(t, router, url.Values{
		"username": {"alice"}, "password": {"s3cret-pass"}, "totp": {""},
	}, &http.Cookie{Name: TrustedDeviceCookieName, Value: trust.Value})
	if resp2.StatusCode != http.StatusFound {
		t.Fatalf("trusted re-login: status = %d, want 302 (TOTP skipped)", resp2.StatusCode)
	}
}

// 无信任 cookie + 空 TOTP → 401(二次验证仍强制)。
func TestLoginHandler_NoTrustEmptyTOTPRejected(t *testing.T) {
	router, _ := newTrustEnv(t)
	resp := postLogin(t, router, url.Values{
		"username": {"alice"}, "password": {"s3cret-pass"}, "totp": {""},
	})
	if resp.StatusCode != http.StatusUnauthorized {
		t.Fatalf("empty TOTP without trust: status = %d, want 401", resp.StatusCode)
	}
}