fixed returning modified current_timestamp and reset_ms correctly cal…

…culated for fixed window in takeeleavated
auth0 · Oct 8, 2024 · d9aea0c · d9aea0c
1 parent 62b6672
commit d9aea0c
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 8 deletions.
diff --git a/lib/take.lua b/lib/take.lua
@@ -8,6 +8,7 @@ local fixed_window         = tonumber(ARGV[6])
 
 local current_time = redis.call('TIME')
 local current_timestamp_ms = current_time[1] * 1000 + current_time[2] / 1000
+local redis_timestamp_ms = current_timestamp_ms
 
 local current = redis.pcall('HMGET', KEYS[1], 'd', 'r')
 
@@ -55,4 +56,4 @@ elseif drip_interval > 0 then
     reset_ms = math.ceil(current_timestamp_ms + (bucket_size - new_content) * drip_interval)
 end
 
-return { new_content, enough_tokens, current_timestamp_ms, reset_ms }
+return { new_content, enough_tokens, redis_timestamp_ms, reset_ms }
diff --git a/lib/take_elevated.lua b/lib/take_elevated.lua
@@ -30,6 +30,7 @@ end
 -- get current time from redis, to be used in new bucket size calculations later
 local current_time = redis.call('TIME')
 local current_timestamp_ms = current_time[1] * 1000 + current_time[2] / 1000
+local redis_timestamp_ms = current_timestamp_ms
 
 local function adjustCurrentTimestampForFixedWindow(current_timestamp_ms)
     if current[1] and tokens_per_ms then
@@ -151,7 +152,9 @@ redis.call('HMSET', lastBucketStateKey,
 redis.call('EXPIRE', lastBucketStateKey, ttl)
 
 local reset_ms = 0
-if drip_interval > 0 then
+if fixed_window > 0 then
+    reset_ms = current_timestamp_ms + fixed_window
+elseif drip_interval > 0 then
     if is_erl_activated == 1 then
         reset_ms = math.ceil(current_timestamp_ms + (erl_bucket_size - bucket_content_after_take) * drip_interval)
     else
@@ -160,4 +163,4 @@ if drip_interval > 0 then
 end
 
 -- Return the current quota
-return { bucket_content_after_take, enough_tokens, current_timestamp_ms, reset_ms, erl_triggered, is_erl_activated, erl_quota_left }
+return { bucket_content_after_take, enough_tokens, redis_timestamp_ms, reset_ms, erl_triggered, is_erl_activated, erl_quota_left }
diff --git a/test/db.tests.js b/test/db.tests.js
@@ -458,7 +458,7 @@ module.exports.tests = (clientCreator) => {
               }
               assert.ok(result.conformant);
               assert.equal(result.remaining, 0);
-              assert.closeTo(result.reset, now / 1000 + 1800, 1);
+              assert.closeTo(result.reset, now / 1000 + 1800, 3);
               assert.closeTo(result.delta_reset_ms, (result.limit - result.remaining) * 3600000/buckets.ip.overrides['10.0.0.1'].per_hour, 1);
               assert.equal(result.limit, 1);
               done();
@@ -878,30 +878,29 @@ module.exports.tests = (clientCreator) => {
 
           describe(`when calling the lua script`, () => {
             it(`should use fixed window when asked`, (done) => {
-              const now = Date.now();
               const interval = 1000;
               const key = '123.123.123.123';
               takeFunc({ type: 'ip', key, count: 1000, fixed_window: true }, (err, response) => {
                 if (err) return done(err);
                 assert.ok(response.conformant);
                 assert.equal(response.remaining, 0);
-                assert.closeTo(response.reset, Math.ceil((now + interval) / 1000), 1);
+                assert.closeTo(response.delta_reset_ms, interval, 100);
                 assert.equal(response.limit, 1000);
                 redisHMGetPromise(`ip:${key}`, ['d', 'r']).then((value) => {
                   const lastDrip = value[0];
                   setTimeout(() => {
                     takeFunc({ type: 'ip', key, count: 1, fixed_window: true }, (err, response) => {
                       assert.notOk(response.conformant);
                       assert.equal(response.remaining, 0);
-                      assert.closeTo(response.reset, Math.ceil((now + interval) / 1000), 1);
+                      assert.closeTo(response.delta_reset_ms, interval/2, 100);
                       assert.equal(response.limit, 1000);
                       redisHMGetPromise(`ip:${key}`, ['d', 'r']).then((value) => {
                         assert.equal(value[0], lastDrip, 'last drip should not have changed');
                         setTimeout(() => {
                           takeFunc({ type: 'ip', key, count: 1, fixed_window: true }, (err, response) => {
                             assert.ok(response.conformant);
                             assert.equal(response.remaining, 999);
-                            assert.closeTo(response.reset, Math.ceil((Date.now() + interval) / 1000), 1);
+                            assert.closeTo(response.delta_reset_ms, interval, 100);
                             assert.equal(response.limit, 1000);
                             redisHMGetPromise(`ip:${key}`, ['d', 'r']).then((value) => {
                               assert.notEqual(value[0], lastDrip, 'last drip should have changed');