diff --git a/src/controllers/two_factor_api.rs b/src/controllers/two_factor_api.rs index abdd66a..369e8dc 100644 --- a/src/controllers/two_factor_api.rs +++ b/src/controllers/two_factor_api.rs @@ -38,9 +38,10 @@ pub async fn save_totp_factor( if !key.check_code(&form.first_code).unwrap_or(false) { return HttpResponse::BadRequest().body(format!( - "Given code is invalid (expected {} or {})!", + "Given code is invalid (expected {}, {} or {})!", + key.previous_code().unwrap_or_default(), key.current_code().unwrap_or_default(), - key.previous_code().unwrap_or_default() + key.following_code().unwrap_or_default(), )); } diff --git a/src/data/totp_key.rs b/src/data/totp_key.rs index 00725ca..307f960 100644 --- a/src/data/totp_key.rs +++ b/src/data/totp_key.rs @@ -73,6 +73,11 @@ impl TotpKey { self.get_code_at(|| time() - PERIOD) } + /// Get following code + pub fn following_code(&self) -> Res { + self.get_code_at(|| time() + PERIOD) + } + /// Get the code at a specific time fn get_code_at u64>(&self, get_time: F) -> Res { let gen = TotpGenerator::new() @@ -98,7 +103,9 @@ impl TotpKey { /// Check a code's validity pub fn check_code(&self, code: &str) -> Res { - Ok(self.current_code()?.eq(code) || self.previous_code()?.eq(code)) + Ok(self.previous_code()?.eq(code) + || self.current_code()?.eq(code) + || self.following_code()?.eq(code)) } } @@ -111,7 +118,10 @@ mod test { let key = TotpKey::new_random(); let code = key.current_code().unwrap(); let old_code = key.previous_code().unwrap(); + let following_code = key.following_code().unwrap(); assert_ne!(code, old_code); + assert_ne!(code, following_code); + assert_ne!(old_code, following_code); } #[test]