83 lines
2.3 KiB
Go
83 lines
2.3 KiB
Go
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)
|
|
})
|
|
}
|
|
}
|