Merge pull request #4 from sourabpramanik/feat/v_0.1.0

Feat/v 0.1.0

.github/workflows/build.yml

@@ -0,0 +1,23 @@
+name: Rust
+
+on:
+  push:
+    branches: [ "main" ]
+  pull_request:
+    branches: [ "main" ]
+
+env:
+  CARGO_TERM_COLOR: always
+  UPSTASH_REDIS_URL: ${{secrets.UPSTASH_REDIS_URL}}
+
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v3
+    - name: Build
+      run: cargo build --verbose
+    - name: Run tests
+      run: cargo test --verbose

.gitignore

@@ -1 +1,5 @@
 /target
+
+Cargo.lock
+
+.env

.rustfmt.toml

@@ -0,0 +1,13 @@
+max_width = 150
+hard_tabs = true
+tab_spaces = 4
+newline_style = "Unix"
+use_small_heuristics = "Default"
+reorder_imports = true
+reorder_modules = true
+remove_nested_parens = true
+edition = "2018"
+merge_derives = true
+use_try_shorthand = false
+use_field_init_shorthand = true
+force_explicit_abi = true

Cargo.toml

@@ -13,7 +13,9 @@ license = "MIT"
 [dependencies]
 assert-panic = "1.0.1"
 chrono = "0.4"
+dotenv = "0.15"
 futures = "0.3"
+nanoid = "0.4.0"
 redis = {version = "0.24.0", features = ["tokio-comp"]}
 regex = "1.10.3"
 tokio = { version = "1.36.0", features = ["full"] }

scripts/single_region/fixed_window.lua

@@ -1,11 +1,12 @@
-local key     = KEYS[1]
-local window  = ARGV[1]
+local key           = KEYS[1]
+local window        = ARGV[1]
+local incrementBy   = ARGV[2] -- increment rate per request at a given value, default is 1
 
-local value = redis.call('INCR', key)
-if value == 1 then 
--- The first time this key is set, the value will be 1.
+local r = redis.call("INCRBY", key, incrementBy)
+if r == incrementBy then
+-- The first time this key is set, the value will be equal to incrementBy.
 -- So we only need the expire command once
-redis.call('PEXPIRE', key, window)
+redis.call("PEXPIRE", key, window)
 end
 
-return value
+return r

scripts/single_region/sliding_window.lua

@@ -0,0 +1,30 @@
+local currentKey  = KEYS[1]           -- identifier including prefixes
+local previousKey = KEYS[2]           -- key of the previous bucket
+local tokens      = tonumber(ARGV[1]) -- tokens per window
+local now         = ARGV[2]           -- current timestamp in milliseconds
+local window      = ARGV[3]           -- interval in milliseconds
+local incrementBy = ARGV[4]           -- increment rate per request at a given value, default is 1
+
+local requestsInCurrentWindow = redis.call("GET", currentKey)
+if requestsInCurrentWindow == false then
+  requestsInCurrentWindow = 0
+end
+
+local requestsInPreviousWindow = redis.call("GET", previousKey)
+if requestsInPreviousWindow == false then
+  requestsInPreviousWindow = 0
+end
+local percentageInCurrent = ( now % window ) / window
+-- weighted requests to consider from the previous window
+requestsInPreviousWindow = math.floor(( incrementBy - percentageInCurrent ) * requestsInPreviousWindow)
+if requestsInPreviousWindow + requestsInCurrentWindow >= tokens then
+  return -1
+end
+
+local newValue = redis.call("INCRBY", currentKey, incrementBy)
+if newValue == incrementBy then
+  -- The first time this key is set, the value will be equal to incrementBy.
+  -- So we only need the expire command once
+  redis.call("PEXPIRE", currentKey, window * 2 + 1000) -- Enough time to overlap with a new window + 1 second
+end
+return tokens - ( newValue + requestsInPreviousWindow )

scripts/single_region/token_window.lua

@@ -0,0 +1,37 @@
+  local key         = KEYS[1]           -- identifier including prefixes
+  local maxTokens   = tonumber(ARGV[1]) -- maximum number of tokens
+  local interval    = tonumber(ARGV[2]) -- size of the window in milliseconds
+  local refillRate  = tonumber(ARGV[3]) -- how many tokens are refilled after each interval
+  local now         = tonumber(ARGV[4]) -- current timestamp in milliseconds
+  local incrementBy = tonumber(ARGV[5]) -- how many tokens to consume, default is 1
+
+  local bucket = redis.call("HMGET", key, "refilledAt", "tokens")
+
+  local refilledAt
+  local tokens
+
+  if bucket[1] == false then
+    refilledAt = now
+    tokens = maxTokens
+  else
+    refilledAt = tonumber(bucket[1])
+    tokens = tonumber(bucket[2])
+  end
+
+  if now >= refilledAt + interval then
+    local numRefills = math.floor((now - refilledAt) / interval)
+    tokens = math.min(maxTokens, tokens + numRefills * refillRate)
+
+    refilledAt = refilledAt + numRefills * interval
+  end
+
+  if tokens == 0 then
+    return {-1, refilledAt + interval}
+  end
+
+  local remaining = tokens - incrementBy
+  local expireAt = math.ceil(((maxTokens - remaining) / refillRate)) * interval
+
+  redis.call("HSET", key, "refilledAt", refilledAt, "tokens", remaining)
+  redis.call("PEXPIRE", key, expireAt)
+  return {remaining, refilledAt + interval}

