fix: key expiration logic (#20)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

Author: ssddOnTop

crates/api-keys-simplified/src/domain.rs

@@ -27,6 +27,8 @@ pub struct ApiKeyManagerV0 {
     validator: KeyValidator,
     #[getter(skip)]
     include_checksum: bool,
+    #[getter(skip)]
+    expiry_grace_period: std::time::Duration,
 }
 
 // FIXME: Need better naming
@@ -53,28 +55,35 @@ impl ApiKeyManagerV0 {
         prefix: impl Into<String>,
         config: KeyConfig,
         hash_config: HashConfig,
+        expiry_grace_period: std::time::Duration,
     ) -> std::result::Result<Self, InitError> {
         let include_checksum = *config.checksum_length() != 0;
         let prefix = KeyPrefix::new(prefix)?;
         let generator = KeyGenerator::new(prefix, config)?;
-        let hasher = KeyHasher::new(hash_config.clone());
+        let hasher = KeyHasher::new(hash_config);
 
         // Generate dummy key and its hash for timing attack protection
         let dummy_key = generator.dummy_key().clone();
         let dummy_hash = hasher.hash(&dummy_key)?;
 
-        let validator = KeyValidator::new(&hash_config, include_checksum, dummy_key, dummy_hash)?;
+        let validator = KeyValidator::new(include_checksum, dummy_key, dummy_hash)?;
 
         Ok(Self {
             generator,
             hasher,
             validator,
             include_checksum,
+            expiry_grace_period,
         })
     }
 
     pub fn init_default_config(prefix: impl Into<String>) -> std::result::Result<Self, InitError> {
-        Self::init(prefix, KeyConfig::default(), HashConfig::default())
+        Self::init(
+            prefix,
+            KeyConfig::default(),
+            HashConfig::default(),
+            std::time::Duration::from_secs(10),
+        )
     }
     pub fn init_high_security_config(
         prefix: impl Into<String>,
@@ -83,6 +92,7 @@ impl ApiKeyManagerV0 {
             prefix,
             KeyConfig::high_security(),
             HashConfig::high_security(),
+            std::time::Duration::from_secs(10),
         )
     }
 
@@ -143,11 +153,23 @@ impl ApiKeyManagerV0 {
     ///
     /// Returns `KeyStatus` indicating whether the key is valid or invalid.
     ///
+    /// # Parameters
+    ///
+    /// * `key` - The API key to verify
+    /// * `stored_hash` - The Argon2 hash stored in your database
+    /// * `expiry_grace_period` - Optional grace period duration after expiry.
+    ///   - `None`: Skip expiry validation (all keys treated as non-expired)
+    ///   - `Some(Duration::ZERO)`: Strict expiry check (no grace period)
+    ///   - `Some(duration)`: Key remains valid for `duration` after its expiry time
+    ///
+    /// The grace period protects against clock skew issues. Once a key expires beyond
+    /// the grace period, it stays expired even if the system clock goes backwards.
+    ///
     /// # Security Flow
     ///
     /// 1. **Checksum validation** (if enabled): Rejects invalid keys in ~20μs
     /// 2. **Argon2 verification**: Verifies hash for valid checksums (~300ms)
-    /// 3. **Expiry check**: Returns `Invalid` if the key's timestamp has passed
+    /// 3. **Expiry check**: Returns `Invalid` if expired beyond grace period
     ///
     /// # Returns
     ///
@@ -165,6 +187,7 @@ impl ApiKeyManagerV0 {
     ///
     /// ```rust
     /// # use api_keys_simplified::{ApiKeyManagerV0, Environment, KeyStatus};
+    /// # use std::time::Duration;
     /// # let manager = ApiKeyManagerV0::init_default_config("sk").unwrap();
     /// # let key = manager.generate(Environment::production()).unwrap();
     /// match manager.verify(key.key(), key.hash())? {
@@ -178,8 +201,11 @@ impl ApiKeyManagerV0 {
             return Ok(KeyStatus::Invalid);
         }
 
-        self.validator
-            .verify(key.expose_secret(), stored_hash.as_ref())
+        self.validator.verify(
+            key.expose_secret(),
+            stored_hash.as_ref(),
+            self.expiry_grace_period,
+        )
     }
 
     pub fn verify_checksum(&self, key: &SecureString) -> Result<bool> {
@@ -282,7 +308,13 @@ mod tests {
     fn test_custom_config() {
         let config = KeyConfig::new().with_entropy(32).unwrap();
 
-        let generator = ApiKeyManagerV0::init("custom", config, HashConfig::default()).unwrap();
+        let generator = ApiKeyManagerV0::init(
+            "custom",
+            config,
+            HashConfig::default(),
+            std::time::Duration::ZERO,
+        )
+        .unwrap();
         let key = generator.generate(Environment::production()).unwrap();
         assert!(generator.verify_checksum(key.key()).unwrap());
     }

crates/api-keys-simplified/src/generator.rs

@@ -373,8 +373,13 @@ mod tests {
 
     #[test]
     fn test_verify_checksum_dos_protection() {
-        let generator =
-            ApiKeyManagerV0::init("sk", KeyConfig::balanced(), HashConfig::default()).unwrap();
+        let generator = ApiKeyManagerV0::init(
+            "sk",
+            KeyConfig::balanced(),
+            HashConfig::default(),
+            std::time::Duration::ZERO,
+        )
+        .unwrap();
 
         // Test oversized key rejection
         let huge_key = SecureString::from("a".repeat(1000));

crates/api-keys-simplified/src/token_parser.rs

@@ -161,7 +161,8 @@ mod tests {
             HashConfig::default()
         };
 
-        let generator = ApiKeyManagerV0::init("text", config, hash_config).unwrap();
+        let generator =
+            ApiKeyManagerV0::init("text", config, hash_config, std::time::Duration::ZERO).unwrap();
         let ts = Utc::now();
         let key = if with_expiry {
             generator

crates/api-keys-simplified/src/validator.rs

@@ -33,7 +33,6 @@ impl KeyValidator {
     const MAX_HASH_LENGTH: usize = 512;
 
     pub fn new(
-        _hash_config: &crate::HashConfig,
         has_checksum: bool,
         dummy_key: SecureString,
         dummy_hash: String,
@@ -48,27 +47,39 @@ impl KeyValidator {
         })
     }
 
-    fn verify_expiry(&self, parts: Parts) -> Result<KeyStatus> {
+    fn verify_expiry(
+        &self,
+        parts: Parts,
+        expiry_grace_period: std::time::Duration,
+    ) -> Result<KeyStatus> {
         if let Some(expiry) = parts.expiry_b64 {
             let decoded = URL_SAFE_NO_PAD
                 .decode(expiry)
                 .or(Err(Error::InvalidFormat))?;
-            let expiry = i64::from_be_bytes(decoded.try_into().or(Err(Error::InvalidFormat))?);
-
-            // TODO(ARCHITECTURE): time libs are platform dependent.
-            // We should set an `infra` layer and abstract
-            // out these libs.
-            if chrono::Utc::now().timestamp() <= expiry {
-                Ok(KeyStatus::Valid)
-            } else {
-                Ok(KeyStatus::Invalid)
+            let expiry_timestamp =
+                i64::from_be_bytes(decoded.try_into().or(Err(Error::InvalidFormat))?);
+
+            let current_time = chrono::Utc::now().timestamp();
+            let grace_seconds = expiry_grace_period.as_secs() as i64;
+
+            // Key is invalid if it expired more than grace_period ago
+            // This ensures once a key expires beyond the grace period, it stays expired
+            // even if the clock goes backwards
+            if expiry_timestamp + grace_seconds < current_time {
+                return Ok(KeyStatus::Invalid);
             }
+            Ok(KeyStatus::Valid)
         } else {
             Ok(KeyStatus::Valid)
         }
     }
 
-    pub fn verify(&self, provided_key: &str, stored_hash: &str) -> Result<KeyStatus> {
+    pub fn verify(
+        &self,
+        provided_key: &str,
+        stored_hash: &str,
+        expiry_grace_period: std::time::Duration,
+    ) -> Result<KeyStatus> {
         // Input length validation to prevent DoS attacks
         if provided_key.len() > Self::MAX_KEY_LENGTH {
             self.dummy_load();
@@ -109,7 +120,7 @@ impl KeyValidator {
         // SECURITY: Force evaluation of expiry check BEFORE the match to ensure
         // constant-time execution. This prevents the compiler from short-circuiting
         // the expiry check when argon_result is Invalid, which would create a timing oracle.
-        let expiry_result = self.verify_expiry(token_parts)?;
+        let expiry_result = self.verify_expiry(token_parts, expiry_grace_period)?;
 
         match (argon_result, expiry_result) {
             (KeyStatus::Invalid, _) | (_, KeyStatus::Invalid) => Ok(KeyStatus::Invalid),
@@ -150,26 +161,30 @@ mod tests {
         let hash = hasher.hash(&key).unwrap();
 
         let (dummy_key, dummy_hash) = dummy_key_and_hash();
-        let validator =
-            KeyValidator::new(&HashConfig::default(), true, dummy_key, dummy_hash).unwrap();
+        let validator = KeyValidator::new(true, dummy_key, dummy_hash).unwrap();
         assert_eq!(
             validator
-                .verify(key.expose_secret(), hash.as_ref())
+                .verify(
+                    key.expose_secret(),
+                    hash.as_ref(),
+                    std::time::Duration::ZERO
+                )
                 .unwrap(),
             KeyStatus::Valid
         );
         assert_eq!(
-            validator.verify("wrong_key", hash.as_ref()).unwrap(),
+            validator
+                .verify("wrong_key", hash.as_ref(), std::time::Duration::ZERO)
+                .unwrap(),
             KeyStatus::Invalid
         );
     }
 
     #[test]
     fn test_invalid_hash_format() {
         let (dummy_key, dummy_hash) = dummy_key_and_hash();
-        let validator =
-            KeyValidator::new(&HashConfig::default(), true, dummy_key, dummy_hash).unwrap();
-        let result = validator.verify("any_key", "invalid_hash");
+        let validator = KeyValidator::new(true, dummy_key, dummy_hash).unwrap();
+        let result = validator.verify("any_key", "invalid_hash", std::time::Duration::ZERO);
         // After timing oracle fix: invalid hash format returns Ok(Invalid) instead of Err
         // to prevent timing-based user enumeration attacks
         assert!(result.is_ok());
@@ -184,9 +199,8 @@ mod tests {
         let hash = hasher.hash(&valid_key).unwrap();
 
         let (dummy_key, dummy_hash) = dummy_key_and_hash();
-        let validator =
-            KeyValidator::new(&HashConfig::default(), true, dummy_key, dummy_hash).unwrap();
-        let result = validator.verify(&oversized_key, hash.as_ref());
+        let validator = KeyValidator::new(true, dummy_key, dummy_hash).unwrap();
+        let result = validator.verify(&oversized_key, hash.as_ref(), std::time::Duration::ZERO);
         assert!(result.is_err());
         assert!(matches!(result.unwrap_err(), Error::InvalidFormat));
     }
@@ -196,9 +210,8 @@ mod tests {
         let oversized_hash = "a".repeat(513); // Exceeds MAX_HASH_LENGTH
 
         let (dummy_key, dummy_hash) = dummy_key_and_hash();
-        let validator =
-            KeyValidator::new(&HashConfig::default(), true, dummy_key, dummy_hash).unwrap();
-        let result = validator.verify("valid_key", &oversized_hash);
+        let validator = KeyValidator::new(true, dummy_key, dummy_hash).unwrap();
+        let result = validator.verify("valid_key", &oversized_hash, std::time::Duration::ZERO);
         assert!(result.is_err());
         assert!(matches!(result.unwrap_err(), Error::InvalidFormat));
     }
@@ -210,17 +223,16 @@ mod tests {
         let hash = hasher.hash(&valid_key).unwrap();
 
         let (dummy_key, dummy_hash) = dummy_key_and_hash();
-        let validator =
-            KeyValidator::new(&HashConfig::default(), true, dummy_key, dummy_hash).unwrap();
+        let validator = KeyValidator::new(true, dummy_key, dummy_hash).unwrap();
 
         // Test at boundary (512 chars - should pass)
         let max_key = "a".repeat(512);
-        let result = validator.verify(&max_key, hash.as_ref());
+        let result = validator.verify(&max_key, hash.as_ref(), std::time::Duration::ZERO);
         assert!(result.is_ok()); // Should not error on length check
 
         // Test just over boundary (513 chars - should fail)
         let over_max_key = "a".repeat(513);
-        let result = validator.verify(&over_max_key, hash.as_ref());
+        let result = validator.verify(&over_max_key, hash.as_ref(), std::time::Duration::ZERO);
         assert!(result.is_err());
         assert!(matches!(result.unwrap_err(), Error::InvalidFormat));
     }
@@ -232,18 +244,25 @@ mod tests {
         let valid_hash = hasher.hash(&valid_key).unwrap();
 
         let (dummy_key, dummy_hash) = dummy_key_and_hash();
-        let validator =
-            KeyValidator::new(&HashConfig::default(), true, dummy_key, dummy_hash).unwrap();
+        let validator = KeyValidator::new(true, dummy_key, dummy_hash).unwrap();
 
-        let result1 = validator.verify("wrong_key", valid_hash.as_ref());
+        let result1 = validator.verify("wrong_key", valid_hash.as_ref(), std::time::Duration::ZERO);
         assert!(result1.is_ok());
         assert_eq!(result1.unwrap(), KeyStatus::Invalid);
 
-        let result2 = validator.verify(valid_key.expose_secret(), "invalid_hash_format");
+        let result2 = validator.verify(
+            valid_key.expose_secret(),
+            "invalid_hash_format",
+            std::time::Duration::ZERO,
+        );
         assert!(result2.is_ok());
         assert_eq!(result2.unwrap(), KeyStatus::Invalid);
 
-        let result3 = validator.verify(valid_key.expose_secret(), "not even close to valid");
+        let result3 = validator.verify(
+            valid_key.expose_secret(),
+            "not even close to valid",
+            std::time::Duration::ZERO,
+        );
         assert!(result3.is_ok());
         assert_eq!(result3.unwrap(), KeyStatus::Invalid);
     }

crates/api-keys-simplified/tests/checksum_dos_protection.rs

@@ -44,7 +44,13 @@ fn test_checksum_prevents_expensive_verification() {
 fn test_valid_checksum_proceeds_to_argon2() {
     // Verify that valid checksums still go through Argon2 verification
     let config = KeyConfig::default().disable_checksum();
-    let generator = ApiKeyManagerV0::init("verify", config, HashConfig::default()).unwrap();
+    let generator = ApiKeyManagerV0::init(
+        "verify",
+        config,
+        HashConfig::default(),
+        std::time::Duration::ZERO,
+    )
+    .unwrap();
 
     let key = generator.generate(Environment::production()).unwrap();
 
@@ -70,13 +76,19 @@ fn test_valid_checksum_proceeds_to_argon2() {
 #[test]
 fn test_dos_protection_comparison() {
     // Compare DoS resistance: with vs without checksum
-    let with_checksum =
-        ApiKeyManagerV0::init("dos1", KeyConfig::default(), HashConfig::default()).unwrap();
+    let with_checksum = ApiKeyManagerV0::init(
+        "dos1",
+        KeyConfig::default(),
+        HashConfig::default(),
+        std::time::Duration::ZERO,
+    )
+    .unwrap();
 
     let without_checksum = ApiKeyManagerV0::init(
         "dos2",
         KeyConfig::default().disable_checksum(),
         HashConfig::default(),
+        std::time::Duration::ZERO,
     )
     .unwrap();
 
@@ -124,6 +136,7 @@ fn test_without_checksum_still_works() {
         "nochk",
         KeyConfig::default().disable_checksum(),
         HashConfig::default(),
+        std::time::Duration::ZERO,
     )
     .unwrap();