[Project] Enable compatibility with the existing buckets

3 files changed, 78 insertions, 42 deletions

diff --git a/lualib/redis_scripts/ratelimit_check.lua b/lualib/redis_scripts/ratelimit_check.lua
index d39cdf148..f24e0daf0 100644
--- a/lualib/redis_scripts/ratelimit_check.lua
+++ b/lualib/redis_scripts/ratelimit_check.lua
@@ -7,6 +7,7 @@
 -- KEYS[4]: The maximum allowed burst
 -- KEYS[5]: The expiration time for a bucket
 -- KEYS[6]: The number of recipients for the message
+-- KEYS[7]: Enable dynamic ratelimits
 
 -- Redis keys used:
 -- l: Last hit (time in milliseconds)
@@ -29,6 +30,7 @@ local nrcpt = tonumber(KEYS[6])
 local leak_rate = tonumber(KEYS[3])
 local max_burst = tonumber(KEYS[4])
 local prefix = KEYS[1]
+local enable_dynamic = KEYS[7] == 'true'
 local dynr, dynb, leaked = 0, 0, 0
 if not last then
   -- New bucket
@@ -52,9 +54,13 @@ pending = pending + nrcpt -- this message
 if burst + pending > 0 then
   -- If we have any time passed
   if burst > 0 and last < now then
-    dynr = tonumber(redis.call('HGET', prefix, 'dr')) / 10000.0
-    if dynr == 0 then
-      dynr = 0.0001
+    if enable_dynamic then
+      dynr = tonumber(redis.call('HGET', prefix, 'dr')) / 10000.0
+      if dynr == 0 then
+        dynr = 0.0001
+      end
+    else
+      dynr = 1.0
     end
     leak_rate = leak_rate * dynr
     leaked = ((now - last) * leak_rate)
@@ -66,9 +72,13 @@ if burst + pending > 0 then
     redis.call('HSET', prefix, 'l', tostring(now))
   end
 
-  dynb = tonumber(redis.call('HGET', prefix, 'db')) / 10000.0
-  if dynb == 0 then
-    dynb = 0.0001
+  if enable_dynamic then
+    dynb = tonumber(redis.call('HGET', prefix, 'db')) / 10000.0
+    if dynb == 0 then
+      dynb = 0.0001
+    end
+  else
+    dynb = 1.0
   end
 
   burst = burst + pending
diff --git a/lualib/redis_scripts/ratelimit_update.lua b/lualib/redis_scripts/ratelimit_update.lua
index caee8fb31..8b7a934dc 100644
--- a/lualib/redis_scripts/ratelimit_update.lua
+++ b/lualib/redis_scripts/ratelimit_update.lua
@@ -9,12 +9,14 @@
 -- KEYS[6] - max_burst_rate: The maximum allowed value for the dynamic burst multiplier.
 -- KEYS[7] - expire: The expiration time for the Redis key storing the bucket information, in seconds.
 -- KEYS[8] - number_of_recipients: The number of requests to be allowed (or the increase rate).
+-- KEYS[9] - Enable dynamic ratelimits
 
 -- 1. Retrieve the last hit time and initialize variables
 local prefix = KEYS[1]
 local last = redis.call('HGET', prefix, 'l')
 local now = tonumber(KEYS[2])
 local nrcpt = tonumber(KEYS[8])
+local enable_dynamic = KEYS[9] == 'true'
 if not last then
   -- 2. Initialize a new bucket if the last hit time is not found (must not happen)
   redis.call('HMSET', prefix, 'l', tostring(now), 'b', tostring(nrcpt), 'dr', '10000', 'db', '10000', 'p', '0')
@@ -25,48 +27,52 @@ end
 -- 3. Update the dynamic rate multiplier based on input parameters
 local dr, db = 1.0, 1.0
 
-local max_dr = tonumber(KEYS[5])
+if enable_dynamic then
+  local max_dr = tonumber(KEYS[5])
 
-if max_dr > 1 then
-  local rate_mult = tonumber(KEYS[3])
-  dr = tonumber(redis.call('HGET', prefix, 'dr')) / 10000
+  if max_dr > 1 then
+    local rate_mult = tonumber(KEYS[3])
+    dr = tonumber(redis.call('HGET', prefix, 'dr')) / 10000
 
-  if rate_mult > 1.0 and dr < max_dr then
-    dr = dr * rate_mult
-    if dr > 0.0001 then
-      redis.call('HSET', prefix, 'dr', tostring(math.floor(dr * 10000)))
-    else
-      redis.call('HSET', prefix, 'dr', '1')
-    end
-  elseif rate_mult < 1.0 and dr > (1.0 / max_dr) then
-    dr = dr * rate_mult
-    if dr > 0.0001 then
-      redis.call('HSET', prefix, 'dr', tostring(math.floor(dr * 10000)))
-    else
-      redis.call('HSET', prefix, 'dr', '1')
+    if rate_mult > 1.0 and dr < max_dr then
+      dr = dr * rate_mult
+      if dr > 0.0001 then
+        redis.call('HSET', prefix, 'dr', tostring(math.floor(dr * 10000)))
+      else
+        redis.call('HSET', prefix, 'dr', '1')
+      end
+    elseif rate_mult < 1.0 and dr > (1.0 / max_dr) then
+      dr = dr * rate_mult
+      if dr > 0.0001 then
+        redis.call('HSET', prefix, 'dr', tostring(math.floor(dr * 10000)))
+      else
+        redis.call('HSET', prefix, 'dr', '1')
+      end
     end
   end
 end
 
 -- 4. Update the dynamic burst multiplier based on input parameters
