From 99f5e59a03c482aa83eabf536aa39a90f245e59d Mon Sep 17 00:00:00 2001
From: Patrick Malouin <patrick.malouin@auth0.com>
Date: Mon, 20 Dec 2021 20:30:56 -0500
Subject: [PATCH] IAMRISK-1455 fix(put): keep full bucket when already full
 (#30)

* fix(put): keep full bucket when already full

* fix: improve readability by making intention clearer
---
 .gitignore       |  3 ++-
 lib/put.lua      | 15 ++++++++-------
 test/db.tests.js | 40 ++++++++++++++++++++++++++++++++++++++--
 3 files changed, 48 insertions(+), 10 deletions(-)

diff --git a/.gitignore b/.gitignore
index 3d5eca0..3314550 100644
--- a/.gitignore
+++ b/.gitignore
@@ -63,6 +63,7 @@ typings/
 
 # End of https://www.gitignore.io/api/node
 .vscode
+.idea
 
 # Profiler
 *.0x
@@ -71,4 +72,4 @@ typings/
 dump.rdb
 
 # lock
-package-lock.json
\ No newline at end of file
+package-lock.json
diff --git a/lib/put.lua b/lib/put.lua
index 859b125..fe7baed 100644
--- a/lib/put.lua
+++ b/lib/put.lua
@@ -5,20 +5,21 @@ local ttl                  = tonumber(ARGV[3])
 local current_time = redis.call('TIME')
 local current_timestamp_ms = current_time[1] * 1000 + current_time[2] / 1000
 
-local current = redis.call('HMGET', KEYS[1], 'r')
-
-if current[1] then
-  tokens_to_add = math.min(current[1] + tokens_to_add, bucket_size)
+local current_remaining = redis.call('HMGET', KEYS[1], 'r')[1]
+if current_remaining == false then
+  current_remaining = bucket_size
 end
 
+local new_content = math.min(current_remaining + tokens_to_add, bucket_size)
+
 redis.replicate_commands()
-if tokens_to_add < bucket_size then
+if new_content < bucket_size then
   redis.call('HMSET', KEYS[1],
             'd', current_timestamp_ms,
-            'r', tokens_to_add)
+            'r', new_content)
   redis.call('EXPIRE', KEYS[1], ttl)
 else
   redis.call('DEL', KEYS[1])
 end
 
-return { tokens_to_add, current_timestamp_ms }
+return { new_content, current_timestamp_ms }
diff --git a/test/db.tests.js b/test/db.tests.js
index 6150745..de44015 100644
--- a/test/db.tests.js
+++ b/test/db.tests.js
@@ -536,6 +536,40 @@ describe('LimitDBRedis', () => {
       });
     });
 
+    it('should do nothing if bucket is already full', (done) => {
+      const key = '1.2.3.4';
+      db.put({ type: 'ip', key, count: 1 }, (err, result) => {
+        if (err) {
+          return done(err);
+        }
+        assert.equal(result.remaining, 10);
+
+        db.take({ type: 'ip', key, count: 1 }, (err, result) => {
+          if (err) {
+            return done(err);
+          }
+          assert.equal(result.remaining, 9);
+          done();
+        });
+      });
+    });
+
+    it('should not put more than the bucket size', (done) => {
+      db.take({ type: 'ip', key: '8.8.8.8', count: 2 }, (err) => {
+        if (err) {
+          return done(err);
+        }
+
+        db.put({ type: 'ip', key: '8.8.8.8', count: 4 }, (err, result) => {
+          if (err) {
+            return done(err);
+          }
+          assert.equal(result.remaining, 10);
+          done();
+        });
+      });
+    });
+
     it('should not override on unlimited buckets', (done) => {
       const bucketKey = { type: 'ip',  key: '0.0.0.0', count: 1000 };
       db.put(bucketKey, (err, result) => {
@@ -640,16 +674,18 @@ describe('LimitDBRedis', () => {
     });
 
     it('should work with negative values', (done) => {
-      db.put({ type: 'ip', key: '8.8.8.1', count: -100 }, (err) => {
+      db.put({ type: 'ip', key: '8.8.8.1', count: -100 }, (err, result) => {
         if (err) {
           return done(err);
         }
+        assert.closeTo(result.remaining, -90, 1);
+
         db.take({ type: 'ip', key: '8.8.8.1' }, (err, result) => {
           if (err) {
             return done(err);
           }
           assert.equal(result.conformant, false);
-          assert.closeTo(result.remaining, -99, 1);
+          assert.closeTo(result.remaining, -89, 1);
           done();
         });
       });