selfpost/internal/store/ratelimits.go

package store

import (
	"database/sql"
	"errors"
	"fmt"
	"strings"
	"time"
)

// Rate-limit scopes (guide § Rate limiting). A level-2 limit is attached
// either to a domain (counted across all its applications and IPs) or to a
// single application.
const (
	RateLimitScopeDomain = "domain"
	RateLimitScopeApp    = "application"
)

// Level-2 rate limit operating modes (plan domain-stats-auto-ratelimit).
const (
	RateLimitModeManual = "manual"
	RateLimitModeAuto   = "auto"
)

// Auto rate-limit multiplier bounds shown in the panel.
const (
	DefaultAutoMultiplier = 2.5
	MinAutoMultiplier     = 1.5
	MaxAutoMultiplier     = 5.0
)

// RateLimit is a differentiated level-2 rate limit (guide § Rate limiting):
// a message ceiling over a sliding window, attached to a domain or an
// application. It is enforced in the journal-milter; level 1 (Postfix anvil,
// architecture.md § Mail path) is the IP backstop that always applies even
// when this is absent or the milter is down.
//
// Domain and application limits apply once max and window are set. Client IP
// restriction for an application is stored on the application row, not here
// (guide § Rate limiting).
type RateLimit struct {
	Scope          string
	RefID          int64
	AllowedIPs     []string // legacy column; unused for new rows
	MaxMessages    int
	WindowSeconds  int
	Mode           string  // manual | auto
	AutoMultiplier float64 // used when Mode == auto
	AutoUpdatedAt  time.Time
}

// Active reports whether the limit is fully configured and should be enforced.
// IsAuto reports whether the limit derives max_messages from send statistics.
func (r RateLimit) IsAuto() bool {
	return r.Mode == RateLimitModeAuto
}

func (r RateLimit) Active() bool {
	return r.MaxMessages > 0 && r.WindowSeconds > 0
}

// GetRateLimit loads the level-2 limit configured for a domain or application by
// its id, for the panel's edit form. ok is false when none is configured.
func (s *Store) GetRateLimit(scope string, refID int64) (RateLimit, bool, error) {
	row := s.db.QueryRow(
		`SELECT allowed_ips, max_messages, window_seconds, mode, auto_multiplier, auto_updated_at
		 FROM rate_limits WHERE scope = ? AND ref_id = ?`, scope, refID)
	rl, err := scanRateLimit(row)
	if errors.Is(err, sql.ErrNoRows) {
		return RateLimit{Scope: scope, RefID: refID}, false, nil
	}
	if err != nil {
		return RateLimit{}, false, fmt.Errorf("get rate limit: %w", err)
	}
	rl.Scope, rl.RefID = scope, refID
	return rl, true, nil
}

// SetRateLimit upserts the level-2 limit for a domain or application. The caller
// (panel) has already validated the IPs and numbers (security.md); values are
// stored via bound parameters and read back live by the milter.
func (s *Store) SetRateLimit(rl RateLimit) error {
	mode := rl.Mode
	if mode == "" {
		mode = RateLimitModeManual
	}
	var autoMult interface{}
	if rl.IsAuto() {
		autoMult = rl.AutoMultiplier
	}
	var autoUpdated interface{}
	if !rl.AutoUpdatedAt.IsZero() {
		autoUpdated = rl.AutoUpdatedAt.UTC().Format(time.RFC3339)
	}
	_, err := s.db.Exec(
		`INSERT INTO rate_limits (scope, ref_id, allowed_ips, max_messages, window_seconds, mode, auto_multiplier, auto_updated_at)
		 VALUES (?, ?, ?, ?, ?, ?, ?, ?)
		 ON CONFLICT(scope, ref_id) DO UPDATE SET
		   allowed_ips       = excluded.allowed_ips,
		   max_messages      = excluded.max_messages,
		   window_seconds    = excluded.window_seconds,
		   mode              = excluded.mode,
		   auto_multiplier   = excluded.auto_multiplier,
		   auto_updated_at   = excluded.auto_updated_at`,
		rl.Scope, rl.RefID, strings.Join(rl.AllowedIPs, ","), rl.MaxMessages, rl.WindowSeconds,
		mode, autoMult, autoUpdated,
	)
	if err != nil {
		return fmt.Errorf("set rate limit: %w", err)
	}
	return nil
}

// DeleteRateLimit removes the level-2 limit for a domain or application, so the
// admin can clear it and fall back to level 1 only.
func (s *Store) DeleteRateLimit(scope string, refID int64) error {
	if _, err := s.db.Exec(`DELETE FROM rate_limits WHERE scope = ? AND ref_id = ?`, scope, refID); err != nil {
		return fmt.Errorf("delete rate limit: %w", err)
	}
	return nil
}

