challenge datatype

Cargo.lock

@@ -303,7 +303,7 @@ dependencies = [
 [[package]]
 name = "libmcaptcha"
 version = "0.1.4"
-source = "git+https://github.com/mCaptcha/libmcaptcha?branch=master#06df4d24d04705fae38bcd80cbf124010bbcb0b7"
+source = "git+https://github.com/mCaptcha/libmcaptcha?branch=master#05dd9bbbb6efa1d47b87b991096d7b4778a14fe1"
 dependencies = [
  "derive_builder",
  "derive_more",

src/challenge.rs

@@ -0,0 +1,163 @@
+/*
+ * 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/>.
+ */
+use std::time::Duration;
+
+use redis_module::native_types::RedisType;
+use redis_module::raw::KeyType;
+use redis_module::NextArg;
+use redis_module::RedisResult;
+use redis_module::REDIS_OK;
+use redis_module::{raw, Context};
+use serde::{Deserialize, Serialize};
+
+use crate::bucket::Format;
+use crate::errors::*;
+use crate::utils::*;
+
+const MCAPTCHA_CHALLENGE_VERSION: i32 = 0;
+
+#[derive(Serialize, Deserialize)]
+pub struct Challenge {
+    difficulty: usize,
+    duration: u64,
+}
+
+#[derive(Serialize, Deserialize)]
+pub struct AddChallenge {
+    difficulty: usize,
+    duration: u64,
+    challenge: String,
+}
+
+impl Challenge {
+    pub fn new(duration: u64, difficulty: usize) -> Self {
+        Self {
+            difficulty,
+            duration,
+        }
+    }
+
+    pub fn create_challenge(ctx: &Context, args: Vec<String>) -> RedisResult {
+        let mut args = args.into_iter().skip(1);
+        let captcha = args.next_string()?;
+        let json = args.next_string()?;
+        let add_challenge: AddChallenge = Format::JSON.from_str(&json)?;
+
+        let challenge_name = get_challenge_name(&captcha, &add_challenge.challenge);
+
+        let key = ctx.open_key_writable(&challenge_name);
+        if key.key_type() != KeyType::Empty {
+            return Err(CacheError::DuplicateChallenge.into());
+        }
+        let challenge = Self::new(add_challenge.duration, add_challenge.difficulty);
+
+        key.set_value(&MCAPTCHA_CHALLENGE_TYPE, challenge)?;
+        key.set_expire(Duration::from_secs(add_challenge.duration))?;
+
+        REDIS_OK
+    }
+
+    pub fn get_challenge(ctx: &Context, args: Vec<String>) -> RedisResult {
+        let mut args = args.into_iter().skip(1);
+        let captcha = args.next_string()?;
+        let challenge = args.next_string()?;
+
+        let challenge_name = get_challenge_name(&captcha, &challenge);
+
+        let key = ctx.open_key_writable(&challenge_name);
+        if key.key_type() == KeyType::Empty {
+            return Err(CacheError::ChallengeNotFound.into());
+        }
+        match key.get_value::<Self>(&MCAPTCHA_CHALLENGE_TYPE)? {
+            Some(challenge) => {
+                let resp = serde_json::to_string(&challenge)?;
+                key.delete()?;
+                Ok(resp.into())
+            }
+            None => Err(CacheError::ChallengeNotFound.into()),
+        }
+    }
+}
+
+pub static MCAPTCHA_CHALLENGE_TYPE: RedisType = RedisType::new(
+    "mcaptchal",
+    MCAPTCHA_CHALLENGE_VERSION,
+    raw::RedisModuleTypeMethods {
+        version: raw::REDISMODULE_TYPE_METHOD_VERSION as u64,
+        rdb_load: Some(type_methods::rdb_load),
+        rdb_save: Some(type_methods::rdb_save),
+        aof_rewrite: None,
+        free: Some(type_methods::free),
+
+        // Currently unused by Redis
+        mem_usage: None,
+        digest: None,
+
+        // Aux data
+        aux_load: None,
+        aux_save: None,
+        aux_save_triggers: 0,
+
+        free_effort: None,
+        unlink: None,
+        copy: None,
+        defrag: None,
+    },
+);
+
+pub mod type_methods {
+    use std::os::raw::c_void;
+
+    use libc::c_int;
+
+    use super::*;
+
+    #[allow(non_snake_case, unused)]
+    pub extern "C" fn rdb_load(rdb: *mut raw::RedisModuleIO, encver: c_int) -> *mut c_void {
+        let challenge = match encver {
+            0 => {
+                let data = raw::load_string(rdb);
+                let challenge: Result<Challenge, CacheError> = Format::JSON.from_str(&data);
+                if challenge.is_err() {
+                    panic!(
+                        "Can't load Challenge from old redis RDB, error while serde {}, data received: {}",
+                        challenge.err().unwrap(),
+                        data
+                    );
+                }
+                challenge.unwrap()
+            }
+            _ => panic!("Can't load mCaptcha from old redis RDB, encver {}", encver),
+        };
+
+        Box::into_raw(Box::new(challenge)) as *mut c_void
+    }
+
+    pub unsafe extern "C" fn free(value: *mut c_void) {
+        let val = value as *mut Challenge;
+        Box::from_raw(val);
+    }
+
+    #[allow(non_snake_case, unused)]
+    pub unsafe extern "C" fn rdb_save(rdb: *mut raw::RedisModuleIO, value: *mut c_void) {
+        let challenge = &*(value as *mut Challenge);
+        match &serde_json::to_string(challenge) {
+            Ok(string) => raw::save_string(rdb, &string),
+            Err(e) => panic!("error while rdb_save: {}", e),
+        }
+    }
+}

src/errors.rs

@@ -30,6 +30,10 @@ pub enum CacheError {
     RedisError(redis_module::RedisError),
     #[display(fmt = "Captcha not found")]
     CaptchaNotFound,
+    #[display(fmt = "Challenge not found")]
+    ChallengeNotFound,
+    #[display(fmt = "Challenge already exists")]
+    DuplicateChallenge,
 }
 
 impl CacheError {
@@ -87,6 +91,8 @@ impl From<CacheError> for RedisError {
             CacheError::Msg(val) => RedisError::String(val),
             CacheError::RedisError(val) => val,
             CacheError::CaptchaNotFound => RedisError::String(format!("{}", e)),
+            CacheError::ChallengeNotFound => RedisError::String(format!("{}", e)),
+            CacheError::DuplicateChallenge => RedisError::String(format!("{}", e)),
         }
     }
 }

src/lib.rs

@@ -22,13 +22,14 @@ use redis_module::{NextArg, RedisResult};
 use redis_module::Context;
 
 mod bucket;
+mod challenge;
 mod errors;
 mod mcaptcha;
-#[allow(dead_code, unused_features, unused_variables)]
 mod safety;
 mod utils;
 
 use bucket::MCAPTCHA_BUCKET_TYPE;
+use challenge::MCAPTCHA_CHALLENGE_TYPE;
 use mcaptcha::MCAPTCHA_MCAPTCHA_TYPE;
 use safety::MCAPTCHA_SAFETY_TYPE;
 
@@ -64,6 +65,8 @@ lazy_static! {
     pub static ref PREFIX_CAPTCHA: String = format!("{}:captcha::", PKG_NAME);
     /// bucket key prefix
     pub static ref PREFIX_BUCKET: String = format!("{}:bucket:{{{}}}:", PKG_NAME, *ID);
+
+    pub static ref PREFIX_CHALLENGE: String = format!("{}:CHALLENGE", PKG_NAME);
 }
 
 pub fn on_delete(ctx: &Context, event_type: NotifyEvent, event: &str, key_name: &str) {
@@ -84,13 +87,15 @@ pub fn on_delete(ctx: &Context, event_type: NotifyEvent, event: &str, key_name:
 redis_module! {
     name: "mcaptcha_cahce",
     version: PKG_VERSION,
-    data_types: [MCAPTCHA_BUCKET_TYPE, MCAPTCHA_MCAPTCHA_TYPE, MCAPTCHA_SAFETY_TYPE],
+    data_types: [MCAPTCHA_BUCKET_TYPE, MCAPTCHA_MCAPTCHA_TYPE, MCAPTCHA_SAFETY_TYPE, MCAPTCHA_CHALLENGE_TYPE],
     commands: [
         ["mcaptcha_cache.add_visitor", bucket::Bucket::counter_create, "write", 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],
         ["mcaptcha_cache.delete_captcha", mcaptcha::MCaptcha::delete_captcha, "write", 1, 1, 1],
         ["mcaptcha_cache.captcha_exists", mcaptcha::MCaptcha::captcha_exists, "readonly", 1, 1, 1],
+        ["mcaptcha_cache.add_challenge", challenge::Challenge::create_challenge, "write", 1, 1, 1],
+        ["mcaptcha_cache.get_challenge", challenge::Challenge::get_challenge, "write", 1, 1, 1],
     ],
    event_handlers: [
         [@EXPIRED @EVICTED: on_delete],

src/mcaptcha.rs

@@ -1,6 +1,3 @@
-use redis_module::key::RedisKey;
-use redis_module::RedisError;
-use redis_module::RedisValue;
 /*
  * Copyright (C) 2021  Aravinth Manivannan <realaravinth@batsense.net>
  *
@@ -18,9 +15,12 @@ use redis_module::RedisValue;
  * along with this program.  If not, see <https://www.gnu.org/licenses/>.
  */
 use libmcaptcha::dev::{AddVisitorResult, CreateMCaptcha, DefenseBuilder, MCaptchaBuilder};
+use redis_module::key::RedisKey;
 use redis_module::key::RedisKeyWritable;
 use redis_module::native_types::RedisType;
 use redis_module::raw::KeyType;
+use redis_module::RedisError;
+use redis_module::RedisValue;
 use redis_module::{Context, RedisResult};
 use redis_module::{NextArg, REDIS_OK};
 //use redis_module::RedisError;

src/safety.rs

@@ -34,7 +34,7 @@ const MCAPTCHA_SAFETY_VERSION: i32 = 0;
 pub struct MCaptchaSafety;
 
 impl MCaptchaSafety {
-    pub fn on_delete(ctx: &Context, event_type: NotifyEvent, event: &str, key_name: &str) {
+    pub fn on_delete(ctx: &Context, _event_type: NotifyEvent, _event: &str, key_name: &str) {
         if !is_mcaptcha_safety(key_name) {
             return;
         }
@@ -143,7 +143,7 @@ impl MCaptchaSafety {
                     return;
                 }
 
-                if let Ok(Some(val)) = MCaptcha::get_mcaptcha(&mcaptcha) {
+                if let Ok(Some(_)) = MCaptcha::get_mcaptcha(&mcaptcha) {
                     let res = Self::new(ctx, duration, mcaptcha_name);
                     if res.is_err() {
                         ctx.log_warning(&format!(
@@ -159,7 +159,7 @@ impl MCaptchaSafety {
 }
 
 pub static MCAPTCHA_SAFETY_TYPE: RedisType = RedisType::new(
-    "mcaptdafe",
+    "mcaptsafe",
     MCAPTCHA_SAFETY_VERSION,
     raw::RedisModuleTypeMethods {
         version: raw::REDISMODULE_TYPE_METHOD_VERSION as u64,

src/utils.rs

@@ -73,6 +73,11 @@ pub fn is_mcaptcha_safety(name: &str) -> bool {
     name.contains(&PREFIX_SAFETY)
 }
 
+#[inline]
+pub fn get_challenge_name(captcha: &str, challenge: &str) -> String {
+    format!("{}:{{{}}}:{}", &*PREFIX_CHALLENGE, captcha, challenge)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

tests/challenge.py

@@ -0,0 +1,131 @@
+#!/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/>.
+from asyncio import sleep
+import json
+
+import redis
+
+import utils
+
+r = utils.connect()
+utils.ping(r)
+
+# 1. Check duplicate challenge
+# 2. Create challenge
+# 3. Read non-existent challenge
+# 4. Read challenge
+# 5. Read expired challenge
+
+
+COMMANDS = {
+ "ADD" :"MCAPTCHA_CACHE.ADD_CHALLENGE",
+ "GET" :"MCAPTCHA_CACHE.GET_CHALLENGE"
+}
+
+CHALLENGE_NOT_FOUND = "Challenge not found"
+DUPLICATE_CHALLENGE = "Challenge already exists"
+
+def add_challenge(captcha, challenge):
+    """Add challenge to Redis"""
+    try :
+        return r.execute_command(COMMANDS["ADD"], captcha, challenge)
+    except Exception as e:
+        return e
+    
+def get_challenge_from_redis(captcha, challenge):
+    """Add challenge to Redis"""
+    try :
+        data = r.execute_command(COMMANDS["GET"], captcha, challenge)
+        return json.loads(data)
+    except Exception as e:
+        return e
+
+def get_challenge(challenge):
+    """Get challenge JSON"""
+    challenge = {
+        "difficulty": 500,
+        "duration": 5,
+        "challenge": challenge,
+    }
+    return json.dumps(challenge)
+
+
+async def add_challenge_works():
+    """Test: Add Challenge"""
+    try:
+        key = "add_challenge"
+        challenge_name = key
+        challenge = get_challenge(challenge_name)
+
+        add_challenge(key, challenge)
+        stored_challenge = get_challenge_from_redis(key, challenge_name)
+        challenge_dict = json.loads(challenge)
+        assert stored_challenge["difficulty"] == challenge_dict["difficulty"]
+        assert stored_challenge["duration"] == challenge_dict["duration"]
+        error = get_challenge_from_redis(key, challenge_name)
+        assert str(error) == CHALLENGE_NOT_FOUND
+        print("[*] Add Challenge works")
+
+    except Exception as e:
+        raise e
+
+async def challenge_ttl_works():
+    """Test: Challenge TTL"""
+    try:
+        key = "ttl_challenge"
+        challenge_name = key
+        challenge = get_challenge(challenge_name)
+
+        add_challenge(key, challenge)
+        await sleep(5 + 2)
+
+        error = get_challenge_from_redis(key, challenge_name)
+        assert str(error) == CHALLENGE_NOT_FOUND
+
+        print("[*] Challenge TTL works")
+    except Exception as e:
+        raise e
+
+
+async def challenge_doesnt_exist():
+    """Test: Non-existent Challenge"""
+    try:
+        challenge_name = "nonexistent_challenge"
+        key  = challenge_name
+
+        error = get_challenge_from_redis(key, challenge_name)
+        assert str(error) == CHALLENGE_NOT_FOUND
+
+        print("[*] Challenge Doesn't Exist works")
+    except Exception as e:
+        raise e
+
+
+async def duplicate_challenge_works():
+    """Test: Duplicate Challenges"""
+    try:
+        challenge_name = "nonexistent_challenge"
+        key  = challenge_name
+        challenge = get_challenge(challenge_name)
+
+        add_challenge(key, challenge)
+        error = add_challenge(key, challenge)
+        assert str(error) == DUPLICATE_CHALLENGE
+
+        print("[*] Duplicate Challenge works")
+    except Exception as e:
+        raise e

tests/runner.py

@@ -18,6 +18,7 @@ import asyncio
 
 import bucket
 import mcaptcha
+import challenge
 
 
 class Runner(object):
@@ -27,7 +28,11 @@ class Runner(object):
         bucket.difficulty_works,
         mcaptcha.delete_captcha_works,
         mcaptcha.captcha_exists_works,
-        mcaptcha.register_captcha_works
+        mcaptcha.register_captcha_works,
+        challenge.add_challenge_works,
+        challenge.challenge_doesnt_exist,
+        challenge.challenge_ttl_works,
+        challenge.duplicate_challenge_works,
     ]
     __tasks = []