gengx

im/backend/internal/app/rate_limit.go

@@ -0,0 +1,56 @@
+package app
+
+import (
+	"context"
+	"crypto/sha256"
+	"fmt"
+	"net"
+	"net/http"
+	"strings"
+	"time"
+)
+
+func (a *App) allowRequest(ctx context.Context, action, subject string, limit int, window time.Duration) bool {
+	if limit < 1 || window < time.Second {
+		return false
+	}
+	windowSeconds := int64(window / time.Second)
+	slot := time.Now().Unix() / windowSeconds
+	key := sha256.Sum256([]byte(fmt.Sprintf("%s|%s|%d", action, subject, slot)))
+	expiresAt := time.Unix((slot+1)*windowSeconds, 0).Add(time.Minute)
+	_, err := a.db.ExecContext(ctx, `INSERT INTO api_rate_limits(bucket_key,action_name,hits,expires_at) VALUES(?,?,1,?) ON DUPLICATE KEY UPDATE hits=hits+1,expires_at=VALUES(expires_at)`, key[:], action, expiresAt)
+	if err != nil {
+		return false
+	}
+	var hits int
+	if err = a.db.QueryRowContext(ctx, `SELECT hits FROM api_rate_limits WHERE bucket_key=?`, key[:]).Scan(&hits); err != nil {
+		return false
+	}
+	if key[0] == 0 {
+		_, _ = a.db.ExecContext(ctx, `DELETE FROM api_rate_limits WHERE expires_at<NOW(3) LIMIT 1000`)
+	}
+	return hits <= limit
+}
+
+func (a *App) rateLimit(w http.ResponseWriter, r *http.Request, action, subject string, limit int, window time.Duration) bool {
+	if a.allowRequest(r.Context(), action, subject, limit, window) {
+		return true
+	}
+	w.Header().Set("Retry-After", fmt.Sprintf("%d", int(window/time.Second)))
+	fail(w, http.StatusTooManyRequests, 20002, "请求过于频繁，请稍后再试")
+	return false
+}
+
+func clientIP(r *http.Request) string {
+	if forwarded := strings.TrimSpace(strings.Split(r.Header.Get("X-Forwarded-For"), ",")[0]); net.ParseIP(forwarded) != nil {
+		return forwarded
+	}
+	if realIP := strings.TrimSpace(r.Header.Get("X-Real-IP")); net.ParseIP(realIP) != nil {
+		return realIP
+	}
+	host, _, err := net.SplitHostPort(strings.TrimSpace(r.RemoteAddr))
+	if err == nil && net.ParseIP(host) != nil {
+		return host
+	}
+	return "unknown"
+}