feat: 部署初版测试

server/internal/middleware/ratelimit.go

@@ -0,0 +1,82 @@
+package middleware
+
+import (
+	"net/http"
+	"sync"
+	"time"
+
+	"golang.org/x/time/rate"
+	"gorm.io/gorm"
+)
+
+var (
+	limiters   = make(map[string]*rate.Limiter)
+	limiterMux sync.RWMutex
+)
+
+// getLimiter retrieves or creates a rate limiter for a specific user.
+// Uses a simple token bucket. For strict "10 per day" with distributed persistence,
+// this should be refactored to use Redis or DB API usage counters.
+func getLimiter(userID string, tier string) *rate.Limiter {
+	limiterMux.Lock()
+	defer limiterMux.Unlock()
+
+	if limiter, exists := limiters[userID]; exists {
+		return limiter
+	}
+
+	var limiter *rate.Limiter
+	if tier == "Pro" || tier == "Premium" {
+		// Unlimited (e.g., 20 requests per second burst)
+		limiter = rate.NewLimiter(rate.Limit(20), 100)
+	} else {
+		// Free: 10 per day -> replenishes 1 token every 2.4 hours, bucket size 10
+		limiter = rate.NewLimiter(rate.Every(24*time.Hour/10), 10)
+	}
+
+	limiters[userID] = limiter
+	return limiter
+}
+
+// RateLimit middleware enforces rate limits based on user tier.
+// It expects JWTAuth to have already populated UserIDKey in the context.
+func RateLimit(db *gorm.DB) func(http.Handler) http.Handler {
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			userIDVal := r.Context().Value(UserIDKey)
+			if userIDVal == nil {
+				// Allow if not authenticated strictly, or rate limit by IP
+				// For now, fallback to generic tight limit for anonymous usage
+				ipLimiter := getLimiter(r.RemoteAddr, "Free")
+				if !ipLimiter.Allow() {
+					http.Error(w, `{"code":429, "message":"Too Many Requests: Rate limit exceeded"}`, http.StatusTooManyRequests)
+					return
+				}
+				next.ServeHTTP(w, r)
+				return
+			}
+
+			userID := userIDVal.(string)
+
+			// Fast DB query to get user tier (ideally cached in Redis in prod)
+			var tier string
+			// Look up active subscription for this user
+			err := db.Table("subscriptions").
+				Select("tier").
+				Where("user_id = ? AND status = 'active'", userID).
+				Scan(&tier).Error
+
+			if err != nil || tier == "" {
+				tier = "Free" // defaults to Free if no active sub
+			}
+
+			limiter := getLimiter(userID, tier)
+			if !limiter.Allow() {
+				http.Error(w, `{"code":429, "message":"Too Many Requests: Daily quota or rate limit exceeded"}`, http.StatusTooManyRequests)
+				return
+			}
+
+			next.ServeHTTP(w, r)
+		})
+	}
+}