Improve rate limit update for catchall

cmd/server/main.go

@@ -259,7 +259,8 @@ func run(ctx context.Context, cfg common.ConfigStore, stderr io.Writer, listener
 	router.Handle("GET "+cdnDomain+"/portal/", http.StripPrefix("/portal/", cdnChain.Then(web.Static(GitCommit))))
 	router.Handle("GET "+cdnDomain+"/widget/", http.StripPrefix("/widget/", cdnChain.Then(widget.Static(GitCommit))))
 	// "protection" (NOTE: different than usual order of monitoring)
-	catchAllChain := alice.New(common.Recovered, metrics.IgnoredHandler, rateLimiter, ipRateLimiter.UpdateLimitsFunc(1, 1*time.Hour))
+	// rate limit update is tricky business due to a possibility for it to be a VPN IP that affects actual users
+	catchAllChain := alice.New(common.Recovered, metrics.IgnoredHandler, rateLimiter, ipRateLimiter.UpdateLimitsFunc(1, 1*time.Minute))
 	portalServer.SetupCatchAll(router, portalDomain, catchAllChain)
 	router.Handle("/", catchAllChain.ThenFunc(common.CatchAll))
 

pkg/ratelimit/http.go

@@ -117,24 +117,20 @@ func (l *httpRateLimiter[TKey]) RateLimit(next http.Handler) http.Handler {
 
 func (l *httpRateLimiter[TKey]) UpdateRequestLimits(r *http.Request, capacity leakybucket.TLevel, leakInterval time.Duration) {
 	ctx := r.Context()
-	if key, ok := ctx.Value(common.RateLimitKeyContextKey).(TKey); ok {
-		l.buckets.Update(key, capacity, leakInterval)
-	} else {
+	key, ok := ctx.Value(common.RateLimitKeyContextKey).(TKey)
+	if !ok {
 		slog.WarnContext(ctx, "Rate limit key not found in http request context")
+		key = l.keyFunc(r)
 	}
+
+	l.buckets.Update(key, capacity, leakInterval)
 }
 
 func (l *httpRateLimiter[TKey]) UpdateLimitsFunc(capacity leakybucket.TLevel, leakInterval time.Duration) func(next http.Handler) http.Handler {
 	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			ctx := r.Context()
-			key, ok := ctx.Value(common.RateLimitKeyContextKey).(TKey)
-			if !ok {
-				slog.WarnContext(ctx, "Rate limit key not found in http request context")
-				key = l.keyFunc(r)
-			}
-
-			l.buckets.Update(key, capacity, leakInterval)
+			l.UpdateRequestLimits(r, capacity, leakInterval)
+			next.ServeHTTP(w, r)
 		})
 	}
 }

pkg/ratelimit/stub.go

@@ -52,3 +52,9 @@ func (srl *StubRateLimiter) UpdateRequestLimits(r *http.Request, capacity leakyb
 func (srl *StubRateLimiter) UpdateLimits(capacity leakybucket.TLevel, leakInterval time.Duration) {
 	// BUMP
 }
+
+func (srl *StubRateLimiter) UpdateLimitsFunc(capacity leakybucket.TLevel, leakInterval time.Duration) func(next http.Handler) http.Handler {
+	return func(next http.Handler) http.Handler {
+		return next
+	}
+}