forked from bot/tmp2
478 lines
14 KiB
Rust
478 lines
14 KiB
Rust
/*
|
|
* mCaptcha - A proof of work based DoS protection system
|
|
* Copyright © 2021 Aravinth Manivannan <realravinth@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 <http://www.gnu.org/licenses/>.
|
|
*/
|
|
|
|
//! Defense datatypes
|
|
//! ```rust
|
|
//! use libmcaptcha::{LevelBuilder, DefenseBuilder};
|
|
//! DefenseBuilder::default()
|
|
//! .add_level(
|
|
//! LevelBuilder::default()
|
|
//! .visitor_threshold(50)
|
|
//! .difficulty_factor(50)
|
|
//! .unwrap()
|
|
//! .build()
|
|
//! .unwrap(),
|
|
//! )
|
|
//! .unwrap()
|
|
//! .add_level(
|
|
//! LevelBuilder::default()
|
|
//! .visitor_threshold(500)
|
|
//! .difficulty_factor(500)
|
|
//! .unwrap()
|
|
//! .build()
|
|
//! .unwrap(),
|
|
//! )
|
|
//! .unwrap()
|
|
//! .build()
|
|
//! .unwrap();
|
|
//! ```
|
|
|
|
use crate::errors::*;
|
|
use serde::{Deserialize, Serialize};
|
|
|
|
/// Level struct that describes threshold-difficulty factor mapping
|
|
#[derive(Debug, Deserialize, Serialize, Copy, Clone, PartialEq)]
|
|
pub struct Level {
|
|
pub visitor_threshold: u32,
|
|
pub difficulty_factor: u32,
|
|
}
|
|
|
|
/// Bulder struct for [Level] to describe threshold-difficulty factor mapping
|
|
#[derive(Debug, Copy, Clone, PartialEq)]
|
|
pub struct LevelBuilder {
|
|
visitor_threshold: Option<u32>,
|
|
difficulty_factor: Option<u32>,
|
|
}
|
|
|
|
impl Default for LevelBuilder {
|
|
fn default() -> Self {
|
|
LevelBuilder {
|
|
visitor_threshold: None,
|
|
difficulty_factor: None,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl LevelBuilder {
|
|
/// set visitor count for level
|
|
pub fn visitor_threshold(&mut self, visitor_threshold: u32) -> &mut Self {
|
|
self.visitor_threshold = Some(visitor_threshold);
|
|
self
|
|
}
|
|
|
|
/// set difficulty factor for level. difficulty_factor can't be zero because
|
|
/// Difficulty is calculated as:
|
|
/// ```no_run
|
|
/// let difficulty_factor = 500;
|
|
/// let difficulty = u128::max_value() - u128::max_value() / difficulty_factor;
|
|
/// ```
|
|
/// the higher the `difficulty_factor`, the higher the difficulty.
|
|
pub fn difficulty_factor(&mut self, difficulty_factor: u32) -> CaptchaResult<&mut Self> {
|
|
if difficulty_factor > 0 {
|
|
self.difficulty_factor = Some(difficulty_factor);
|
|
Ok(self)
|
|
} else {
|
|
Err(CaptchaError::DifficultyFactorZero)
|
|
}
|
|
}
|
|
|
|
/// build Level struct
|
|
pub fn build(&mut self) -> CaptchaResult<Level> {
|
|
if self.visitor_threshold.is_none() {
|
|
Err(CaptchaError::SetVisitorThreshold)
|
|
} else if self.difficulty_factor.is_none() {
|
|
Err(CaptchaError::SetDifficultyFactor)
|
|
} else {
|
|
Ok(Level {
|
|
difficulty_factor: self.difficulty_factor.unwrap(),
|
|
visitor_threshold: self.visitor_threshold.unwrap(),
|
|
})
|
|
}
|
|
}
|
|
}
|
|
|
|
/// struct describes all the different [Level]s at which an mCaptcha system operates
|
|
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
|
|
pub struct Defense {
|
|
levels: Vec<Level>,
|
|
// index of current visitor threshold
|
|
current_visitor_threshold: usize,
|
|
}
|
|
|
|
impl From<Defense> for Vec<Level> {
|
|
fn from(d: Defense) -> Self {
|
|
d.levels
|
|
}
|
|
}
|
|
|
|
/// Builder struct for [Defense]
|
|
#[derive(Debug, Clone, PartialEq)]
|
|
pub struct DefenseBuilder {
|
|
levels: Vec<Level>,
|
|
}
|
|
|
|
impl Default for DefenseBuilder {
|
|
fn default() -> Self {
|
|
DefenseBuilder { levels: vec![] }
|
|
}
|
|
}
|
|
|
|
impl DefenseBuilder {
|
|
/// add a level to [Defense]
|
|
pub fn add_level(&mut self, level: Level) -> CaptchaResult<&mut Self> {
|
|
for i in self.levels.iter() {
|
|
if i.visitor_threshold == level.visitor_threshold {
|
|
return Err(CaptchaError::DuplicateVisitorCount);
|
|
}
|
|
}
|
|
self.levels.push(level);
|
|
Ok(self)
|
|
}
|
|
|
|
/// Build [Defense]
|
|
pub fn build(&mut self) -> CaptchaResult<Defense> {
|
|
if !self.levels.is_empty() {
|
|
// sort levels to arrange in ascending order
|
|
self.levels.sort_by_key(|a| a.visitor_threshold);
|
|
|
|
for level in self.levels.iter() {
|
|
if level.difficulty_factor == 0 {
|
|
return Err(CaptchaError::DifficultyFactorZero);
|
|
}
|
|
}
|
|
|
|
// as visitor count increases, difficulty_factor too should increse
|
|
// if it decreses, an error must be thrown
|
|
for i in 0..self.levels.len() - 1 {
|
|
if self.levels[i].difficulty_factor > self.levels[i + 1].difficulty_factor {
|
|
return Err(CaptchaError::DecreaseingDifficultyFactor);
|
|
}
|
|
}
|
|
|
|
Ok(Defense {
|
|
levels: self.levels.to_owned(),
|
|
current_visitor_threshold: 0,
|
|
})
|
|
} else {
|
|
Err(CaptchaError::LevelEmpty)
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Defense {
|
|
///! Difficulty is calculated as:
|
|
///! ```rust
|
|
///! let difficulty = u128::max_value() - u128::max_value() / difficulty_factor;
|
|
///! ```
|
|
///! The higher the `difficulty_factor`, the higher the difficulty.
|
|
|
|
/// Get difficulty factor of current level of defense
|
|
pub fn get_difficulty(&self) -> u32 {
|
|
self.levels[self.current_visitor_threshold].difficulty_factor
|
|
}
|
|
|
|
/// tighten up defense. Increases defense level by a factor of one.
|
|
/// When defense is at max level, calling this method will have no effect
|
|
pub fn tighten_up(&mut self) {
|
|
if self.current_visitor_threshold < self.levels.len() - 1 {
|
|
self.current_visitor_threshold += 1;
|
|
}
|
|
}
|
|
/// Loosen up defense. Decreases defense level by a factor of one.
|
|
/// When defense is at the lowest level, calling this method will have no effect.
|
|
pub fn loosen_up(&mut self) {
|
|
if self.current_visitor_threshold > 0 {
|
|
self.current_visitor_threshold -= 1;
|
|
}
|
|
}
|
|
|
|
/// Set defense to maximum level
|
|
pub fn max_defense(&mut self) {
|
|
self.current_visitor_threshold = self.levels.len() - 1;
|
|
}
|
|
|
|
/// Set defense to minimum level
|
|
pub fn min_defense(&mut self) {
|
|
self.current_visitor_threshold = 0;
|
|
}
|
|
|
|
/// Get current level's visitor threshold
|
|
pub fn visitor_threshold(&self) -> u32 {
|
|
self.levels[self.current_visitor_threshold].difficulty_factor
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn level_builder_works() {
|
|
let level = LevelBuilder::default()
|
|
.difficulty_factor(1)
|
|
.unwrap()
|
|
.visitor_threshold(0)
|
|
.build()
|
|
.unwrap();
|
|
|
|
assert_eq!(level.visitor_threshold, 0);
|
|
assert_eq!(level.difficulty_factor, 1);
|
|
|
|
assert_eq!(
|
|
LevelBuilder::default().difficulty_factor(0),
|
|
Err(CaptchaError::DifficultyFactorZero)
|
|
);
|
|
|
|
assert_eq!(
|
|
LevelBuilder::default()
|
|
.difficulty_factor(1)
|
|
.unwrap()
|
|
.build(),
|
|
Err(CaptchaError::SetVisitorThreshold)
|
|
);
|
|
assert_eq!(
|
|
LevelBuilder::default().visitor_threshold(10).build(),
|
|
Err(CaptchaError::SetDifficultyFactor)
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn defense_builder_duplicate_visitor_threshold() {
|
|
let mut defense_builder = DefenseBuilder::default();
|
|
let err = defense_builder
|
|
.add_level(
|
|
LevelBuilder::default()
|
|
.visitor_threshold(50)
|
|
.difficulty_factor(50)
|
|
.unwrap()
|
|
.build()
|
|
.unwrap(),
|
|
)
|
|
.unwrap()
|
|
.add_level(
|
|
LevelBuilder::default()
|
|
.visitor_threshold(50)
|
|
.difficulty_factor(50)
|
|
.unwrap()
|
|
.build()
|
|
.unwrap(),
|
|
);
|
|
assert_eq!(err, Err(CaptchaError::DuplicateVisitorCount));
|
|
}
|
|
|
|
#[test]
|
|
fn defense_builder_decreasing_difficulty_factor() {
|
|
let mut defense_builder = DefenseBuilder::default();
|
|
let err = defense_builder
|
|
.add_level(
|
|
LevelBuilder::default()
|
|
.visitor_threshold(50)
|
|
.difficulty_factor(50)
|
|
.unwrap()
|
|
.build()
|
|
.unwrap(),
|
|
)
|
|
.unwrap()
|
|
.add_level(
|
|
LevelBuilder::default()
|
|
.visitor_threshold(500)
|
|
.difficulty_factor(10)
|
|
.unwrap()
|
|
.build()
|
|
.unwrap(),
|
|
)
|
|
.unwrap()
|
|
.build();
|
|
assert_eq!(err, Err(CaptchaError::DecreaseingDifficultyFactor));
|
|
}
|
|
|
|
#[test]
|
|
fn checking_for_integer_overflow() {
|
|
let mut defense = DefenseBuilder::default()
|
|
.add_level(
|
|
LevelBuilder::default()
|
|
.visitor_threshold(5)
|
|
.difficulty_factor(5)
|
|
.unwrap()
|
|
.build()
|
|
.unwrap(),
|
|
)
|
|
.unwrap()
|
|
.add_level(
|
|
LevelBuilder::default()
|
|
.visitor_threshold(10)
|
|
.difficulty_factor(50)
|
|
.unwrap()
|
|
.build()
|
|
.unwrap(),
|
|
)
|
|
.unwrap()
|
|
.add_level(
|
|
LevelBuilder::default()
|
|
.visitor_threshold(20)
|
|
.difficulty_factor(60)
|
|
.unwrap()
|
|
.build()
|
|
.unwrap(),
|
|
)
|
|
.unwrap()
|
|
.add_level(
|
|
LevelBuilder::default()
|
|
.visitor_threshold(30)
|
|
.difficulty_factor(65)
|
|
.unwrap()
|
|
.build()
|
|
.unwrap(),
|
|
)
|
|
.unwrap()
|
|
.build()
|
|
.unwrap();
|
|
|
|
for _ in 0..500 {
|
|
defense.tighten_up();
|
|
}
|
|
|
|
defense.get_difficulty();
|
|
for _ in 0..500000 {
|
|
defense.tighten_up();
|
|
}
|
|
|
|
defense.get_difficulty();
|
|
}
|
|
|
|
fn get_defense() -> Defense {
|
|
DefenseBuilder::default()
|
|
.add_level(
|
|
LevelBuilder::default()
|
|
.visitor_threshold(50)
|
|
.difficulty_factor(50)
|
|
.unwrap()
|
|
.build()
|
|
.unwrap(),
|
|
)
|
|
.unwrap()
|
|
.add_level(
|
|
LevelBuilder::default()
|
|
.visitor_threshold(500)
|
|
.difficulty_factor(5000)
|
|
.unwrap()
|
|
.build()
|
|
.unwrap(),
|
|
)
|
|
.unwrap()
|
|
.add_level(
|
|
LevelBuilder::default()
|
|
.visitor_threshold(5000)
|
|
.difficulty_factor(50000)
|
|
.unwrap()
|
|
.build()
|
|
.unwrap(),
|
|
)
|
|
.unwrap()
|
|
.add_level(
|
|
LevelBuilder::default()
|
|
.visitor_threshold(50000)
|
|
.difficulty_factor(500000)
|
|
.unwrap()
|
|
.build()
|
|
.unwrap(),
|
|
)
|
|
.unwrap()
|
|
.add_level(
|
|
LevelBuilder::default()
|
|
.visitor_threshold(500000)
|
|
.difficulty_factor(5000000)
|
|
.unwrap()
|
|
.build()
|
|
.unwrap(),
|
|
)
|
|
.unwrap()
|
|
.build()
|
|
.unwrap()
|
|
}
|
|
#[test]
|
|
fn defense_builder_works() {
|
|
let defense = get_defense();
|
|
|
|
assert_eq!(defense.levels[0].difficulty_factor, 50);
|
|
assert_eq!(defense.levels[1].difficulty_factor, 5000);
|
|
assert_eq!(defense.levels[2].difficulty_factor, 50_000);
|
|
assert_eq!(defense.levels[3].difficulty_factor, 500_000);
|
|
assert_eq!(defense.levels[4].difficulty_factor, 5_000_000);
|
|
}
|
|
|
|
#[test]
|
|
fn tighten_up_works() {
|
|
let mut defense = get_defense();
|
|
|
|
assert_eq!(defense.get_difficulty(), 50);
|
|
|
|
defense.tighten_up();
|
|
assert_eq!(defense.get_difficulty(), 5_000);
|
|
|
|
defense.tighten_up();
|
|
assert_eq!(defense.get_difficulty(), 50_000);
|
|
|
|
defense.tighten_up();
|
|
assert_eq!(defense.get_difficulty(), 500_000);
|
|
|
|
defense.tighten_up();
|
|
assert_eq!(defense.get_difficulty(), 5_000_000);
|
|
|
|
defense.tighten_up();
|
|
assert_eq!(defense.get_difficulty(), 5_000_000);
|
|
}
|
|
|
|
#[test]
|
|
fn max_defense_works() {
|
|
let mut defense = get_defense();
|
|
defense.max_defense();
|
|
assert_eq!(defense.get_difficulty(), 5_000_000);
|
|
}
|
|
|
|
#[test]
|
|
fn minimum_defense_works() {
|
|
let mut defense = get_defense();
|
|
defense.min_defense();
|
|
assert_eq!(defense.get_difficulty(), 50);
|
|
}
|
|
|
|
#[test]
|
|
fn loosen_up_works() {
|
|
let mut defense = get_defense();
|
|
defense.max_defense();
|
|
|
|
assert_eq!(defense.get_difficulty(), 5_000_000);
|
|
|
|
defense.loosen_up();
|
|
assert_eq!(defense.get_difficulty(), 500_000);
|
|
|
|
defense.loosen_up();
|
|
assert_eq!(defense.get_difficulty(), 50_000);
|
|
|
|
defense.loosen_up();
|
|
assert_eq!(defense.get_difficulty(), 5_000);
|
|
|
|
defense.loosen_up();
|
|
assert_eq!(defense.get_difficulty(), 50);
|
|
|
|
defense.loosen_up();
|
|
assert_eq!(defense.get_difficulty(), 50);
|
|
}
|
|
}
|