rolling_token_auth/
lib.rs

1use hmac::{Hmac, Mac};
2use sha2::Sha256;
3use std::time::{SystemTime, UNIX_EPOCH};
4
5type HmacSha256 = Hmac<Sha256>;
6
7#[derive(Debug, Clone)]
8pub struct Token {
9    pub token: String,
10    pub timestamp: i64,
11}
12
13impl Token {
14    fn get_offset(&self, manager: &RollingTokenManager) -> i64 {
15        self.timestamp - manager.current_timestamp()
16    }
17}
18
19#[derive(Clone)]
20pub struct RollingTokenManager {
21    secret: Vec<u8>,
22    interval: i64,
23    tolerance: i64,
24    active_tokens: Vec<Token>,
25}
26
27impl RollingTokenManager {
28    pub fn new(secret: impl Into<Vec<u8>>, interval: i64, tolerance: Option<i64>) -> Self {
29        Self {
30            secret: secret.into(),
31            interval,
32            tolerance: tolerance.unwrap_or(1),
33            active_tokens: Vec::new(),
34        }
35    }
36
37    fn current_timestamp(&self) -> i64 {
38        SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs() as i64 / self.interval
39    }
40
41    pub fn generate_token_with_offset(&self, offset: i64) -> Token {
42        let timestamp = self.current_timestamp() + offset;
43        let encoded_timestamp = timestamp.to_string();
44
45        let mut mac = HmacSha256::new_from_slice(&self.secret).expect("HMAC can take key of any size");
46
47        mac.update(encoded_timestamp.as_bytes());
48        let result = mac.finalize();
49        let token = hex::encode(result.into_bytes());
50
51        Token { token, timestamp }
52    }
53
54    pub fn generate_token(&self) -> Token {
55        self.generate_token_with_offset(0)
56    }
57
58    fn refresh_tokens(&mut self) {
59        let current_time = self.current_timestamp();
60
61        // Remove tokens outside tolerance
62        self.active_tokens
63            .retain(|token| (token.timestamp - current_time).abs() <= self.tolerance);
64
65        if self.active_tokens.len() as i64 == 1 + 2 * self.tolerance {
66            return;
67        }
68
69        // Create a set of timestamps we need to generate
70        let mut needed_timestamps: Vec<i64> = (-self.tolerance..=self.tolerance).map(|offset| current_time + offset).collect();
71
72        // Remove timestamps we already have
73        for token in &self.active_tokens {
74            needed_timestamps.retain(|&t| t != token.timestamp);
75        }
76
77        // Generate missing tokens
78        for timestamp in needed_timestamps {
79            let offset = timestamp - current_time;
80            let token = self.generate_token_with_offset(offset);
81            self.active_tokens.push(token);
82        }
83    }
84
85    pub fn is_valid(&mut self, token: &str) -> bool {
86        self.refresh_tokens();
87        self.active_tokens.iter().any(|t| t.token == token)
88    }
89}
90
91#[cfg(test)]
92mod tests {
93    use super::*;
94
95    #[test]
96    fn test_token_validation() {
97        let mut manager = RollingTokenManager::new("test_secret", 30, Some(1));
98        let token = manager.generate_token();
99        assert!(manager.is_valid(&token.token));
100        assert!(token.get_offset(&manager) == 0);
101
102        let token_offset_1 = manager.generate_token_with_offset(1);
103        assert!(manager.is_valid(&token_offset_1.token));
104        assert!(token_offset_1.get_offset(&manager) == 1);
105
106        let token_offset_2 = manager.generate_token_with_offset(2);
107        assert!(!manager.is_valid(&token_offset_2.token)); // token is too far in the future -> invalid
108        assert!(token_offset_2.get_offset(&manager) == 2);
109    }
110
111    #[test]
112    fn test_invalid_token() {
113        let mut manager = RollingTokenManager::new("test_secret", 30, Some(1));
114        assert!(!manager.is_valid("invalid_token"));
115    }
116}
rolling_token_auth/lib.rs

rolling_token_auth/
lib.rs
