feat: add cache, outbound HTTP, request protection, tasks and module scaffolding

2026-09-23 13:38:25 +08:00
parent 467c890363
commit fe66221a9c
34 changed files with 3100 additions and 1 deletions
--- a/middleware/rateLimit.go
+++ b/middleware/rateLimit.go
@@ -0,0 +1,109 @@
+package middleware
+
+import (
+	"container/list"
+	"errors"
+	"math"
+	"strconv"
+	"sync"
+	"time"
+
+	apierrors "git.yuntuo.hk/k2safe/go-omni/errors"
+	"git.yuntuo.hk/k2safe/go-omni/httpx"
+)
+
+type RateLimitConfig struct {
+	Rate    float64                     `json:"rate"` // Tokens per second, per key.
+	Burst   int                         `json:"burst"`
+	MaxKeys int                         `json:"maxKeys"`
+	IdleTTL time.Duration               `json:"idleTtl"`
+	Key     func(*httpx.Context) string `json:"-"` // Default: trusted ClientIP.
+}
+
+type bucket struct {
+	key     string
+	tokens  float64
+	updated time.Time
+}
+type RateLimiter struct {
+	mu     sync.Mutex
+	config RateLimitConfig
+	items  map[string]*list.Element
+	lru    *list.List
+	now    func() time.Time
+}
+
+func NewRateLimiter(cfg RateLimitConfig) (*RateLimiter, error) {
+	if math.IsNaN(cfg.Rate) || math.IsInf(cfg.Rate, 0) || cfg.Rate <= 0 || cfg.Burst < 1 || cfg.MaxKeys < 0 || cfg.IdleTTL < 0 {
+		return nil, errors.New("middleware: invalid rate limit configuration")
+	}
+	refill := float64(cfg.Burst) / cfg.Rate
+	if refill > (24 * time.Hour).Seconds() {
+		return nil, errors.New("middleware: bucket refill time must not exceed 24h")
+	}
+	minimum := time.Duration(math.Ceil(refill * float64(time.Second)))
+	if cfg.IdleTTL == 0 {
+		cfg.IdleTTL = max(5*time.Minute, minimum)
+	}
+	if cfg.IdleTTL < minimum {
+		return nil, errors.New("middleware: idleTtl must allow a full bucket refill")
+	}
+	if cfg.MaxKeys == 0 {
+		cfg.MaxKeys = 10000
+	}
+	if cfg.Key == nil {
+		cfg.Key = func(c *httpx.Context) string { return c.ClientIP() }
+	}
+	return &RateLimiter{config: cfg, items: make(map[string]*list.Element), lru: list.New(), now: time.Now}, nil
+}
+
+func (l *RateLimiter) Middleware() httpx.Middleware {
+	return func(next httpx.Handler) httpx.Handler {
+		return func(c *httpx.Context) error {
+			key := l.config.Key(c)
+			if key == "" || len(key) > 1024 {
+				return apierrors.BadRequest("INVALID_RATE_KEY", "invalid rate limit key")
+			}
+			delay, full := l.allow(key)
+			if full {
+				return apierrors.New(503, "RATE_LIMIT_CAPACITY", "rate limiter capacity reached")
+			}
+			if delay > 0 {
+				c.SetResponseHeader("Retry-After", strconv.Itoa(max(1, int(math.Ceil(delay.Seconds())))))
+				return apierrors.New(429, "RATE_LIMITED", "too many requests")
+			}
+			return next(c)
+		}
+	}
+}
+
+func (l *RateLimiter) allow(key string) (time.Duration, bool) {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	now := l.now()
+	e := l.items[key]
+	if e == nil {
+		for l.lru.Len() > 0 {
+			old := l.lru.Back()
+			if now.Sub(old.Value.(*bucket).updated) < l.config.IdleTTL {
+				break
+			}
+			delete(l.items, old.Value.(*bucket).key)
+			l.lru.Remove(old)
+		}
+		if len(l.items) >= l.config.MaxKeys {
+			return 0, true
+		}
+		e = l.lru.PushFront(&bucket{key: key, tokens: float64(l.config.Burst), updated: now})
+		l.items[key] = e
+	}
+	b := e.Value.(*bucket)
+	b.tokens = min(float64(l.config.Burst), b.tokens+max(0, now.Sub(b.updated).Seconds())*l.config.Rate)
+	b.updated = now
+	l.lru.MoveToFront(e)
+	if b.tokens >= 1 {
+		b.tokens--
+		return 0, false
+	}
+	return time.Duration(math.Ceil((1 - b.tokens) / l.config.Rate * float64(time.Second))), false
+}