Ything-RepoMirrors/mcaptcha-cache
2
0
Fork 0
mirror of https://github.com/mCaptcha/cache.git synced 2026-01-25 04:16:29 +00:00
Code Issues Packages Projects Releases Wiki Activity

delete and check exists captcha

Browse Source
This commit is contained in:
realaravinth
2021-06-07 17:14:39 +05:30
parent 17c2fcacd9
commit b1cdb6f221
11 changed files with 327 additions and 119 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
Cargo.lock generated
View File

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

4
Cargo.toml
View File

@@ -20,8 +20,8 @@ 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}
#libmcaptcha = { branch = "master", git = "https://github.com/mCaptcha/libmcaptcha", features = ["minimal"], default-features = false }
libmcaptcha = { path = "../libmcaptcha", features = ["minimal"], default-features = false}
#[target.x86_64-unknown-linux-musl]
#linker = "cc"

39
src/bucket.rs
View File

@@ -32,6 +32,9 @@ use crate::mcaptcha::MCaptcha;
use crate::utils::*;
use crate::*;
/// Bucket type version, aka encoding version
const REDIS_MCAPTCHA_BUCKET_TYPE_VERSION: i32 = 0;
#[derive(Debug, PartialEq)]
/// encoding formats for persistence
pub enum Format {
@@ -114,9 +117,8 @@ impl Bucket {
/// 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) => {
match key.get_value::<Bucket>(&MCAPTCHA_BUCKET_TYPE) {
Ok(Some(bucket)) => {
ctx.log_debug(&format!("entering loop hashmap "));
for (captcha, count) in bucket.decrement.drain() {
ctx.log_debug(&format!(
@@ -127,13 +129,12 @@ impl Bucket {
if stored_captcha.key_type() == KeyType::Empty {
continue;
}
let captcha = MCaptcha::get_mut_mcaptcha(&stored_captcha)
.unwrap()
.unwrap();
captcha.decrement_visitor_by(count);
if let Ok(Some(captcha)) = MCaptcha::get_mut_mcaptcha(&stored_captcha) {
captcha.decrement_visitor_by(count);
}
}
}
None => {
_ => {
ctx.log_debug(&format!("bucket not found, can't decrement"));
}
}
@@ -160,18 +161,27 @@ impl Bucket {
/// increments count of key = captcha and registers for auto decrement
#[inline]
fn increment(ctx: &Context, duration: u64, captcha: &str) -> CacheResult<()> {
fn increment(ctx: &Context, captcha: &str) -> CacheResult<()> {
let captcha_name = get_captcha_key(captcha);
ctx.log_debug(&captcha_name);
// increment
let captcha = ctx.open_key_writable(&captcha_name);
ctx.log_debug("loading mcaptcha");
let captcha = MCaptcha::get_mut_mcaptcha(&captcha)?;
match captcha {
Some(val) => val.add_visitor(),
None => return Err(CacheError::new("Captcha not found".into())),
ctx.log_debug("loaded mcaptcha");
if captcha.is_none() {
return Err(CacheError::new("Captcha not found".into()));
}
let captcha = captcha.unwrap();
ctx.log_debug(&format!(
"current visitor count: {}",
captcha.get_visitors()
));
captcha.add_visitor();
ctx.log_debug("visitor added");
let duration = captcha.get_duration();
let bucket_instant = get_bucket_instant(duration)?;
let bucket_name = get_bucket_name(bucket_instant);
@@ -207,8 +217,7 @@ impl Bucket {
// mcaptcha captcha key name
let key_name = args.next_string()?;
// expiry
let duration = args.next_u64()?;
bucket::Bucket::increment(ctx, duration, &key_name)?;
bucket::Bucket::increment(ctx, &key_name)?;
REDIS_OK
}
}
@@ -254,7 +263,7 @@ pub mod type_methods {
let bucket: Bucket = Format::JSON.from_str(&data).unwrap();
bucket
}
_ => panic!("Can't load bucket from old redis RDB"),
_ => panic!("Can't load bucket from old redis RDB, encver: {}", encver,),
};
// if bucket.

7
src/errors.rs
View File

@@ -18,6 +18,7 @@
use std::num::ParseIntError;
use derive_more::Display;
use libmcaptcha::errors::CaptchaError;
use redis_module::RedisError;
use redis_module::RedisResult;
@@ -74,6 +75,12 @@ impl From<CacheError> for RedisResult {
}
}
impl From<CaptchaError> for CacheError {
fn from(e: CaptchaError) -> Self {
CacheError::Msg(format!("{}", e))
}
}
impl From<CacheError> for RedisError {
fn from(e: CacheError) -> Self {
match e {

11
src/lib.rs
View File

@@ -15,7 +15,7 @@
* along with this program. If not, see <https://www.gnu.org/licenses/>.
*/
use lazy_static::lazy_static;
use redis_module::{redis_command, redis_event_handler, redis_module};
use redis_module::{redis_command, redis_event_handler, redis_module, RedisValue};
use redis_module::{NextArg, RedisResult, REDIS_OK};
mod bucket;
@@ -32,11 +32,8 @@ use safety::MCAPTCHA_SAFETY_TYPE;
/// Initial allocation ammount of bucket[bucket::Bucket]
pub const HIT_PER_SECOND: usize = 100;
/// Bucket[bucket::Bucket] type version
pub const REDIS_MCAPTCHA_BUCKET_TYPE_VERSION: i32 = 1;
pub const PKG_NAME: &str = "mcap";
pub const PKG_VERSION: usize = 1;
pub const PKG_VERSION: usize = 0;
/// bucket timer key prefix
// PREFIX_BUCKET_TIMER is used like this:
@@ -52,6 +49,8 @@ pub const PREFIX_SAFETY: &str = "safety:";
pub const BUCKET_EXPIRY_OFFSET: u64 = 30;
lazy_static! {
/// node unique identifier, useful when running in cluster mode
pub static ref ID: usize = {
use rand::prelude::*;
@@ -72,6 +71,8 @@ redis_module! {
["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],
],
event_handlers: [
[@EXPIRED @EVICTED: bucket::Bucket::on_delete],

174
src/mcaptcha.rs
View File

@@ -1,4 +1,5 @@
use redis_module::key::RedisKey;
use redis_module::RedisError;
use redis_module::RedisValue;
/*
* Copyright (C) 2021 Aravinth Manivannan <realaravinth@batsense.net>
@@ -16,6 +17,7 @@ use redis_module::RedisValue;
* 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 libmcaptcha::{defense::Level, DefenseBuilder, MCaptchaBuilder};
use redis_module::key::RedisKeyWritable;
use redis_module::native_types::RedisType;
use redis_module::raw::KeyType;
@@ -28,6 +30,7 @@ use serde::{Deserialize, Serialize};
use crate::bucket::Format;
use crate::errors::*;
use crate::safety::MCaptchaSafety;
use crate::utils::*;
const REDIS_MCPATCHA_MCAPTCHA_TYPE_VERSION: i32 = 0;
@@ -37,11 +40,30 @@ pub struct MCaptcha {
m: libmcaptcha::MCaptcha,
}
#[derive(Serialize, Deserialize)]
pub struct CreateMCaptcha {
levels: Vec<Level>,
duration: u64,
}
impl MCaptcha {
#[inline]
fn new(m: libmcaptcha::MCaptcha) -> Self {
MCaptcha { m }
fn new(mut m: CreateMCaptcha) -> CacheResult<Self> {
let mut defense_builder = DefenseBuilder::default();
let mut defense_builder = &mut defense_builder;
for l in m.levels.drain(0..) {
defense_builder = defense_builder.add_level(l)?;
}
let defense = defense_builder.build()?;
let m = MCaptchaBuilder::default()
.defense(defense)
.duration(m.duration)
.build()?;
Ok(MCaptcha { m })
}
/// increments the visitor count by one
#[inline]
pub fn add_visitor(&mut self) {
@@ -50,24 +72,28 @@ impl MCaptcha {
/// decrements the visitor count by one
#[inline]
#[allow(dead_code)]
pub fn decrement_visitor(&mut self) {
self.m.decrement_visitor()
}
/// get current difficulty factor
#[inline]
#[allow(dead_code)]
pub fn get_difficulty(&self) -> u32 {
self.m.get_difficulty()
}
/// get [MCaptcha]'s lifetime
#[inline]
#[allow(dead_code)]
pub fn get_duration(&self) -> u64 {
self.m.get_duration()
}
/// get [MCaptcha]'s current visitor_threshold
#[inline]
#[allow(dead_code)]
pub fn get_visitors(&self) -> u32 {
self.m.get_visitors()
}
@@ -112,13 +138,50 @@ impl MCaptcha {
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 mcaptcha: CreateMCaptcha = Format::JSON.from_str(&json)?;
let duration = mcaptcha.duration;
let mcaptcha = Self::new(mcaptcha)?;
let key = ctx.open_key_writable(&&key_name);
key.set_value(&MCAPTCHA_MCAPTCHA_TYPE, mcaptcha)?;
let key = ctx.open_key_writable(&key_name);
if key.key_type() == KeyType::Empty {
key.set_value(&MCAPTCHA_MCAPTCHA_TYPE, mcaptcha)?;
ctx.log_debug(&format!("mcaptcha {} created", key_name));
MCaptchaSafety::new(ctx, duration, &key_name)?;
REDIS_OK
} else {
let msg = format!("mcaptcha {} exists", key_name);
ctx.log_debug(&msg);
Err(CacheError::new(msg).into())
}
}
REDIS_OK
/// check if captcha exists
pub fn captcha_exists(ctx: &Context, args: Vec<String>) -> RedisResult {
let mut args = args.into_iter().skip(1);
let key_name = get_captcha_key(&args.next_string()?);
let key = ctx.open_key(&key_name);
if key.key_type() == KeyType::Empty {
// 1 is false
Ok(RedisValue::Integer(1))
} else {
// 0 is true
Ok(RedisValue::Integer(0))
}
}
/// Add captcha to redis
pub fn delete_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 key = ctx.open_key_writable(&key_name);
if key.key_type() == KeyType::Empty {
Err(RedisError::nonexistent_key())
} else {
key.delete()?;
REDIS_OK
}
}
}
@@ -160,10 +223,17 @@ pub mod type_methods {
let mcaptcha = match encver {
0 => {
let data = raw::load_string(rdb);
let mcaptcha: MCaptcha = Format::JSON.from_str(&data).unwrap();
mcaptcha
let mcaptcha: Result<MCaptcha, CacheError> = Format::JSON.from_str(&data);
if mcaptcha.is_err() {
panic!(
"Can't load mCaptcha from old redis RDB, error while serde {}, data received: {}",
mcaptcha.err().unwrap(),
data
);
}
mcaptcha.unwrap()
}
_ => panic!("Can't load mCaptcha from old redis RDB"),
_ => panic!("Can't load mCaptcha from old redis RDB, encver {}", encver),
};
Box::into_raw(Box::new(mcaptcha)) as *mut c_void
@@ -183,3 +253,87 @@ pub mod type_methods {
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use libmcaptcha::defense::LevelBuilder;
fn get_levels() -> Vec<Level> {
let mut levels = Vec::default();
levels.push(
LevelBuilder::default()
.visitor_threshold(50)
.difficulty_factor(50)
.unwrap()
.build()
.unwrap(),
);
levels.push(
LevelBuilder::default()
.visitor_threshold(500)
.difficulty_factor(5000)
.unwrap()
.build()
.unwrap(),
);
levels.push(
LevelBuilder::default()
.visitor_threshold(5000)
.difficulty_factor(50000)
.unwrap()
.build()
.unwrap(),
);
levels.push(
LevelBuilder::default()
.visitor_threshold(50000)
.difficulty_factor(500000)
.unwrap()
.build()
.unwrap(),
);
levels.push(
LevelBuilder::default()
.visitor_threshold(500000)
.difficulty_factor(5000000)
.unwrap()
.build()
.unwrap(),
);
levels
}
#[test]
fn create_mcaptcha_works() {
let levels = get_levels();
let payload = CreateMCaptcha {
levels,
duration: 30,
};
let mcaptcha = MCaptcha::new(payload);
assert!(mcaptcha.is_ok());
let mut mcaptcha = mcaptcha.unwrap();
for _ in 0..50 {
mcaptcha.add_visitor();
}
assert_eq!(mcaptcha.get_visitors(), 50);
assert_eq!(mcaptcha.get_difficulty(), 50);
for _ in 0..451 {
mcaptcha.add_visitor();
}
assert_eq!(mcaptcha.get_visitors(), 501);
assert_eq!(mcaptcha.get_difficulty(), 5000);
mcaptcha.decrement_visitor_by(501);
for _ in 0..5002 {
mcaptcha.add_visitor();
}
assert_eq!(mcaptcha.get_visitors(), 5002);
assert_eq!(mcaptcha.get_difficulty(), 50000);
}
}

56
src/safety.rs
View File

@@ -36,42 +36,43 @@ impl MCaptchaSafety {
pub fn new(ctx: &Context, duration: u64, mcaptcha_name: &str) -> CacheResult<()> {
let safety_name = get_safety_name(mcaptcha_name);
let safety = ctx.open_key_writable(&safety_name);
Self::set_timer(ctx, &safety, (&safety_name, duration))?;
if safety.key_type() == KeyType::Empty {
let safety_val = MCaptchaSafety {};
safety.set_value(&MCAPTCHA_SAFETY_TYPE, safety_val)?;
ctx.log_debug(&format!("mcaptcha safety created: {}", safety_name));
Self::set_timer(ctx, &safety, (safety_name, duration))?;
} else {
ctx.log_debug(&format!("mcaptcha safety exists: {}", safety_name));
}
Ok(())
}
fn set_timer(
ctx: &Context,
safety: &RedisKeyWritable,
(safety_name, duration): (&str, u64),
(safety_name, duration): (String, u64),
) -> CacheResult<()> {
let _ = ctx.create_timer(
Duration::from_secs(duration),
Self::boost,
(&safety_name, duration),
(safety_name, duration),
);
safety.set_expire(Duration::from_secs(duration * 2))?;
Ok(())
}
/// executes when timer goes off. Refreshes expiry timer and resets timer
fn boost(ctx: &Context, (safety_name, duration): (&str, u64)) {
let safety = ctx.open_key_writable(safety_name);
let x = safety.get_value::<Self>(&MCAPTCHA_SAFETY_TYPE);
// Result<Option<&mut Safety>, RedisError>
// Ok(Some(val)) => refresh
// _ => check if corresponding captcha is available => Yes -> create timer
// NO -> Ignore
//
fn boost(ctx: &Context, (safety_name, duration): (String, u64)) {
let safety = ctx.open_key_writable(&safety_name);
match safety.get_value::<Self>(&MCAPTCHA_SAFETY_TYPE) {
Ok(Some(_safety_val)) => {
Self::set_timer(ctx, &safety, (&safety_name, duration)).unwrap()
}
Ok(Some(_safety_val)) => match Self::set_timer(ctx, &safety, (safety_name, duration)) {
Ok(_) => (),
Err(e) => ctx.log_warning(&format!("{}", e)),
},
_ => {
let mcaptcha_name = get_mcaptcha_from_safety(safety_name);
let mcaptcha_name = get_mcaptcha_from_safety(&safety_name);
if mcaptcha_name.is_none() {
return;
}
@@ -127,19 +128,24 @@ pub mod type_methods {
use libc::c_int;
use crate::bucket::Format;
const SAFETY_RDB_VAL: &str = "SAFETY";
use super::*;
#[allow(non_snake_case, unused)]
pub extern "C" fn rdb_load(rdb: *mut raw::RedisModuleIO, encver: c_int) -> *mut c_void {
let bucket = match encver {
0 => {
let data = raw::load_string(rdb);
let bucket: MCaptchaSafety = Format::JSON.from_str(&data).unwrap();
bucket
if data == SAFETY_RDB_VAL {
MCaptchaSafety {}
} else {
panic!("Can't safety from old redis RDB, data received : {}", data);
}
}
_ => panic!("Can't load bucket from old redis RDB"),
_ => panic!(
"Can't safety from old redis RDB, encoding version: {}",
encver
),
};
// if bucket.
@@ -153,10 +159,6 @@ pub mod type_methods {
#[allow(non_snake_case, unused)]
pub unsafe extern "C" fn rdb_save(rdb: *mut raw::RedisModuleIO, value: *mut c_void) {
let bucket = &*(value as *mut MCaptchaSafety);
match &serde_json::to_string(bucket) {
Ok(string) => raw::save_string(rdb, &string),
Err(e) => eprintln!("error while rdb_save: {}", e),
}
raw::save_string(rdb, &SAFETY_RDB_VAL)
}
}

25
tests/bucket.py
View File

@@ -16,18 +16,20 @@
from asyncio import sleep
import sys
import test
from mcaptcha import register
from test import REDIS_URL
import utils
r = test.r
r = utils.connect(REDIS_URL)
utils.ping(r)
COMMANDS = {
"COUNT" : "mcaptcha_cache.add_visitor",
"GET" : "mcaptcha_cache.get",
}
def incr(key, time):
r.execute_command(COMMANDS["COUNT"], key, time)
def incr(key):
r.execute_command(COMMANDS["COUNT"], key)
def get_count(key):
try:
@@ -43,14 +45,13 @@ def assert_count(expect, key):
async def incr_one_works():
try:
key = "incr_one"
register(r, key)
time = 2
register(key)
initial_count = get_count(key)
# incriment
incr(key, time)
incr(key)
assert_count(initial_count + 1, key)
# wait till expiry
await sleep(time + 2)
await sleep(5 + 2)
assert_count(initial_count, key)
print("Incr one works")
except Exception as e:
@@ -60,16 +61,14 @@ async def incr_one_works():
async def race_works():
key = "race_works"
try:
register(r, key)
register(key)
initial_count = get_count(key)
race_num = 200
time = 3
for _ in range(race_num):
incr(key, time)
incr(key)
assert_count(initial_count + race_num, key)
# wait till expiry
await sleep(time + 2)
await sleep(5 + 2)
assert_count(initial_count, key)
print("Race works")
except Exception as e:

61
tests/mcaptcha.py
View File

@@ -17,26 +17,71 @@
import json
import utils
from test import REDIS_URL
r = utils.connect(REDIS_URL)
utils.ping(r)
MCAPTCHA = {
"visitor_threshold": 0,
"defense": {
"levels": [
"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",
"DELETE_CAPTCHA": "MCAPTCHA_CACHE.DELETE_CAPTCHA",
"CAPTCHA_EXISTS": "MCAPTCHA_CACHE.CAPTCHA_EXISTS",
}
payload = json.dumps(MCAPTCHA)
def register(r, key):
if r.exists(key):
r.delete(key)
def delete_captcha(key):
r.execute_command(COMMANDS["DELETE_CAPTCHA"], key)
def add_captcha(key):
r.execute_command(COMMANDS["ADD_CAPTCHA"], key, payload)
def captcha_exists(key):
exists = r.execute_command(COMMANDS["CAPTCHA_EXISTS"], key)
if exists == 0:
return True
if exists == 1:
return False
def register(key):
if captcha_exists(key):
delete_captcha(key)
add_captcha(key)
async def captcha_exists_works():
key = "captcha_delete_works"
if captcha_exists(key):
delete_captcha(key)
assert captcha_exists(key) is False
register(key)
assert captcha_exists(key) is True
print("Captcha delete works")
async def register_captcha_works():
key = "register_captcha_works"
register(key)
assert captcha_exists(key) is True
print("Add captcha works")
async def delete_captcha_works():
key = "delete_captcha_works"
register(key)
exists = captcha_exists(key)
print("captcha exists stauts", exists)
assert exists is True
delete_captcha(key)
assert captcha_exists(key) is False
print("Delete captcha works")

30
tests/runner.py
View File

@@ -14,35 +14,25 @@
# 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 threading import Thread
import asyncio
class Runner(object):
_functions = []
_threads = []
_tasks = []
""" Register functions to be run"""
def register(self, fn):
async def register(self, fn):
""" Register functions to be run"""
self._functions.append(fn)
t = Thread(target=fn)
self._threads.append(t)
task = asyncio.create_task(fn())
self._tasks.append(task)
"""Wait for registered functions to finish executing"""
def __run__(self):
for thread in self._threads:
try:
thread.start()
except:
print("yo")
async def wait(self):
"""Wait for registered functions to finish executing"""
def wait(self):
self.__run__()
for thread in self._threads:
try:
thread.join()
except:
print("yo")
for task in self._tasks:
await task
"""Runs in seperate threads"""
def __init__(self):
super(Runner, self).__init__()
# self.arg = arg

38
tests/test.py
View File

@@ -23,31 +23,33 @@ from redis import BlockingConnectionPool
import utils
from runner import Runner
import bucket
import mcaptcha
REDIS_URL = "redis://localhost:6350"
r = utils.connect(REDIS_URL)
utils.ping(r)
async def main():
#runner = Runner()
#fn = [bucket.incr_one_works]#, bucket.race_works]
runner = Runner()
task1 = asyncio.create_task(bucket.incr_one_works())
task2 = asyncio.create_task(bucket.race_works())
await task1
await task2
fn = [
bucket.incr_one_works,
bucket.race_works,
#mcaptcha.delete_captcha_works,
mcaptcha.captcha_exists_works,
mcaptcha.register_captcha_works
]
#try:
# for r in fn:
# runner.register(r)
#tasts = []
#task1 = asyncio.create_task(bucket.incr_one_works())
#task2 = asyncio.create_task(bucket.race_works())
#await task1
#await task2
# runner.wait()
# print("All tests passed")
#except Exception as e:
# raise e
for r in fn:
await runner.register(r)
await runner.wait()
print("All tests passed")
if __name__ == "__main__":
asyncio.run(main())