mirrors/mCaptcha-cache
1
0
Fork 0
mirror of https://github.com/mCaptcha/cache.git synced 2024-11-21 16:25:19 +03:00
Code Issues Projects Releases Packages Wiki Activity

MCaptcha: increment, decrement & rdb persist

Browse source
This commit is contained in:
realaravinth 2021-06-06 16:35:14 +05:30
parent fd9b6f75ae
commit d3bce7669a
No known key found for this signature in database
GPG key ID: AD9F0F08E855ED88
13 changed files with 274 additions and 146 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
Cargo.lock generated
View file

@ -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",

1
Cargo.toml
View file

@ -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}

41
docs/mechanism.md Normal file
View file

@ -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.

132
src/bucket.rs
View file

@ -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()? {
Some(val) => {
if val.trim().is_empty() {
captcha.write("1")?;
} else {
let mut val: usize = val.parse()?;
val += 1;
captcha.write(&val.to_string())?;
}
}
None => {
captcha.write("1")?;
}
match captcha {
Some(val) => val.add_visitor(),
None => return Err(CacheError::new("Captcha not found".into())),
}
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 = fmt.from_str(&data).unwrap();
let bucket: Bucket = Format::JSON.from_str(&data).unwrap();
bucket
}
_ => panic!("Can't load old RedisJSON RDB"),
_ => panic!("Can't load bucket from old redis RDB"),
};
// if bucket.

29
src/errors.rs
View file

@ -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::RedisError(val) => Err(val),
CacheError::Msg(val) => RedisError::String(val),
CacheError::RedisError(val) => val,
CacheError::CaptchaNotFound => RedisError::String(format!("{}", e)),
}
}
}

7
src/lib.rs
View file

@ -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.get", mcaptcha::MCaptcha::get, "readonly", 1, 1, 1],
["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],
],
event_handlers: [
[@EXPIRED @EVICTED: bucket::Bucket::on_delete],

63
src/mcaptcha.rs
View file

@ -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
pub fn get(ctx: &Context, args: Vec<String>) -> RedisResult {
use redis_module::NextArg;
/// decrement [MCaptcha]'s current visitor_threshold by specified count
#[inline]
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 fmt = Format::JSON;
let mcaptcha: MCaptcha = fmt.from_str(&data).unwrap();
let mcaptcha: MCaptcha = Format::JSON.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),
}
}
}

2
src/utils.rs
View file

@ -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]

60
tests/bucket.py
View file

@ -1,6 +1,5 @@
#!/bin/env python3
#
# Copyright (C) 2021 Aravinth Manivannan <realaravinth@batsense.net>
#!/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
@ -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"
time = 2
initial_count = get_count(key)
# incriment
incr(key, time)
assert_count(initial_count + 1, key)
# wait till expiry
sleep(time + 2)
assert_count(initial_count, key)
print("Incr one works")
try:
key = "incr_one"
register(r, key)
time = 2
initial_count = get_count(key)
# incriment
incr(key, time)
assert_count(initial_count + 1, key)
# wait till expiry
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)
race_num = 200
time = 3
try:
register(r, key)
initial_count = get_count(key)
race_num = 200
time = 3
for _ in range(race_num):
incr(key, time)
assert_count(initial_count + race_num, key)
# wait till expiry
sleep(time + 2)
assert_count(initial_count, key)
print("Race works")
for _ in range(race_num):
incr(key, time)
assert_count(initial_count + race_num, key)
# wait till expiry
sleep(time + 2)
assert_count(initial_count, key)
print("Race works")
except Exception as e:
raise e

42
tests/mcaptcha.py Normal file
View file

@ -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)

18
tests/runner.py
View file

@ -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):

19
tests/test.py
View file

@ -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()
fn = [bucket.incr_one_works, bucket.race_works]
for r in fn:
runner.register(r)
#runner = Runner()
#fn = [bucket.incr_one_works]#, bucket.race_works]
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()

3
tests/utils.py
View file

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