From 2b7a0e24f80a79155f9fe31b1f8ef0ff9c24bcca Mon Sep 17 00:00:00 2001
From: Sam Calder-Mason <sam@puritydev.io>
Date: Wed, 30 Sep 2026 16:25:29 +1000
Subject: [PATCH] fix(ratelimit): set the counter TTL atomically on every
 request

EXPIRE only ran when INCR returned 1, so a counter that lost its TTL (an
EXPIRE not yet replicated when Redis failed over) never got one again and
blocked its IP forever. 60 of 73 production keys had no TTL, including the
node address kubelet probes come from, which failed every probe on that node.
---
 internal/ratelimit/ratelimit.go      | 34 ++++++++++++++-----------
 internal/ratelimit/ratelimit_test.go | 38 ++++++++++++++++++++++++++++
 2 files changed, 57 insertions(+), 15 deletions(-)

diff --git a/internal/ratelimit/ratelimit.go b/internal/ratelimit/ratelimit.go
index 6aa333e..07c1620 100644
--- a/internal/ratelimit/ratelimit.go
+++ b/internal/ratelimit/ratelimit.go
@@ -77,10 +77,22 @@ func (s *service) Allow(
 ) (bool, int, time.Time, error) {
 	redisKey := fmt.Sprintf("rate_limit:%s:%s", ip, key)
 
-	// Increment counter
-	count, err := s.redis.Incr(ctx, redisKey).Result()
+	// EXPIRE NX runs on every request, not only the first, so a counter that lost its TTL
+	// (e.g. an EXPIRE not yet replicated when Redis failed over) heals instead of growing forever.
+	var (
+		incrCmd *redis.IntCmd
+		ttlCmd  *redis.DurationCmd
+	)
+
+	_, err := s.redis.TxPipelined(ctx, func(pipe redis.Pipeliner) error {
+		incrCmd = pipe.Incr(ctx, redisKey)
+		pipe.ExpireNX(ctx, redisKey, window)
+		ttlCmd = pipe.TTL(ctx, redisKey)
+
+		return nil
+	})
 	if err != nil {
-		s.log.WithError(err).Error("failed to increment rate limit counter in Redis")
+		s.log.WithError(err).Error("failed to update rate limit counter in Redis")
 
 		// Handle failure based on configured mode
 		if s.failureMode == "fail_closed" {
@@ -91,19 +103,11 @@ func (s *service) Allow(
 		return true, 0, time.Time{}, nil
 	}
 
-	// Set expiry on first request
-	if count == 1 {
-		if expireErr := s.redis.Expire(ctx, redisKey, window).Err(); expireErr != nil {
-			s.log.WithError(expireErr).Warn("failed to set rate limit TTL")
-		}
-	}
-
-	// Calculate reset time
-	ttl, err := s.redis.TTL(ctx, redisKey).Result()
-	if err != nil {
-		s.log.WithError(err).Warn("failed to get rate limit TTL")
+	count := incrCmd.Val()
 
-		ttl = window // Fallback
+	ttl := ttlCmd.Val()
+	if ttl <= 0 {
+		ttl = window
 	}
 
 	resetAt := time.Now().Add(ttl)
diff --git a/internal/ratelimit/ratelimit_test.go b/internal/ratelimit/ratelimit_test.go
index ada8d19..bea375f 100644
--- a/internal/ratelimit/ratelimit_test.go
+++ b/internal/ratelimit/ratelimit_test.go
@@ -426,3 +426,41 @@ func TestService_Allow_ConcurrentRequests(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, int64(numGoroutines), count, "all requests should be counted")
 }
+
+// TestService_Allow_RestoresMissingTTL verifies that a counter left without a TTL (as after a
+// Redis failover drops an unreplicated EXPIRE) gets one again instead of blocking forever.
+func TestService_Allow_RestoresMissingTTL(t *testing.T) {
+	mr := miniredis.RunT(t)
+	defer mr.Close()
+
+	client := redis.NewClient(&redis.Options{Addr: mr.Addr()})
+	defer client.Close()
+
+	logger := logrus.New()
+	logger.SetOutput(io.Discard)
+
+	svc := NewService(logger, client, "fail_open")
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	key := "rate_limit:10.42.25.1:health"
+	require.NoError(t, mr.Set(key, "106440"))
+	assert.Equal(t, time.Duration(0), mr.TTL(key), "precondition: counter has no TTL")
+
+	allowed, _, resetAt, err := svc.Allow(ctx, "10.42.25.1", "health", 100000, time.Minute)
+	require.NoError(t, err)
+	assert.False(t, allowed, "counter is still over the limit for the current window")
+	assert.True(t, resetAt.After(time.Now()), "reset time should be in the future")
+
+	ttl := mr.TTL(key)
+	assert.Greater(t, ttl, 50*time.Second, "TTL should be restored")
+	assert.LessOrEqual(t, ttl, time.Minute)
+
+	mr.FastForward(time.Minute + time.Second)
+
+	allowed, remaining, _, err := svc.Allow(ctx, "10.42.25.1", "health", 100000, time.Minute)
+	require.NoError(t, err)
+	assert.True(t, allowed, "counter should reset once the restored TTL expires")
+	assert.Equal(t, 99999, remaining)
+}