Smail/backend/services/verification/rate_limit_test.go

package verification

import (
	"context"
	"errors"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"git.pcux.cn/xuans/Smail/models"
	_ "modernc.org/sqlite"
	"xorm.io/xorm"
)

func newRateLimitEngine(t *testing.T) *xorm.Engine {
	t.Helper()
	engine, err := xorm.NewEngine("sqlite", filepath.Join(t.TempDir(), "rate-limit.db"))
	if err != nil {
		t.Fatal(err)
	}
	if err := engine.Sync2(&models.VerificationRateLimit{}); err != nil {
		_ = engine.Close()
		t.Fatal(err)
	}
	return engine
}

func TestPersistentRateLimiterSurvivesServiceRestart(t *testing.T) {
	engine := newRateLimitEngine(t)
	defer engine.Close()
	limiter, err := NewPersistentRateLimiter(engine)
	if err != nil {
		t.Fatal(err)
	}
	first, err := NewStore(32)
	if err != nil {
		t.Fatal(err)
	}
	second, err := NewStore(32)
	if err != nil {
		t.Fatal(err)
	}
	service1, err := NewService(first, func(_ Challenge, _ string) error { return nil }, nil)
	if err != nil {
		t.Fatal(err)
	}
	service2, err := NewService(second, func(_ Challenge, _ string) error { return nil }, nil)
	if err != nil {
		t.Fatal(err)
	}
	service1.SetRateLimitStore(limiter)
	service2.SetRateLimitStore(limiter)
	now := time.Date(2026, 8, 26, 10, 0, 0, 0, time.UTC)
	service1.now = func() time.Time { return now }
	service2.now = func() time.Time { return now }
	for i := 0; i < maxPerHour; i++ {
		if _, err := service1.Issue(IssueRequest{Purpose: PurposeLogin, Account: "smail", Email: "owner@example.com", TTL: time.Second}, "[2001:0db8::1]:443"); err != nil {
			t.Fatalf("issue %d: %v", i, err)
		}
		now = now.Add(issueInterval)
	}
	if _, err := service2.Issue(IssueRequest{Purpose: PurposeLogin, Account: "smail", Email: "owner@example.com"}, "2001:db8::1"); !errors.Is(err, ErrRateLimited) {
		t.Fatalf("persistent account limit was bypassed after restart: %v", err)
	}
	if _, err := limiter.Cleanup(context.Background(), now); err != nil {
		t.Fatal(err)
	}
}

func TestPersistentRateLimiterAppliesIPLimitAcrossIdentities(t *testing.T) {
	engine := newRateLimitEngine(t)
	defer engine.Close()
	limiter, err := NewPersistentRateLimiter(engine)
	if err != nil {
		t.Fatal(err)
	}
	now := time.Date(2026, 8, 26, 12, 0, 0, 0, time.UTC)
	for i := 0; i < maxPerIPHour; i++ {
		allowed, err := limiter.Allow(context.Background(), RateLimitRequest{
			Account:  "user-" + string(rune('a'+i)),
			Email:    "user-" + string(rune('a'+i)) + "@example.com",
			RemoteIP: "192.0.2.10",
			Purpose:  PurposeLogin,
		}, now)
		if err != nil || !allowed {
			t.Fatalf("allow %d = %v, %v", i, allowed, err)
		}
		now = now.Add(issueInterval)
	}
	allowed, err := limiter.Allow(context.Background(), RateLimitRequest{Account: "other", Email: "other@example.com", RemoteIP: "192.0.2.10", Purpose: PurposePasswordRecover}, now)
	if err != nil {
		t.Fatal(err)
	}
	if allowed {
		t.Fatal("IP limit was bypassed by changing account, email, and purpose")
	}
}

func TestPersistentRateLimiterDoesNotPersistSensitiveValues(t *testing.T) {
	engine := newRateLimitEngine(t)
	defer engine.Close()
	limiter, err := NewPersistentRateLimiter(engine)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := limiter.Allow(context.Background(), RateLimitRequest{Account: "smail", Email: "owner@example.com", RemoteIP: "127.0.0.1", Purpose: PurposePasswordRecover}, time.Now()); err != nil {
		t.Fatal(err)
	}
	var raw string
	if _, err := engine.SQL("SELECT dimension_key || '|' || dimension_type FROM verification_rate_limit").Get(&raw); err != nil {
		t.Fatal(err)
	}
	if raw == "" || containsAny(raw, "smail", "owner@example.com", "127.0.0.1") {
		t.Fatalf("sensitive value persisted: %q", raw)
	}
}

func containsAny(value string, values ...string) bool {
	for _, item := range values {
		if strings.Contains(value, item) {
			return true
		}
	}
	return false
}