src/lib.rs

@@ -1 +1 @@
-mod ratelimit;
+pub mod ratelimit;

src/ratelimit/cache.rs

@@ -1,79 +1,56 @@
 use std::{
-    collections::HashMap,
-    time::{SystemTime, UNIX_EPOCH},
+	collections::HashMap,
+	time::{SystemTime, UNIX_EPOCH},
 };
 
 #[derive(Debug)]
 pub struct Blocked {
-    pub blocked: bool,
-    pub reset: u128,
+	pub blocked: bool,
+	pub reset: u128,
 }
 
 #[derive(Debug, Clone)]
 pub struct EphemeralCache {
-    pub cache: HashMap<String, u128>,
+	pub cache: HashMap<String, u128>,
 }
-
-impl EphemeralCache {
-    pub fn new() -> EphemeralCache {
-        EphemeralCache {
-            cache: HashMap::new(),
-        }
-    }
-
-    pub fn is_blocked(&self, identifier: &str) -> Blocked {
-        let reset = self.cache.get(identifier);
-
-        match reset {
-            Some(value) => {
-                if *value
-                    < SystemTime::now()
-                        .duration_since(UNIX_EPOCH)
-                        .expect("Time went backwards")
-                        .as_millis()
-                {
-                    return Blocked {
-                        blocked: false,
-                        reset: 0,
-                    };
-                }
-                Blocked {
-                    blocked: true,
-                    reset: value.to_owned(),
-                }
-            }
-            None => Blocked {
-                blocked: false,
-                reset: 0,
-            },
-        }
-    }
-    pub fn block_until(&mut self, identifier: &str, reset: u128) {
-        self.cache.insert(identifier.to_owned(), reset);
-    }
-    pub fn set(&mut self, identifier: &str, reset: u128) {
-        self.cache.insert(identifier.to_owned(), reset);
-    }
-    pub fn get(&self, identifier: &str) -> Option<u128> {
-        self.cache.get(identifier).copied()
-    }
-    pub fn incr(&mut self, identifier: &str) {
-        let mut value = self.get(identifier).unwrap_or_else(|| 0);
-        value += 1;
-        self.cache.insert(identifier.to_owned(), value);
-    }
+impl Default for EphemeralCache {
+	fn default() -> Self {
+		Self::new()
+	}
 }
+impl EphemeralCache {
+	pub fn new() -> EphemeralCache {
+		EphemeralCache { cache: HashMap::new() }
+	}
 
-#[cfg(test)]
-mod tests {
-    use super::EphemeralCache;
+	pub fn is_blocked(&self, identifier: &str) -> Blocked {
+		let reset = self.cache.get(identifier);
 
-    #[test]
-    fn test_get_cache() {
-        let key = "ip_address";
-        let mut test_cache = EphemeralCache::new();
-        test_cache.set(key, 2000);
-        assert_eq!(2000, test_cache.get(key).unwrap());
-        assert!(test_cache.get("unknown_key").is_none());
-    }
+		match reset {
+			Some(value) => {
+				if *value < SystemTime::now().duration_since(UNIX_EPOCH).expect("Time went backwards").as_millis() {
+					return Blocked { blocked: false, reset: 0 };
+				}
+				Blocked {
+					blocked: true,
+					reset: value.to_owned(),
+				}
+			}
+			None => Blocked { blocked: false, reset: 0 },
+		}
+	}
+	pub fn block_until(&mut self, identifier: &str, reset: u128) {
+		self.cache.insert(identifier.to_owned(), reset);
+	}
+	pub fn set(&mut self, identifier: &str, reset: u128) {
+		self.cache.insert(identifier.to_owned(), reset);
+	}
+	pub fn get(&self, identifier: &str) -> Option<u128> {
+		self.cache.get(identifier).copied()
+	}
+	pub fn incr(&mut self, identifier: &str) {
+		let mut value = self.get(identifier).unwrap_or(0);
+		value += 1;
+		self.cache.insert(identifier.to_owned(), value);
+	}
 }

src/ratelimit/common.rs

@@ -1,7 +1,7 @@
 #[derive(Debug)]
 pub struct RatelimitResponse {
-    pub success: bool,
-    pub limit: u32,
-    pub remaining: u32,
-    pub reset: u128,
+	pub success: bool,
+	pub limit: u32,
+	pub remaining: u32,
+	pub reset: u128,
 }