Smail/backend/services/verification/rate_limit.go

package verification

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"net"
	"net/netip"
	"strings"
	"sync"
	"time"

	"git.pcux.cn/xuans/Smail/models"
	"xorm.io/xorm"
)

const (
	RateLimitDimensionAccount     = "account"
	RateLimitDimensionEmail       = "email"
	RateLimitDimensionIP          = "ip"
	RateLimitDimensionUser        = "user"
	RateLimitDimensionPurpose     = "purpose"
	RateLimitDimensionTransaction = "transaction"
	RateLimitDimensionEnvironment = "environment"
	RateLimitDimensionCredential  = "credential"

	PersistentRateLimitStateRetention = 30 * 24 * time.Hour
)

var ErrRateLimitStoreUnavailable = errors.New("verification rate limit store unavailable")

type RateLimitRequest struct {
	Account            string
	Email              string
	RemoteIP           string
	UserID             int
	Purpose            Purpose
	BrowserTransaction string
	Environment        string
	CredentialVersion  int
}

type rateLimitDimension struct {
	typ       string
	value     string
	hourLimit int
	dayLimit  int
}

// PersistentRateLimiter 使用数据库事务保存限流计数，数据库中只保存哈希后的维度键。
type PersistentRateLimiter struct {
	engine *xorm.Engine
	mu     sync.Mutex
	now    func() time.Time
}

func NewPersistentRateLimiter(engine *xorm.Engine) (*PersistentRateLimiter, error) {
	if engine == nil {
		return nil, ErrRateLimitStoreUnavailable
	}
	return &PersistentRateLimiter{engine: engine, now: time.Now}, nil
}

func (r *PersistentRateLimiter) Allow(ctx context.Context, request RateLimitRequest, now time.Time) (bool, error) {
	if r == nil || r.engine == nil {
		return false, ErrRateLimitStoreUnavailable
	}
	if ctx == nil {
		ctx = context.Background()
	}
	if now.IsZero() {
		now = r.now()
	}
	dimensions := buildRateLimitDimensions(request)
	if len(dimensions) == 0 {
		return false, ErrRateLimitStoreUnavailable
	}

	// SQLite 使用单连接和数据库锁，MySQL/PostgreSQL 通过行锁保证跨进程原子更新。
	r.mu.Lock()
	defer r.mu.Unlock()
	session := r.engine.NewSession().Context(ctx)
	if err := session.Begin(); err != nil {
		return false, fmt.Errorf("begin verification rate limit transaction: %w", err)
	}
	rollback := func(err error) (bool, error) {
		_ = session.Rollback()
		return false, err
	}
	rows := make([]models.VerificationRateLimit, 0, len(dimensions))
	for _, dimension := range dimensions {
		key := hashRateLimitKey(dimension.typ, dimension.value)
		if err := insertRateLimitRow(session, key, dimension.typ, now); err != nil {
			return rollback(fmt.Errorf("initialize verification rate limit: %w", err))
		}
		var row models.VerificationRateLimit
		query := session.Where("dimension_key = ?", key)
		if r.engine.DriverName() != "sqlite3" && r.engine.DriverName() != "sqlite" {
			query = query.ForUpdate()
		}
		has, err := query.Get(&row)
		if err != nil {
			return rollback(fmt.Errorf("load verification rate limit: %w", err))
		}
		if !has {
			return rollback(ErrRateLimitStoreUnavailable)
		}
		if !rateLimitRowAllowed(row, now, dimension.hourLimit, dimension.dayLimit) {
			return rollback(nil)
		}
		rows = append(rows, row)
	}
	for i, dimension := range dimensions {
		row := rows[i]
		updated := now.UTC()
		if row.HourStartedAt == nil || now.Sub(*row.HourStartedAt) >= issueWindow {
			row.HourStartedAt = timePtr(now)
			row.HourCount = 0
		}
		if row.DayStartedAt == nil || now.Sub(*row.DayStartedAt) >= issueDaily {
			row.DayStartedAt = timePtr(now)
			row.DayCount = 0
		}
		row.HourCount++
		row.DayCount++
		row.LastIssuedAt = timePtr(now)
		row.ExpiresAt = updated.Add(PersistentRateLimitStateRetention)
		row.UpdatedAt = updated
		if _, err := session.ID(row.ID).Cols("hour_started_at", "hour_count", "day_started_at", "day_count", "last_issued_at", "expires_at", "updated_at").Update(&row); err != nil {
			return rollback(fmt.Errorf("update verification rate limit %s: %w", dimension.typ, err))
		}
	}
	if err := session.Commit(); err != nil {
		return false, fmt.Errorf("commit verification rate limit transaction: %w", err)
	}
	return true, nil
}

