MCaptcha: increment, decrement & rdb persist

Cargo.lock

@@ -303,7 +303,7 @@ dependencies = [
 [[package]]
 name = "libmcaptcha"
 version = "0.1.4"
-source = "git+https://github.com/mCaptcha/libmcaptcha?branch=master#895d54d26d7325f83df963794c815b00f19990d0"
+source = "git+https://github.com/mCaptcha/libmcaptcha?branch=master#68f95f99c28753a7725cd4107078978477ed2f63"
 dependencies = [
  "derive_builder",
  "derive_more",
@@ -324,6 +324,7 @@ dependencies = [
 name = "mcaptcha-cache"
 version = "0.1.0"
 dependencies = [
+ "derive_more",
  "lazy_static",
  "libc",
  "libmcaptcha",

Cargo.toml

@@ -19,6 +19,7 @@ serde_json = "1.0.64"
 serde = {version = "1.0.126", features = ["derive"]}
 lazy_static = "1.4.0"
 rand = "0.8.3"
+derive_more = "0.99"
 libmcaptcha = { branch = "master", git = "https://github.com/mCaptcha/libmcaptcha", features = ["minimal"], default-features = false }
 #libmcaptcha = { path = "../libmcaptcha", features = ["minimal"], default-features = false}
 

docs/mechanism.md

@@ -0,0 +1,41 @@
+# mCaptcha Redis Module Mechanism
+
+## Data types
+
+1. Bucket
+2. Bucket safety
+3. mCaptcha
+4. mCaptcha safety
+
+## Bucket
+
+- Timer queue, used for scheduling decrements.
+- Timer used for scheduling can't be persisted so requires a safety
+
+## Bucket Safety
+
+- Has expiry timer(Redis `EXIPRE`)
+- Key name includes bucket name for which this safety was created for
+- When Redis is started after crash, bucket safety expires and expire
+  event is fired, the corresponding bucket is fetched and executed.
+
+## mCaptcha
+
+- Contains mCaptcha defense details and current state
+- When redis performs resharding, it's possible that mCaptcha gets
+  separated from its bucket.
+- We can't pin mCaptcha by using hashtags because the client figures out
+  where a key should be placed/is available and will have no
+  knowledge about node IDs that we use for pinning.
+
+- So this too requires a safety to make sure that when recovering from a
+  crash, it's counter doesn't have residues permanently.
+
+## mCaptcha safety
+
+- Has expire timer(Redis `EXIPRE`)
+- Key name includes mCaptcha name for which this safety was created for
+- when expiry event is fired for this type, a bucket for corresponding
+  mCaptcha is created to decrement it by `x`(where `x` is the count of mCaptcha
+  when this event is fired). It's not perfect but at least we get
+  eventual consistency.

src/bucket.rs

@@ -28,6 +28,7 @@ use redis_module::{raw, Context};
 use serde::{Deserialize, Serialize};
 
 use crate::errors::*;
+use crate::mcaptcha::MCaptcha;
 use crate::utils::*;
 use crate::*;
 
@@ -58,7 +59,7 @@ pub struct Bucket {
     /// instant(seconds from UNIX_EPOCH) at which time bucket begins decrement process
     bucket_instant: u64,
     /// a list of captcha keys that should be decremented during clean up
-    decrement: HashMap<String, usize>,
+    decrement: HashMap<String, u32>,
 }
 
 impl Bucket {
@@ -109,6 +110,54 @@ impl Bucket {
         Ok(bucket)
     }
 
+    /// decrement runner that decrements all registered counts _without_ cleaning after itself
+    /// use [decrement] when you require auto cleanup. Internally, it calls this method.
+    #[inline]
+    fn decrement_runner(ctx: &Context, key: &RedisKeyWritable) {
+        let val = key.get_value::<Bucket>(&MCAPTCHA_BUCKET_TYPE).unwrap();
+        match val {
+            Some(bucket) => {
+                ctx.log_debug(&format!("entering loop hashmap "));
+                for (captcha, count) in bucket.decrement.drain() {
+                    ctx.log_debug(&format!(
+                        "reading captcha: {} with decr count {}",
+                        &captcha, count
+                    ));
+                    let stored_captcha = ctx.open_key_writable(&captcha);
+                    if stored_captcha.key_type() == KeyType::Empty {
+                        continue;
+                    }
+                    let captcha = MCaptcha::get_mcaptcha(&ctx, &stored_captcha)
+                        .unwrap()
+                        .unwrap();
+                    captcha.decrement_visitor_by(count);
+                }
+            }
+            None => {
+                ctx.log_debug(&format!("bucket not found, can't decrement"));
+            }
+        }
+    }
+
+    /// executes when timer goes off. Decrements all registered counts and cleans itself up
+    fn decrement(ctx: &Context, bucket_instant: u64) {
+        // get  bucket
+        let bucket_name = get_bucket_name(bucket_instant);
+
+        let timer = ctx.open_key_writable(&get_timer_name_from_bucket_name(&bucket_name));
+        let _ = timer.delete();
+
+        ctx.log_debug(&format!("Bucket instant: {}", &bucket_instant));
+
+        let bucket = ctx.open_key_writable(&bucket_name);
+        Bucket::decrement_runner(ctx, &bucket);
+
+        match bucket.delete() {
+            Err(e) => ctx.log_warning(&format!("enountered error while deleting hashmap: {:?}", e)),
+            Ok(_) => (),
+        }
+    }
+
     /// increments count of key = captcha and registers for auto decrement
     #[inline]
     fn increment(ctx: &Context, duration: u64, captcha: &str) -> CacheResult<()> {
@@ -116,20 +165,11 @@ impl Bucket {
         ctx.log_debug(&captcha_name);
         // increment
         let captcha = ctx.open_key_writable(&captcha_name);
+        let captcha = MCaptcha::get_mcaptcha(ctx, &captcha)?;
 
-        match captcha.read()? {
+        match captcha {
-            Some(val) => {
+            Some(val) => val.add_visitor(),
-                if val.trim().is_empty() {
+            None => return Err(CacheError::new("Captcha not found".into())),
-                    captcha.write("1")?;
-                } else {
-                    let mut val: usize = val.parse()?;
-                    val += 1;
-                    captcha.write(&val.to_string())?;
-                }
-            }
-            None => {
-                captcha.write("1")?;
-            }
         }
 
         let bucket_instant = get_bucket_instant(duration)?;
@@ -161,65 +201,6 @@ impl Bucket {
         Ok(())
     }
 
-    /// decrement runner that decrements all registered counts _without_ cleaning after itself
-    /// use [decrement] when you require auto cleanup. Internally, it calls this method.
-    #[inline]
-    fn decrement_runner(ctx: &Context, key: &RedisKeyWritable) {
-        let val = key.get_value::<Bucket>(&MCAPTCHA_BUCKET_TYPE).unwrap();
-        match val {
-            Some(bucket) => {
-                ctx.log_debug(&format!("entering loop hashmap "));
-                for (captcha, count) in bucket.decrement.drain() {
-                    ctx.log_debug(&format!(
-                        "reading captcha: {} with decr count {}",
-                        &captcha, count
-                    ));
-                    let stored_captcha = ctx.open_key_writable(&captcha);
-                    if stored_captcha.key_type() == KeyType::Empty {
-                        continue;
-                    }
-
-                    let mut stored_count: usize =
-                        stored_captcha.read().unwrap().unwrap().parse().unwrap();
-                    stored_count -= count;
-                    if stored_count == 0 {
-                        match stored_captcha.delete() {
-                            Err(e) => ctx.log_warning(&format!(
-                                "Error occured while cleaning up captcha when it became 0: {}",
-                                e
-                            )),
-                            Ok(_) => (),
-                        }
-                    } else {
-                        stored_captcha.write(&stored_count.to_string()).unwrap();
-                    }
-                }
-            }
-            None => {
-                ctx.log_debug(&format!("bucket not found, can't decrement"));
-            }
-        }
-    }
-
-    /// executes when timer goes off. Decrements all registered counts and cleans itself up
-    fn decrement(ctx: &Context, bucket_instant: u64) {
-        // get  bucket
-        let bucket_name = get_bucket_name(bucket_instant);
-
-        let timer = ctx.open_key_writable(&get_timer_name_from_bucket_name(&bucket_name));
-        let _ = timer.delete();
-
-        ctx.log_debug(&format!("Bucket instant: {}", &bucket_instant));
-
-        let bucket = ctx.open_key_writable(&bucket_name);
-        Bucket::decrement_runner(ctx, &bucket);
-
-        match bucket.delete() {
-            Err(e) => ctx.log_warning(&format!("enountered error while deleting hashmap: {:?}", e)),
-            Ok(_) => (),
-        }
-    }
-
     /// Create new counter
     pub fn counter_create(ctx: &Context, args: Vec<String>) -> RedisResult {
         let mut args = args.into_iter().skip(1);
@@ -270,11 +251,10 @@ pub mod type_methods {
         let bucket = match encver {
             0 => {
                 let data = raw::load_string(rdb);
-                let fmt = Format::JSON;
+                let bucket: Bucket = Format::JSON.from_str(&data).unwrap();
-                let bucket: Bucket = fmt.from_str(&data).unwrap();
                 bucket
             }
-            _ => panic!("Can't load old RedisJSON RDB"),
+            _ => panic!("Can't load bucket from old redis RDB"),
         };
 
         //        if bucket.

src/errors.rs

@@ -17,13 +17,18 @@
 
 use std::num::ParseIntError;
 
+use derive_more::Display;
 use redis_module::RedisError;
 use redis_module::RedisResult;
 
-#[derive(Debug)]
+#[derive(Debug, Display)]
 pub enum CacheError {
+    #[display(fmt = "{}", &_0)]
     Msg(String),
+    #[display(fmt = "{}", &_0.to_string)]
     RedisError(redis_module::RedisError),
+    #[display(fmt = "Captcha not found")]
+    CaptchaNotFound,
 }
 
 impl CacheError {
@@ -50,21 +55,12 @@ impl From<serde_json::Error> for CacheError {
     }
 }
 
-impl From<redis_module::RedisError> for CacheError {
+impl From<RedisError> for CacheError {
     fn from(e: redis_module::RedisError) -> Self {
         CacheError::RedisError(e)
     }
 }
 
-impl From<CacheError> for RedisError {
-    fn from(e: CacheError) -> Self {
-        match e {
-            CacheError::Msg(val) => RedisError::String(val),
-            CacheError::RedisError(val) => val,
-        }
-    }
-}
-
 impl From<ParseIntError> for CacheError {
     fn from(e: ParseIntError) -> Self {
         let err: RedisError = e.into();
@@ -73,10 +69,17 @@ impl From<ParseIntError> for CacheError {
 }
 
 impl From<CacheError> for RedisResult {
+    fn from(e: CacheError) -> Self {
+        Self::Err(e.into())
+    }
+}
+
+impl From<CacheError> for RedisError {
     fn from(e: CacheError) -> Self {
         match e {
-            CacheError::Msg(val) => Err(RedisError::String(val)),
+            CacheError::Msg(val) => RedisError::String(val),
-            CacheError::RedisError(val) => Err(val),
+            CacheError::RedisError(val) => val,
+            CacheError::CaptchaNotFound => RedisError::String(format!("{}", e)),
         }
     }
 }

src/lib.rs

@@ -56,7 +56,7 @@ lazy_static! {
         rng.gen()
     };
     /// counter/captcha key prefix
-    pub static ref PREFIX_COUNTER: String = format!("{}:captcha:{}:", PKG_NAME, *ID);
+    pub static ref PREFIX_CAPTCHA: String = format!("{}:captcha::", PKG_NAME);
     /// bucket key prefix
     pub static ref PREFIX_BUCKET: String = format!("{}:bucket:{{{}}}:", PKG_NAME, *ID);
 }
@@ -66,8 +66,9 @@ redis_module! {
     version: PKG_VERSION,
     data_types: [MCAPTCHA_BUCKET_TYPE, MCAPTCHA_MCAPTCHA_TYPE],
     commands: [
-        ["mcaptcha_cache.count", bucket::Bucket::counter_create, "write", 1, 1, 1],
+        ["mcaptcha_cache.add_visitor", bucket::Bucket::counter_create, "write", 1, 1, 1],
-        ["mcaptcha_cache.get", mcaptcha::MCaptcha::get, "readonly", 1, 1, 1],
+        ["mcaptcha_cache.get", mcaptcha::MCaptcha::get_count, "readonly", 1, 1, 1],
+        ["mcaptcha_cache.add_captcha", mcaptcha::MCaptcha::add_captcha, "readonly", 1, 1, 1],
     ],
    event_handlers: [
         [@EXPIRED @EVICTED: bucket::Bucket::on_delete],

src/mcaptcha.rs

@@ -1,3 +1,4 @@
+use redis_module::RedisValue;
 /*
  * Copyright (C) 2021  Aravinth Manivannan <realaravinth@batsense.net>
  *
@@ -14,10 +15,11 @@
  * You should have received a copy of the GNU Affero General Public License
  * along with this program.  If not, see <https://www.gnu.org/licenses/>.
  */
-//use redis_module::key::RedisKeyWritable;
+use redis_module::key::RedisKeyWritable;
 use redis_module::native_types::RedisType;
 use redis_module::raw::KeyType;
 use redis_module::{Context, RedisResult};
+use redis_module::{NextArg, REDIS_OK};
 //use redis_module::RedisError;
 use redis_module::raw;
 
@@ -25,9 +27,9 @@ use serde::{Deserialize, Serialize};
 
 use crate::bucket::Format;
 use crate::errors::*;
-use crate::utils;
+use crate::utils::*;
 
-const REDIS_MCPATCHA_MCAPTCHA_TYPE_VERSION: i32 = 1;
+const REDIS_MCPATCHA_MCAPTCHA_TYPE_VERSION: i32 = 0;
 
 #[derive(Serialize, Deserialize)]
 pub struct MCaptcha {
@@ -35,6 +37,10 @@ pub struct MCaptcha {
 }
 
 impl MCaptcha {
+    #[inline]
+    fn new(m: libmcaptcha::MCaptcha) -> Self {
+        MCaptcha { m }
+    }
     /// increments the visitor count by one
     #[inline]
     pub fn add_visitor(&mut self) {
@@ -53,32 +59,61 @@ impl MCaptcha {
         self.m.get_difficulty()
     }
 
-    /// get [Counter]'s lifetime
+    /// get [MCaptcha]'s lifetime
     #[inline]
     pub fn get_duration(&self) -> u64 {
         self.m.get_duration()
     }
 
-    /// get [Counter]'s current visitor_threshold
+    /// get [MCaptcha]'s current visitor_threshold
     #[inline]
     pub fn get_visitors(&self) -> u32 {
         self.m.get_visitors()
     }
 
-    /// Get counter value
+    /// decrement [MCaptcha]'s current visitor_threshold by specified count
-    pub fn get(ctx: &Context, args: Vec<String>) -> RedisResult {
+    #[inline]
-        use redis_module::NextArg;
+    pub fn decrement_visitor_by(&mut self, count: u32) {
+        self.m.decrement_visitor_by(count)
+    }
 
+    /// get mcaptcha from redis key writable
+    pub fn get_mcaptcha<'a>(
+        ctx: &Context,
+        key: &'a RedisKeyWritable,
+    ) -> CacheResult<Option<&'a mut Self>> {
+        Ok(key.get_value::<Self>(&MCAPTCHA_MCAPTCHA_TYPE)?)
+    }
+
+    /// Get counter value
+    pub fn get_count(ctx: &Context, args: Vec<String>) -> RedisResult {
         let mut args = args.into_iter().skip(1);
         let key_name = args.next_string()?;
-        let key_name = utils::get_captcha_key(&key_name);
+        let key_name = get_captcha_key(&key_name);
 
         let stored_captcha = ctx.open_key(&key_name);
         if stored_captcha.key_type() == KeyType::Empty {
             return CacheError::new(format!("key {} not found", key_name)).into();
         }
 
-        Ok(stored_captcha.read()?.unwrap().into())
+        match stored_captcha.get_value::<Self>(&MCAPTCHA_MCAPTCHA_TYPE)? {
+            Some(val) => Ok(RedisValue::Integer(val.get_visitors().into())),
+            None => return Err(CacheError::CaptchaNotFound.into()),
+        }
+    }
+
+    /// Add captcha to redis
+    pub fn add_captcha(ctx: &Context, args: Vec<String>) -> RedisResult {
+        let mut args = args.into_iter().skip(1);
+        let key_name = get_captcha_key(&args.next_string()?);
+        let json = args.next_string()?;
+        let mcaptcha: libmcaptcha::MCaptcha = Format::JSON.from_str(&json)?;
+        let mcaptcha = Self::new(mcaptcha);
+
+        let key = ctx.open_key_writable(&&key_name);
+        key.set_value(&MCAPTCHA_MCAPTCHA_TYPE, mcaptcha)?;
+
+        REDIS_OK
     }
 }
 
@@ -120,12 +155,10 @@ pub mod type_methods {
         let mcaptcha = match encver {
             0 => {
                 let data = raw::load_string(rdb);
-
+                let mcaptcha: MCaptcha = Format::JSON.from_str(&data).unwrap();
-                let fmt = Format::JSON;
-                let mcaptcha: MCaptcha = fmt.from_str(&data).unwrap();
                 mcaptcha
             }
-            _ => panic!("Can't load old RedisJSON RDB"),
+            _ => panic!("Can't load mCaptcha from old redis RDB"),
         };
 
         Box::into_raw(Box::new(mcaptcha)) as *mut c_void
@@ -141,7 +174,7 @@ pub mod type_methods {
         let mcaptcha = &*(value as *mut MCaptcha);
         match &serde_json::to_string(mcaptcha) {
             Ok(string) => raw::save_string(rdb, &string),
-            Err(e) => eprintln!("error while rdb_save: {}", e),
+            Err(e) => panic!("error while rdb_save: {}", e),
         }
     }
 }

src/utils.rs

@@ -50,7 +50,7 @@ pub fn get_bucket_instant(duration: u64) -> CacheResult<u64> {
 
 #[inline]
 pub fn get_captcha_key(name: &str) -> String {
-    format!("{}{}", &*PREFIX_COUNTER, name)
+    format!("{}{}", &*PREFIX_CAPTCHA, name)
 }
 
 #[inline]

tests/bucket.py

@@ -1,6 +1,5 @@
-#!/bin/env python3 
+#!/bin/env /usr/bin/python3
-#
+# # Copyright (C) 2021  Aravinth Manivannan <realaravinth@batsense.net>
-# Copyright (C) 2021  Aravinth Manivannan <realaravinth@batsense.net>
 # 
 # This program is free software: you can redistribute it and/or modify
 # it under the terms of the GNU Affero General Public License as
@@ -15,13 +14,15 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 from time import sleep
+import sys
 
-import test;
+import test
+from mcaptcha import register
 
 r = test.r
 
 COMMANDS = {
-"COUNT" : "mcaptcha_cache.count",
+"COUNT" : "mcaptcha_cache.add_visitor",
 "GET" : "mcaptcha_cache.get",
 }
 
@@ -40,27 +41,36 @@ def assert_count(expect, key):
     assert count == expect
 
 def incr_one_works():
-    key = "incr_one"
+    try:
-    time = 2
+        key = "incr_one"
-    initial_count = get_count(key)
+        register(r, key)
-    # incriment
+        time = 2
-    incr(key, time)
+        initial_count = get_count(key)
-    assert_count(initial_count + 1, key)
+        # incriment
-    # wait till expiry
+        incr(key, time)
-    sleep(time + 2)
+        assert_count(initial_count + 1, key)
-    assert_count(initial_count, key)
+        # wait till expiry
-    print("Incr one works")
+        sleep(time + 2)
+        assert_count(initial_count, key)
+        print("Incr one works")
+    except Exception as e:
+        raise e
+
 
 def race_works():
     key = "race_works"
-    initial_count = get_count(key)
+    try:
-    race_num = 200
+        register(r, key)
-    time = 3
+        initial_count = get_count(key)
+        race_num = 200
+        time = 3
 
-    for _ in range(race_num):
+        for _ in range(race_num):
-        incr(key, time)
+            incr(key, time)
-    assert_count(initial_count + race_num, key)
+        assert_count(initial_count + race_num, key)
-    # wait till expiry
+        # wait till expiry
-    sleep(time + 2)
+        sleep(time + 2)
-    assert_count(initial_count, key)
+        assert_count(initial_count, key)
-    print("Race works")
+        print("Race works")
+    except Exception as e:
+        raise e

tests/mcaptcha.py

@@ -0,0 +1,42 @@
+#!/bin/env /usr/bin/python3
+#
+# Copyright (C) 2021  Aravinth Manivannan <realaravinth@batsense.net>
+# 
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+# 
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+# 
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+import json
+
+MCAPTCHA = {
+  "visitor_threshold": 0,
+  "defense": {
+    "levels": [
+      {"visitor_threshold": 50, "difficulty_factor": 50},
+      {"visitor_threshold": 500, "difficulty_factor": 500}
+    ],
+    "current_visitor_threshold": 0
+  },
+  "duration": 5
+}
+
+COMMANDS = {
+    "ADD_CAPTCHA": "MCAPTCHA_CACHE.ADD_CAPTCHA",
+}
+
+payload = json.dumps(MCAPTCHA)
+
+def register(r, key):
+    if r.exists(key):
+        r.delete(key)
+
+    r.execute_command(COMMANDS["ADD_CAPTCHA"], key, payload)

tests/runner.py

@@ -1,3 +1,4 @@
+#!/bin/env /usr/bin/python3
 # Copyright (C) 2021  Aravinth Manivannan <realaravinth@batsense.net>
 # 
 # This program is free software: you can redistribute it and/or modify
@@ -22,13 +23,24 @@ class Runner(object):
     def register(self, fn):
         self._functions.append(fn)
         t = Thread(target=fn)
-        t.start()
         self._threads.append(t)
 
     """Wait for registered functions to finish executing"""
-    def wait(self):
+
+    def __run__(self):
         for thread in self._threads:
-            thread.join()
+            try:
+                thread.start()
+            except:
+                print("yo")
+
+    def wait(self):
+        self.__run__()
+        for thread in self._threads:
+            try:
+                thread.join()
+            except:
+                print("yo")
 
     """Runs in seperate threads"""
     def __init__(self):

tests/test.py

@@ -1,4 +1,4 @@
-#!/bin/env python3 
+#!/bin/env /usr/bin/python3
 #
 # Copyright (C) 2021  Aravinth Manivannan <realaravinth@batsense.net>
 # 
@@ -32,15 +32,20 @@ utils.ping(r)
 
 
 def main():
-    runner = Runner()
+    #runner = Runner()
-    fn = [bucket.incr_one_works, bucket.race_works]
+    #fn = [bucket.incr_one_works]#, bucket.race_works]
-    for r in fn:
-        runner.register(r)
 
-    runner.wait()
+    bucket.incr_one_works()
+    bucket.race_works()
 
-    print("All tests passed")
+    #try:
+    #    for r in fn:
+    #        runner.register(r)
 
+    #    runner.wait()
+    #    print("All tests passed")
+    #except Exception as e:
+    #    raise e
 
 if __name__ == "__main__":
     main()

tests/utils.py

@@ -1,5 +1,4 @@
-
+#!/bin/env /usr/bin/python3
-#!/bin/env python3 
 #
 # Copyright (C) 2021  Aravinth Manivannan <realaravinth@batsense.net>
 #