Go upd

rate_limiter.go

@@ -0,0 +1,71 @@
+package main
+
+import (
+	"time"
+
+	"golang.org/x/time/rate"
+)
+
+type userLimiter struct {
+	hourlyLimiter   *rate.Limiter
+	dailyLimiter    *rate.Limiter
+	lastHourlyReset time.Time
+	lastDailyReset  time.Time
+	banUntil        time.Time
+	clock           Clock
+}
+
+func (b *Bot) checkRateLimits(userID int64) bool {
+	b.userLimitersMu.Lock()
+	defer b.userLimitersMu.Unlock()
+
+	limiter, exists := b.userLimiters[userID]
+	if !exists {
+		limiter = &userLimiter{
+			hourlyLimiter:   rate.NewLimiter(rate.Every(time.Hour/time.Duration(b.config.MessagePerHour)), b.config.MessagePerHour),
+			dailyLimiter:    rate.NewLimiter(rate.Every(24*time.Hour/time.Duration(b.config.MessagePerDay)), b.config.MessagePerDay),
+			lastHourlyReset: b.clock.Now(),
+			lastDailyReset:  b.clock.Now(),
+			clock:           b.clock,
+		}
+		b.userLimiters[userID] = limiter
+	}
+
+	now := limiter.clock.Now()
+
+	// Check if the user is currently banned
+	if now.Before(limiter.banUntil) {
+		return false
+	}
+
+	// Reset hourly limiter if an hour has passed since the last reset
+	if now.Sub(limiter.lastHourlyReset) >= time.Hour {
+		limiter.hourlyLimiter = rate.NewLimiter(rate.Every(time.Hour/time.Duration(b.config.MessagePerHour)), b.config.MessagePerHour)
+		limiter.lastHourlyReset = now
+	}
+
+	// Reset daily limiter if 24 hours have passed since the last reset
+	if now.Sub(limiter.lastDailyReset) >= 24*time.Hour {
+		limiter.dailyLimiter = rate.NewLimiter(rate.Every(24*time.Hour/time.Duration(b.config.MessagePerDay)), b.config.MessagePerDay)
+		limiter.lastDailyReset = now
+	}
+
+	// Check if the message exceeds rate limits.
+	// Reserve from both limiters first, then cancel both if either is over budget.
+	// This prevents consuming a token from one limiter when the other rejects.
+	dailyRes := limiter.dailyLimiter.ReserveN(now, 1)
+	hourlyRes := limiter.hourlyLimiter.ReserveN(now, 1)
+	if dailyRes.DelayFrom(now) > 0 || hourlyRes.DelayFrom(now) > 0 {
+		dailyRes.CancelAt(now)
+		hourlyRes.CancelAt(now)
+		banDuration, err := time.ParseDuration(b.config.TempBanDuration)
+		if err != nil {
+			// If parsing fails, default to a 24-hour ban
+			banDuration = 24 * time.Hour
+		}
+		limiter.banUntil = now.Add(banDuration)
+		return false
+	}
+
+	return true
+}