func (r *PersistentRateLimiter) Cleanup(ctx context.Context, now time.Time) (int, error) {
	if r == nil || r.engine == nil {
		return 0, ErrRateLimitStoreUnavailable
	}
	if now.IsZero() {
		now = r.now()
	}
	cutoff := now.UTC()
	result, err := r.engine.Context(ctx).Where("expires_at < ?", cutoff).Delete(&models.VerificationRateLimit{})
	if err != nil {
		return 0, fmt.Errorf("cleanup verification rate limits: %w", err)
	}
	return int(result), nil
}

func buildRateLimitDimensions(request RateLimitRequest) []rateLimitDimension {
	account := strings.ToLower(strings.TrimSpace(request.Account))
	email := strings.ToLower(strings.TrimSpace(request.Email))
	remoteIP := normalizeRateLimitIP(request.RemoteIP)
	purpose := strings.ToLower(strings.TrimSpace(string(request.Purpose)))
	environment := strings.TrimSpace(request.Environment)
	if environment == "" {
		environment = "default"
	}
	scope := func(value string) string {
		return environment + "\x00" + value
	}
	dimensions := make([]rateLimitDimension, 0, 8)
	appendDimension := func(typ, value string, hourLimit, dayLimit int) {
		if strings.TrimSpace(value) != "" {
			dimensions = append(dimensions, rateLimitDimension{typ: typ, value: value, hourLimit: hourLimit, dayLimit: dayLimit})
		}
	}
	appendDimension(RateLimitDimensionAccount, scope(account), maxPerHour, maxPerDay)
	appendDimension(RateLimitDimensionEmail, scope(email), maxPerHour, maxPerDay)
	appendDimension(RateLimitDimensionIP, scope(remoteIP), maxPerIPHour, maxPerIPDay)
	if request.UserID > 0 {
		appendDimension(RateLimitDimensionUser, scope(fmt.Sprintf("%d", request.UserID)), maxPerHour, maxPerDay)
	}
	if purpose != "" && account != "" {
		appendDimension(RateLimitDimensionPurpose, scope(purpose+"\x00"+account), maxPerHour, maxPerDay)
	}
	if transaction := strings.TrimSpace(request.BrowserTransaction); transaction != "" {
		appendDimension(RateLimitDimensionTransaction, scope(transaction), maxPerHour, maxPerDay)
	}
	if request.CredentialVersion > 0 && account != "" {
		appendDimension(RateLimitDimensionCredential, scope(fmt.Sprintf("%d\x00%s", request.CredentialVersion, account)), maxPerHour, maxPerDay)
	}
	return dimensions
}
func hashRateLimitKey(typ, value string) string {
	hash := sha256.New()
	_, _ = hash.Write([]byte("smail:verification-rate-limit:v1:"))
	_, _ = hash.Write([]byte(typ))
	_, _ = hash.Write([]byte{0})
	_, _ = hash.Write([]byte(value))
	return hex.EncodeToString(hash.Sum(nil))
}

func rateLimitRowAllowed(row models.VerificationRateLimit, now time.Time, hourLimit, dayLimit int) bool {
	if row.LastIssuedAt != nil && now.Sub(*row.LastIssuedAt) < issueInterval {
		return false
	}
	hourCount := row.HourCount
	if row.HourStartedAt == nil || now.Sub(*row.HourStartedAt) >= issueWindow {
		hourCount = 0
	}
	dayCount := row.DayCount
	if row.DayStartedAt == nil || now.Sub(*row.DayStartedAt) >= issueDaily {
		dayCount = 0
	}
	return hourCount < hourLimit && dayCount < dayLimit
}

func insertRateLimitRow(session *xorm.Session, key, typ string, now time.Time) error {
	var query string
	switch session.Engine().DriverName() {
	case "mysql":
		query = "INSERT IGNORE INTO verification_rate_limit (dimension_key, dimension_type, expires_at, updated_at) VALUES (?, ?, ?, ?)"
	case "postgres":
		query = "INSERT INTO verification_rate_limit (dimension_key, dimension_type, expires_at, updated_at) VALUES ($1, $2, $3, $4) ON CONFLICT (dimension_key) DO NOTHING"
	default:
		query = "INSERT OR IGNORE INTO verification_rate_limit (dimension_key, dimension_type, expires_at, updated_at) VALUES (?, ?, ?, ?)"
	}
	updated := now.UTC()
	_, err := session.Exec(query, key, typ, updated.Add(PersistentRateLimitStateRetention), updated)
	return err
}

func normalizeRateLimitIP(value string) string {
	value = strings.TrimSpace(value)
	if value == "" {
		return ""
	}
	if host, _, err := net.SplitHostPort(value); err == nil {
		value = host
	} else {
		value = strings.Trim(value, "[]")
	}
	if address, err := netip.ParseAddr(value); err == nil {
		return address.Unmap().String()
	}
	return strings.ToLower(value)
}

func timePtr(value time.Time) *time.Time {
	value = value.UTC()
	return &value
}