diff --git a/src/defense.rs b/src/defense.rs index 32c39d1..8294292 100644 --- a/src/defense.rs +++ b/src/defense.rs @@ -206,12 +206,8 @@ impl Defense { } /// Get current level's visitor threshold - pub fn visitor_threshold(&self) -> Option { - if let Some(level) = self.levels.get(self.current_visitor_threshold) { - Some(level.visitor_threshold) - } else { - None - } + pub fn visitor_threshold(&self) -> u32 { + self.levels[self.current_visitor_threshold].difficulty_factor } } @@ -299,6 +295,60 @@ mod tests { 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( diff --git a/src/mcaptcha.rs b/src/mcaptcha.rs index d653c53..a5919c2 100644 --- a/src/mcaptcha.rs +++ b/src/mcaptcha.rs @@ -83,7 +83,7 @@ impl MCaptcha { #[inline] pub fn add_visitor(&mut self) { self.visitor_threshold += 1; - if self.visitor_threshold > self.defense.visitor_threshold().unwrap() { + if self.visitor_threshold > self.defense.visitor_threshold() { self.defense.tighten_up(); } else { self.defense.loosen_up();