Accept future OTP code

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(),
         ));
     }
 

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<String> {
+        self.get_code_at(|| time() + PERIOD)
+    }
+
     /// Get the code at a specific time
     fn get_code_at<F: Fn() -> u64>(&self, get_time: F) -> Res<String> {
         let gen = TotpGenerator::new()
@@ -98,7 +103,9 @@ impl TotpKey {
 
     /// Check a code's validity
     pub fn check_code(&self, code: &str) -> Res<bool> {
-        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]