// DeleteRateLimitsForDomain removes the domain's own limit and the limits of all
// its applications in one statement. It is called on domain deletion, before the
// application rows are cascade-deleted, to avoid leaving orphan limit rows.
func (s *Store) DeleteRateLimitsForDomain(domainID int64) error {
	_, err := s.db.Exec(
		`DELETE FROM rate_limits
		 WHERE (scope = 'domain' AND ref_id = ?)
		    OR (scope = 'application' AND ref_id IN (SELECT id FROM applications WHERE domain_id = ?))`,
		domainID, domainID,
	)
	if err != nil {
		return fmt.Errorf("delete domain rate limits: %w", err)
	}
	return nil
}

// RateLimit resolves the level-2 limit that applies to a message, keyed by the
// human-readable reference the milter has on the receive path: the sending
// domain name (scope "domain") or the SASL login (scope "application"). ok is
// false when no limit is configured for that reference.
func (s *Store) RateLimit(scope, ref string) (RateLimit, bool, error) {
	var query string
	switch scope {
	case RateLimitScopeDomain:
		query = `SELECT rl.allowed_ips, rl.max_messages, rl.window_seconds, rl.mode, rl.auto_multiplier, rl.auto_updated_at
		         FROM rate_limits rl JOIN domains d ON d.id = rl.ref_id
		         WHERE rl.scope = 'domain' AND d.name = ?`
	case RateLimitScopeApp:
		query = `SELECT rl.allowed_ips, rl.max_messages, rl.window_seconds, rl.mode, rl.auto_multiplier, rl.auto_updated_at
		         FROM rate_limits rl JOIN applications a ON a.id = rl.ref_id
		         WHERE rl.scope = 'application' AND a.login = ?`
	default:
		return RateLimit{}, false, fmt.Errorf("unknown rate-limit scope %q", scope)
	}
	rl, err := scanRateLimit(s.db.QueryRow(query, ref))
	if errors.Is(err, sql.ErrNoRows) {
		return RateLimit{}, false, nil
	}
	if err != nil {
		return RateLimit{}, false, fmt.Errorf("rate limit for %s %q: %w", scope, ref, err)
	}
	rl.Scope = scope
	return rl, true, nil
}

// CountMessages returns how many distinct messages the reference (a domain
// name or an application login) has queued since t, for the level-2 sliding
// window (guide § Rate limiting). It counts distinct queue-ids — one message
// with many recipients is one message, matching level 1's per-message
// semantics — and excludes rows that were themselves rejected by a limit (they
// were never sent). It reuses the send log the journal already writes (README
// § Rate limiting — the limiter reuses the send log).
func (s *Store) CountMessages(scope, ref string, since time.Time) (int64, error) {
	var column string
	switch scope {
	case RateLimitScopeDomain:
		column = "domain"
	case RateLimitScopeApp:
		column = "app_login"
	default:
		return 0, fmt.Errorf("unknown rate-limit scope %q", scope)
	}
	var n int64
	// created_at is stored as RFC3339 UTC, so a lexical comparison against the
	// same format is chronologically correct (as in DeleteSendLogBefore).
	err := s.db.QueryRow(
		`SELECT COUNT(DISTINCT queue_id) FROM send_log
		 WHERE `+column+` = ? AND status != ? AND created_at >= ?`,
		ref, StatusRejected, since.UTC().Format(time.RFC3339),
	).Scan(&n)
	if err != nil {
		return 0, fmt.Errorf("count messages for %s %q: %w", scope, ref, err)
	}
	return n, nil
}

// scanRateLimit reads the three stored columns, tolerating NULL numeric columns
// by leaving the corresponding field zero, which makes the limit inert via
// Active() until max and window are both set.
func scanRateLimit(r scanRow) (RateLimit, error) {
	var (
		ips         sql.NullString
		maxMsgs     sql.NullInt64
		windowSecs  sql.NullInt64
		mode        sql.NullString
		autoMult    sql.NullFloat64
		autoUpdated sql.NullString
	)
	if err := r.Scan(&ips, &maxMsgs, &windowSecs, &mode, &autoMult, &autoUpdated); err != nil {
		return RateLimit{}, err
	}
	rl := RateLimit{
		AllowedIPs:    splitIPs(ips.String),
		MaxMessages:   int(maxMsgs.Int64),
		WindowSeconds: int(windowSecs.Int64),
		Mode:          mode.String,
	}
	if autoMult.Valid {
		rl.AutoMultiplier = autoMult.Float64
	}
	if autoUpdated.Valid {
		rl.AutoUpdatedAt, _ = time.Parse(time.RFC3339, autoUpdated.String)
	}
	return rl, nil
}

// splitIPs parses the comma-separated storage form back into a slice, dropping
// empties so an empty column yields nil (an inactive limit).
func splitIPs(s string) []string {
	if strings.TrimSpace(s) == "" {
		return nil
	}
	parts := strings.Split(s, ",")
	out := make([]string, 0, len(parts))
	for _, p := range parts {
		if p = strings.TrimSpace(p); p != "" {
			out = append(out, p)
		}
	}
	return out
}