-local max_db = tonumber(KEYS[6])
-if max_db > 1 then
-  local rate_mult = tonumber(KEYS[4])
-  db = tonumber(redis.call('HGET', prefix, 'db')) / 10000
+if enable_dynamic then
+  local max_db = tonumber(KEYS[6])
+  if max_db > 1 then
+    local rate_mult = tonumber(KEYS[4])
+    db = tonumber(redis.call('HGET', prefix, 'db')) / 10000
 
-  if rate_mult > 1.0 and db < max_db then
-    db = db * rate_mult
-    if db > 0.0001 then
-      redis.call('HSET', prefix, 'db', tostring(math.floor(db * 10000)))
-    else
-      redis.call('HSET', prefix, 'db', '1')
-    end
-  elseif rate_mult < 1.0 and db > (1.0 / max_db) then
-    db = db * rate_mult
-    if db > 0.0001 then
-      redis.call('HSET', prefix, 'db', tostring(math.floor(db * 10000)))
-    else
-      redis.call('HSET', prefix, 'db', '1')
+    if rate_mult > 1.0 and db < max_db then
+      db = db * rate_mult
+      if db > 0.0001 then
+        redis.call('HSET', prefix, 'db', tostring(math.floor(db * 10000)))
+      else
+        redis.call('HSET', prefix, 'db', '1')
+      end
+    elseif rate_mult < 1.0 and db > (1.0 / max_db) then
+      db = db * rate_mult
+      if db > 0.0001 then
+        redis.call('HSET', prefix, 'db', tostring(math.floor(db * 10000)))
+      else
+        redis.call('HSET', prefix, 'db', '1')
+      end
     end
   end
 end
diff --git a/src/plugins/lua/ratelimit.lua b/src/plugins/lua/ratelimit.lua
index 470ac5f07..c3602f915 100644
--- a/src/plugins/lua/ratelimit.lua
+++ b/src/plugins/lua/ratelimit.lua
@@ -365,17 +365,33 @@ local function make_prefix(redis_key, name, bucket)
   -- Fill defaults
   -- If settings.dynamic_rate_limit is false, then the default dynamic rate limits are 1.0
   -- We always allow per-bucket overrides of the dyn rate limits
+
+  local seen_specific_dyn_rate = false
+
   if not bucket.spam_factor_rate then
     bucket.spam_factor_rate = settings.dynamic_rate_limit and settings.spam_factor_rate or 1.0
+  else
+    seen_specific_dyn_rate = true
   end
   if not bucket.ham_factor_rate then
     bucket.ham_factor_rate = settings.dynamic_rate_limit and settings.ham_factor_rate or 1.0
+  else
+    seen_specific_dyn_rate = true
   end
   if not bucket.spam_factor_burst then
     bucket.spam_factor_burst = settings.dynamic_rate_limit and settings.spam_factor_burst or 1.0
+  else
+    seen_specific_dyn_rate = true
   end
   if not bucket.ham_factor_burst then
     bucket.ham_factor_burst = settings.dynamic_rate_limit and settings.ham_factor_burst or 1.0
+  else
+    seen_specific_dyn_rate = true
+  end
+
+  if seen_specific_dyn_rate then
+    -- Use if afterwards in case we don't use global dyn rates
+    bucket.specific_dyn_rate = true
   end
 
   return {
@@ -555,13 +571,15 @@ local function ratelimit_cb(task)
         bincr = 1
       end
 
+      local dyn_rate_enabled = settings.dynamic_rate_limit or bucket.specific_dyn_rate
+
       lua_util.debugm(N, task, "check limit %s:%s -> %s (%s/%s)",
           value.name, pr, value.hash, bucket.burst, bucket.rate)
       lua_redis.exec_redis_script(bucket_check_id,
           { key = value.hash, task = task, is_write = true },
           gen_check_cb(pr, bucket, value.name, value.hash),
           { value.hash, tostring(now), tostring(rate), tostring(bucket.burst),
-            tostring(settings.expire), tostring(bincr) })
+            tostring(settings.expire), tostring(bincr), tostring(dyn_rate_enabled) })
     end
   end
 end
@@ -661,12 +679,14 @@ local function ratelimit_update_cb(task)
         bincr = 1
       end
 
+      local dyn_rate_enabled = settings.dynamic_rate_limit or bucket.specific_dyn_rate
+
       lua_redis.exec_redis_script(bucket_update_id,
           { key = v.hash, task = task, is_write = true },
           update_bucket_cb,
           { v.hash, tostring(now), tostring(mult_rate), tostring(mult_burst),
             tostring(settings.max_rate_mult), tostring(settings.max_bucket_mult),
-            tostring(settings.expire), tostring(bincr) })
+            tostring(settings.expire), tostring(bincr), tostring(dyn_rate_enabled) })
     end
   end
 end