rns_crypto/
token.rs

1use alloc::vec::Vec;
2use core::fmt;
3
4use crate::aes128::Aes128;
5use crate::aes256::Aes256;
6use crate::hmac::{hmac_sha256, hmac_sha256_verify};
7use crate::pkcs7;
8use crate::Rng;
9
10pub const TOKEN_OVERHEAD: usize = 48; // 16 IV + 32 HMAC
11
12#[derive(Debug, PartialEq)]
13pub enum TokenError {
14    InvalidKeyLength,
15    InvalidToken,
16    HmacMismatch,
17    DecryptionFailed,
18}
19
20impl fmt::Display for TokenError {
21    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
22        match self {
23            TokenError::InvalidKeyLength => write!(f, "Token key must be 32 or 64 bytes"),
24            TokenError::InvalidToken => write!(f, "Token too short"),
25            TokenError::HmacMismatch => write!(f, "Token HMAC was invalid"),
26            TokenError::DecryptionFailed => write!(f, "Could not decrypt token"),
27        }
28    }
29}
30
31enum AesMode {
32    Aes128(Aes128),
33    Aes256(Aes256),
34}
35
36pub struct Token {
37    signing_key: Vec<u8>,
38    mode: AesMode,
39}
40
41impl Token {
42    pub fn new(key: &[u8]) -> Result<Self, TokenError> {
43        match key.len() {
44            32 => {
45                let signing_key = key[..16].to_vec();
46                let encryption_key: [u8; 16] = key[16..32].try_into().unwrap();
47                Ok(Token {
48                    signing_key,
49                    mode: AesMode::Aes128(Aes128::new(&encryption_key)),
50                })
51            }
52            64 => {
53                let signing_key = key[..32].to_vec();
54                let encryption_key: [u8; 32] = key[32..64].try_into().unwrap();
55                Ok(Token {
56                    signing_key,
57                    mode: AesMode::Aes256(Aes256::new(&encryption_key)),
58                })
59            }
60            _ => Err(TokenError::InvalidKeyLength),
61        }
62    }
63
64    pub fn encrypt(&self, plaintext: &[u8], rng: &mut dyn Rng) -> Vec<u8> {
65        let mut iv = [0u8; 16];
66        rng.fill_bytes(&mut iv);
67        self.encrypt_with_iv(plaintext, &iv)
68    }
69
70    pub fn encrypt_with_iv(&self, plaintext: &[u8], iv: &[u8; 16]) -> Vec<u8> {
71        let padded = pkcs7::pad(plaintext, 16);
72        let ciphertext = match &self.mode {
73            AesMode::Aes128(aes) => aes.encrypt_cbc(&padded, iv),
74            AesMode::Aes256(aes) => aes.encrypt_cbc(&padded, iv),
75        }
76        .expect("PKCS7 padding always produces block-aligned plaintext");
77
78        let mut signed_parts = Vec::with_capacity(16 + ciphertext.len());
79        signed_parts.extend_from_slice(iv);
80        signed_parts.extend_from_slice(&ciphertext);
81
82        let mac = hmac_sha256(&self.signing_key, &signed_parts);
83
84        let mut result = Vec::with_capacity(signed_parts.len() + 32);
85        result.extend_from_slice(&signed_parts);
86        result.extend_from_slice(&mac);
87        result
88    }
89
90    pub fn verify_hmac(&self, token: &[u8]) -> Result<bool, TokenError> {
91        if token.len() <= 32 {
92            return Err(TokenError::InvalidToken);
93        }
94        let received_hmac = &token[token.len() - 32..];
95        Ok(hmac_sha256_verify(
96            &self.signing_key,
97            &token[..token.len() - 32],
98            received_hmac,
99        ))
100    }
101
102    pub fn decrypt(&self, token: &[u8]) -> Result<Vec<u8>, TokenError> {
103        if token.len() <= TOKEN_OVERHEAD {
104            return Err(TokenError::InvalidToken);
105        }
106
107        if !self
108            .verify_hmac(token)
109            .map_err(|_| TokenError::InvalidToken)?
110        {
111            return Err(TokenError::HmacMismatch);
112        }
113
114        let iv: [u8; 16] = token[..16].try_into().unwrap();
115        let ciphertext = &token[16..token.len() - 32];
116
117        // CBC ciphertext must be a non-empty whole number of 16-byte blocks.
118        if ciphertext.is_empty() || !ciphertext.len().is_multiple_of(16) {
119            return Err(TokenError::DecryptionFailed);
120        }
121
122        let mut decrypted = match &self.mode {
123            AesMode::Aes128(aes) => aes.decrypt_cbc(ciphertext, &iv),
124            AesMode::Aes256(aes) => aes.decrypt_cbc(ciphertext, &iv),
125        }
126        .map_err(|_| TokenError::DecryptionFailed)?;
127
128        let unpadded_len = pkcs7::unpad(&decrypted, 16)
129            .map_err(|_| TokenError::DecryptionFailed)?
130            .len();
131        decrypted.truncate(unpadded_len);
132        Ok(decrypted)
133    }
134}
135
136#[cfg(test)]
137mod tests {
138    use super::*;
139    use crate::FixedRng;
140
141    #[test]
142    fn test_token_new_32byte_key() {
143        let key = [0u8; 32];
144        assert!(Token::new(&key).is_ok());
145    }
146
147    #[test]
148    fn test_token_new_64byte_key() {
149        let key = [0u8; 64];
150        assert!(Token::new(&key).is_ok());
151    }
152
153    #[test]
154    fn test_token_new_48byte_key() {
155        let key = [0u8; 48];
156        assert!(matches!(
157            Token::new(&key),
158            Err(TokenError::InvalidKeyLength)
159        ));
160    }
161
162    #[test]
163    fn test_token_roundtrip_32() {
164        let key = [0x42u8; 32];
165        let token = Token::new(&key).unwrap();
166        let mut rng = FixedRng::new(&[0xAA; 16]);
167        let plaintext = b"Hello, Reticulum!";
168        let encrypted = token.encrypt(plaintext, &mut rng);
169        let decrypted = token.decrypt(&encrypted).unwrap();
170        assert_eq!(decrypted, plaintext);
171    }
172
173    #[test]
174    fn test_token_roundtrip_64() {
175        let key = [0x42u8; 64];
176        let token = Token::new(&key).unwrap();
177        let mut rng = FixedRng::new(&[0xBB; 16]);
178        let plaintext = b"Hello, Reticulum!";
179        let encrypted = token.encrypt(plaintext, &mut rng);
180        let decrypted = token.decrypt(&encrypted).unwrap();
181        assert_eq!(decrypted, plaintext);
182    }
183
184    #[test]
185    fn test_token_hmac_reject_tampered() {
186        let key = [0x42u8; 64];
187        let token = Token::new(&key).unwrap();
188        let mut rng = FixedRng::new(&[0xCC; 16]);
189        let encrypted = token.encrypt(b"test", &mut rng);
190        let mut tampered = encrypted.clone();
191        tampered[20] ^= 0xFF; // flip a bit in ciphertext
192        assert!(token.decrypt(&tampered).is_err());
193    }
194
195    #[test]
196    fn token_hmac_rejects_mismatches_across_the_entire_tag() {
197        let key = [0x42u8; 64];
198        let token = Token::new(&key).unwrap();
199        let mut rng = FixedRng::new(&[0xCC; 16]);
200        let encrypted = token.encrypt(b"timing regression", &mut rng);
201
202        assert_eq!(token.verify_hmac(&encrypted), Ok(true));
203        for tag_index in 0..32 {
204            let mut corrupted = encrypted.clone();
205            let index = corrupted.len() - 32 + tag_index;
206            corrupted[index] ^= 0x01;
207            assert_eq!(
208                token.verify_hmac(&corrupted),
209                Ok(false),
210                "tag mismatch at byte {tag_index} must fail"
211            );
212        }
213    }
214
215    #[test]
216    fn test_token_decrypt_truncated() {
217        let key = [0x42u8; 64];
218        let token = Token::new(&key).unwrap();
219        assert!(matches!(
220            token.decrypt(&[0u8; 10]),
221            Err(TokenError::InvalidToken)
222        ));
223    }
224
225    fn sign_token_parts(key: &[u8], iv_and_ciphertext: &[u8]) -> Vec<u8> {
226        let signing_key = &key[..key.len() / 2];
227        let mut token = iv_and_ciphertext.to_vec();
228        token.extend_from_slice(&hmac_sha256(signing_key, iv_and_ciphertext));
229        token
230    }
231
232    #[test]
233    fn test_token_decrypt_rejects_misaligned_ciphertext_with_valid_hmac() {
234        for key in [&[0x42u8; 32][..], &[0x42u8; 64][..]] {
235            let token = Token::new(key).unwrap();
236            for ciphertext_len in [1usize, 15, 17, 31, 33] {
237                let mut parts = alloc::vec![0xAAu8; 16];
238                parts.extend(core::iter::repeat_n(0x55u8, ciphertext_len));
239                let signed = sign_token_parts(key, &parts);
240                assert_eq!(token.verify_hmac(&signed), Ok(true));
241                assert_eq!(
242                    token.decrypt(&signed),
243                    Err(TokenError::DecryptionFailed),
244                    "ciphertext length {ciphertext_len} must be rejected"
245                );
246            }
247        }
248    }
249
250    #[test]
251    fn test_token_well_formed_roundtrip_after_alignment_check() {
252        for key in [&[0x24u8; 32][..], &[0x24u8; 64][..]] {
253            let token = Token::new(key).unwrap();
254            for plaintext_len in [0usize, 1, 15, 16, 17, 32, 100] {
255                let plaintext = alloc::vec![0x5Au8; plaintext_len];
256                let encrypted = token.encrypt_with_iv(&plaintext, &[0x11; 16]);
257                assert_eq!((encrypted.len() - TOKEN_OVERHEAD) % 16, 0);
258                assert_eq!(token.decrypt(&encrypted).unwrap(), plaintext);
259            }
260        }
261    }
262
263    #[test]
264    fn test_token_overhead() {
265        assert_eq!(TOKEN_OVERHEAD, 48);
266    }
267}
rns_crypto/token.rs

rns_crypto/
token.rs
