From 4f32055d8bade0495f5cd5f37e38a0b51948b714 Mon Sep 17 00:00:00 2001
From: constantoine <cleo.rebert@gmail.com>
Date: Fri, 20 May 2022 17:43:39 +0200
Subject: [PATCH] Add next_step and next_step_current methods

Signed-off-by: constantoine <cleo.rebert-ext@treezor.com>
---
 src/lib.rs | 30 ++++++++++++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/src/lib.rs b/src/lib.rs
index 92fce39..f29533d 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -216,6 +216,21 @@ impl<T: AsRef<[u8]>> TOTP<T> {
         )
     }
 
+    /// Returns the timestamp of the first second for the next step
+    /// given the provided timestamp in seconds
+    pub fn next_step(&self, time: u64) -> u64 {
+        let step = time / self.step;
+
+        (step + 1) * self.step
+    }
+
+    /// Returns the timestamp of the first second of the next step
+    /// According to system time
+    pub fn next_step_current(&self)-> Result<u64, SystemTimeError> {
+        let t = system_time()?;
+        Ok(self.next_step(t))
+    }
+
     /// Generate a token from the current system time
     pub fn generate_current(&self) -> Result<String, SystemTimeError> {
         let t = system_time()?;
@@ -573,6 +588,21 @@ mod tests {
         );
     }
 
+    #[test]
+    fn next_step() {
+        let totp = TOTP::new(Algorithm::SHA1, 6, 1, 30, "TestSecret", Some("Github".to_string()), "constantoine@github.com".to_string()).unwrap();
+        assert!(totp.next_step(0) == 30);
+        assert!(totp.next_step(29) == 30);
+        assert!(totp.next_step(30) == 60);
+    }
+
+    #[test]
+    fn next_step_current() {
+        let totp = TOTP::new(Algorithm::SHA1, 6, 1, 30, "TestSecret", Some("Github".to_string()), "constantoine@github.com".to_string()).unwrap();
+        let t = system_time().unwrap();
+        assert!(totp.next_step_current().unwrap() == totp.next_step(t));
+    }
+
     #[test]
     #[cfg(feature = "otpauth")]
     fn from_url_err() {