Skip to main content

shield_core/
identity.rs

1//! Identity and SSO without public-key cryptography.
2//!
3//! Provides identity management, session tokens, and service tokens.
4
5// Token timestamps and indices use intentional truncation for compact encoding
6#![allow(clippy::cast_possible_truncation)]
7
8use base64::Engine;
9use ring::hmac;
10use std::collections::HashMap;
11use std::num::NonZeroU32;
12use std::time::{SystemTime, UNIX_EPOCH};
13use subtle::ConstantTimeEq;
14use zeroize::Zeroize;
15
16use crate::error::{Result, ShieldError};
17
18/// Generate keystream using HMAC-SHA256 (keyed PRF).
19fn generate_keystream(key: &[u8], nonce: &[u8], length: usize) -> Result<Vec<u8>> {
20    let num_blocks = length.div_ceil(32);
21    if u32::try_from(num_blocks).is_err() {
22        return Err(ShieldError::StreamError(
23            "keystream too long: counter overflow".into(),
24        ));
25    }
26    let mut keystream = Vec::with_capacity(num_blocks * 32);
27    let hmac_key = hmac::Key::new(hmac::HMAC_SHA256, key);
28
29    for i in 0..num_blocks {
30        let counter = (i as u32).to_le_bytes();
31        let mut data = Vec::with_capacity(nonce.len() + 4);
32        data.extend_from_slice(nonce);
33        data.extend_from_slice(&counter);
34
35        let tag = hmac::sign(&hmac_key, &data);
36        keystream.extend_from_slice(tag.as_ref());
37    }
38
39    keystream.truncate(length);
40    Ok(keystream)
41}
42
43/// User identity.
44#[derive(Clone)]
45pub struct Identity {
46    pub user_id: String,
47    pub display_name: String,
48    pub verification_key: [u8; 32],
49    pub attributes: HashMap<String, String>,
50    pub created_at: u64,
51}
52
53/// Session information.
54#[derive(Clone)]
55pub struct Session {
56    pub user_id: String,
57    pub permissions: Vec<String>,
58    pub expires_at: Option<u64>,
59    pub attributes: HashMap<String, String>,
60}
61
62impl Session {
63    /// Check if session is expired.
64    #[must_use]
65    pub fn is_expired(&self) -> bool {
66        match self.expires_at {
67            None => false,
68            Some(expires) => {
69                let now = SystemTime::now()
70                    .duration_since(UNIX_EPOCH)
71                    .map_or(0, |d| d.as_secs());
72                now > expires
73            }
74        }
75    }
76
77    /// Check if session has permission.
78    #[must_use]
79    pub fn has_permission(&self, permission: &str) -> bool {
80        self.permissions.contains(&permission.to_string())
81    }
82}
83
84/// Stored user data.
85struct UserData {
86    password_hash: [u8; 32],
87    salt: [u8; 16],
88    identity: Identity,
89}
90
91/// Identity provider for managing users and sessions.
92pub struct IdentityProvider {
93    master_key: [u8; 32],
94    token_ttl: u64,
95    users: HashMap<String, UserData>,
96}
97
98impl IdentityProvider {
99    const ITERATIONS: u32 = 100_000;
100
101    /// Create new identity provider.
102    #[must_use]
103    pub fn new(master_key: [u8; 32], token_ttl: u64) -> Self {
104        Self {
105            master_key,
106            token_ttl: if token_ttl == 0 { 3600 } else { token_ttl },
107            users: HashMap::new(),
108        }
109    }
110
111    /// Derive key for specific purpose using HMAC-SHA256 (keyed PRF).
112    fn derive_key(&self, purpose: &str) -> [u8; 32] {
113        let hmac_key = hmac::Key::new(hmac::HMAC_SHA256, &self.master_key);
114        let tag = hmac::sign(&hmac_key, purpose.as_bytes());
115        let mut key = [0u8; 32];
116        key.copy_from_slice(&tag.as_ref()[..32]);
117        key
118    }
119
120    /// Derive separated encryption and MAC subkeys for token operations.
121    /// Prevents key reuse between encryption and authentication.
122    fn derive_token_subkeys(&self, purpose: &str) -> ([u8; 32], [u8; 32]) {
123        let base_key = self.derive_key(purpose);
124        let hmac_key = hmac::Key::new(hmac::HMAC_SHA256, &base_key);
125
126        let enc_tag = hmac::sign(&hmac_key, b"shield-token-encrypt");
127        let mut enc_key = [0u8; 32];
128        enc_key.copy_from_slice(&enc_tag.as_ref()[..32]);
129
130        let mac_tag = hmac::sign(&hmac_key, b"shield-token-authenticate");
131        let mut mac_key = [0u8; 32];
132        mac_key.copy_from_slice(&mac_tag.as_ref()[..32]);
133
134        (enc_key, mac_key)
135    }
136
137    /// Register a new user.
138    pub fn register(
139        &mut self,
140        user_id: &str,
141        password: &str,
142        display_name: Option<&str>,
143        attributes: HashMap<String, String>,
144    ) -> Result<Identity> {
145        if self.users.contains_key(user_id) {
146            return Err(ShieldError::UserExists(user_id.to_string()));
147        }
148
149        let salt: [u8; 16] = crate::random::random_bytes()?;
150
151        let mut password_hash = [0u8; 32];
152        ring::pbkdf2::derive(
153            ring::pbkdf2::PBKDF2_HMAC_SHA256,
154            NonZeroU32::new(Self::ITERATIONS).unwrap(),
155            &salt,
156            password.as_bytes(),
157            &mut password_hash,
158        );
159
160        // Generate verification key using HMAC-SHA256 (keyed PRF)
161        let verify_key = self.derive_key("verify");
162        let vk_hmac = hmac::Key::new(hmac::HMAC_SHA256, &verify_key);
163        let vk_tag = hmac::sign(&vk_hmac, user_id.as_bytes());
164        let mut verification_key = [0u8; 32];
165        verification_key.copy_from_slice(&vk_tag.as_ref()[..32]);
166
167        let now = SystemTime::now()
168            .duration_since(UNIX_EPOCH)
169            .map_or(0, |d| d.as_secs());
170
171        let identity = Identity {
172            user_id: user_id.to_string(),
173            display_name: display_name.unwrap_or(user_id).to_string(),
174            verification_key,
175            attributes, // Consume owned value, no clone needed
176            created_at: now,
177        };
178
179        self.users.insert(
180            user_id.to_string(),
181            UserData {
182                password_hash,
183                salt,
184                identity: identity.clone(),
185            },
186        );
187
188        Ok(identity)
189    }
190
191    /// Authenticate user and return session token.
192    ///
193    /// Runs PBKDF2 even for non-existent users to prevent timing-based
194    /// user enumeration (CWE-203).
195    #[must_use]
196    pub fn authenticate(
197        &self,
198        user_id: &str,
199        password: &str,
200        permissions: &[String],
201        ttl: Option<u64>,
202    ) -> Option<String> {
203        // Use real salt if user exists, otherwise derive a stable dummy salt
204        // to ensure constant-time behavior regardless of user existence.
205        let dummy_salt = self.derive_key("dummy-salt");
206        let (salt, expected_hash, user_exists) = match self.users.get(user_id) {
207            Some(user) => (user.salt, user.password_hash, true),
208            None => {
209                let mut dummy = [0u8; 16];
210                dummy.copy_from_slice(&dummy_salt[..16]);
211                (dummy, [0u8; 32], false)
212            }
213        };
214
215        let mut password_hash = [0u8; 32];
216        ring::pbkdf2::derive(
217            ring::pbkdf2::PBKDF2_HMAC_SHA256,
218            NonZeroU32::new(Self::ITERATIONS).unwrap(),
219            &salt,
220            password.as_bytes(),
221            &mut password_hash,
222        );
223
224        if !user_exists || password_hash.ct_eq(&expected_hash).unwrap_u8() != 1 {
225            return None;
226        }
227
228        self.create_token(user_id, permissions, ttl.unwrap_or(self.token_ttl))
229            .ok()
230    }
231
232    /// Create session token.
233    fn create_token(&self, user_id: &str, permissions: &[String], ttl: u64) -> Result<String> {
234        let nonce: [u8; 16] = crate::random::random_bytes()?;
235
236        let now = SystemTime::now()
237            .duration_since(UNIX_EPOCH)
238            .map_or(0, |d| d.as_secs());
239        let expires_at = now + ttl;
240
241        // Serialize token data
242        let user_id_bytes = user_id.as_bytes();
243        let perms_json = serde_json::to_string(permissions).unwrap_or_default();
244        let perms_bytes = perms_json.as_bytes();
245
246        let mut token_data = Vec::new();
247        token_data.extend_from_slice(&(user_id_bytes.len() as u16).to_le_bytes());
248        token_data.extend_from_slice(user_id_bytes);
249        token_data.extend_from_slice(&(perms_bytes.len() as u16).to_le_bytes());
250        token_data.extend_from_slice(perms_bytes);
251        token_data.extend_from_slice(&expires_at.to_le_bytes());
252
253        // Derive separate encryption and MAC keys
254        let (enc_key, mac_key) = self.derive_token_subkeys("session");
255
256        // Encrypt with enc_key
257        let keystream = generate_keystream(&enc_key, &nonce, token_data.len())?;
258        let encrypted: Vec<u8> = token_data
259            .iter()
260            .zip(keystream.iter())
261            .map(|(p, k)| p ^ k)
262            .collect();
263
264        // MAC with mac_key
265        let hmac_key = hmac::Key::new(hmac::HMAC_SHA256, &mac_key);
266        let mut hmac_data = Vec::with_capacity(16 + encrypted.len());
267        hmac_data.extend_from_slice(&nonce);
268        hmac_data.extend_from_slice(&encrypted);
269        let tag = hmac::sign(&hmac_key, &hmac_data);
270
271        let mut result = Vec::with_capacity(16 + encrypted.len() + 16);
272        result.extend_from_slice(&nonce);
273        result.extend_from_slice(&encrypted);
274        result.extend_from_slice(&tag.as_ref()[..16]);
275
276        Ok(base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(&result))
277    }
278
279    /// Validate session token.
280    #[must_use]
281    pub fn validate_token(&self, token: &str) -> Option<Session> {
282        let data = base64::engine::general_purpose::URL_SAFE_NO_PAD
283            .decode(token)
284            .ok()?;
285
286        if data.len() < 34 {
287            return None;
288        }
289
290        let nonce = &data[..16];
291        let encrypted = &data[16..data.len() - 16];
292        let mac = &data[data.len() - 16..];
293
294        // Derive separate encryption and MAC keys
295        let (enc_key, mac_key) = self.derive_token_subkeys("session");
296
297        // Verify MAC with mac_key
298        let hmac_key = hmac::Key::new(hmac::HMAC_SHA256, &mac_key);
299        let mut hmac_data = Vec::with_capacity(16 + encrypted.len());
300        hmac_data.extend_from_slice(nonce);
301        hmac_data.extend_from_slice(encrypted);
302        let expected_tag = hmac::sign(&hmac_key, &hmac_data);
303
304        if mac.ct_eq(&expected_tag.as_ref()[..16]).unwrap_u8() != 1 {
305            return None;
306        }
307
308        // Decrypt with enc_key
309        let keystream = generate_keystream(&enc_key, nonce, encrypted.len()).ok()?;
310        let token_data: Vec<u8> = encrypted
311            .iter()
312            .zip(keystream.iter())
313            .map(|(c, k)| c ^ k)
314            .collect();
315
316        // Parse with bounds checks to prevent panics on malformed data
317        if token_data.len() < 2 {
318            return None;
319        }
320        let user_id_len = u16::from_le_bytes([token_data[0], token_data[1]]) as usize;
321        if token_data.len() < 2 + user_id_len + 2 {
322            return None;
323        }
324        let user_id = String::from_utf8(token_data[2..2 + user_id_len].to_vec()).ok()?;
325
326        let offset = 2 + user_id_len;
327        let perms_len = u16::from_le_bytes([token_data[offset], token_data[offset + 1]]) as usize;
328        if token_data.len() < offset + 2 + perms_len + 8 {
329            return None;
330        }
331        let perms_json =
332            String::from_utf8(token_data[offset + 2..offset + 2 + perms_len].to_vec()).ok()?;
333        let permissions: Vec<String> = serde_json::from_str(&perms_json).ok()?;
334
335        let exp_offset = offset + 2 + perms_len;
336        let expires_at =
337            u64::from_le_bytes(token_data[exp_offset..exp_offset + 8].try_into().ok()?);
338
339        // Reject tokens for revoked/unregistered users
340        if !self.users.contains_key(&user_id) {
341            return None;
342        }
343
344        let session = Session {
345            user_id,
346            permissions,
347            expires_at: Some(expires_at),
348            attributes: HashMap::new(),
349        };
350
351        if session.is_expired() {
352            return None;
353        }
354
355        Some(session)
356    }
357
358    /// Create service-specific token.
359    #[must_use]
360    pub fn create_service_token(
361        &self,
362        session_token: &str,
363        service: &str,
364        permissions: &[String],
365        ttl: u64,
366    ) -> Option<String> {
367        let session = self.validate_token(session_token)?;
368
369        let nonce: [u8; 16] = crate::random::random_bytes().ok()?;
370
371        let now = SystemTime::now()
372            .duration_since(UNIX_EPOCH)
373            .map_or(0, |d| d.as_secs());
374        let expires_at = now + ttl;
375
376        // Serialize
377        let user_id_bytes = session.user_id.as_bytes();
378        let service_bytes = service.as_bytes();
379        let perms_json = serde_json::to_string(permissions).unwrap_or_default();
380        let perms_bytes = perms_json.as_bytes();
381
382        let mut token_data = Vec::new();
383        token_data.extend_from_slice(&(user_id_bytes.len() as u16).to_le_bytes());
384        token_data.extend_from_slice(user_id_bytes);
385        token_data.extend_from_slice(&(service_bytes.len() as u16).to_le_bytes());
386        token_data.extend_from_slice(service_bytes);
387        token_data.extend_from_slice(&(perms_bytes.len() as u16).to_le_bytes());
388        token_data.extend_from_slice(perms_bytes);
389        token_data.extend_from_slice(&expires_at.to_le_bytes());
390
391        // Derive separate enc/mac keys for service-specific token
392        let (enc_key, mac_key) = self.derive_token_subkeys(&format!("service:{service}"));
393        let keystream = generate_keystream(&enc_key, &nonce, token_data.len()).ok()?;
394        let encrypted: Vec<u8> = token_data
395            .iter()
396            .zip(keystream.iter())
397            .map(|(p, k)| p ^ k)
398            .collect();
399
400        let hmac_key = hmac::Key::new(hmac::HMAC_SHA256, &mac_key);
401        let mut hmac_data = Vec::with_capacity(16 + encrypted.len());
402        hmac_data.extend_from_slice(&nonce);
403        hmac_data.extend_from_slice(&encrypted);
404        let tag = hmac::sign(&hmac_key, &hmac_data);
405
406        let mut result = Vec::with_capacity(16 + encrypted.len() + 16);
407        result.extend_from_slice(&nonce);
408        result.extend_from_slice(&encrypted);
409        result.extend_from_slice(&tag.as_ref()[..16]);
410
411        Some(base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(&result))
412    }
413
414    /// Validate service token.
415    #[must_use]
416    pub fn validate_service_token(&self, token: &str, service: &str) -> Option<Session> {
417        let data = base64::engine::general_purpose::URL_SAFE_NO_PAD
418            .decode(token)
419            .ok()?;
420
421        if data.len() < 34 {
422            return None;
423        }
424
425        let nonce = &data[..16];
426        let encrypted = &data[16..data.len() - 16];
427        let mac = &data[data.len() - 16..];
428
429        // Derive separate enc/mac keys for service token
430        let (enc_key, mac_key) = self.derive_token_subkeys(&format!("service:{service}"));
431
432        // Verify MAC with mac_key
433        let hmac_key = hmac::Key::new(hmac::HMAC_SHA256, &mac_key);
434        let mut hmac_data = Vec::with_capacity(16 + encrypted.len());
435        hmac_data.extend_from_slice(nonce);
436        hmac_data.extend_from_slice(encrypted);
437        let expected_tag = hmac::sign(&hmac_key, &hmac_data);
438
439        if mac.ct_eq(&expected_tag.as_ref()[..16]).unwrap_u8() != 1 {
440            return None;
441        }
442
443        // Decrypt with enc_key
444        let keystream = generate_keystream(&enc_key, nonce, encrypted.len()).ok()?;
445        let token_data: Vec<u8> = encrypted
446            .iter()
447            .zip(keystream.iter())
448            .map(|(c, k)| c ^ k)
449            .collect();
450
451        // Parse with bounds checks
452        let mut offset = 0;
453        if token_data.len() < offset + 2 {
454            return None;
455        }
456        let user_id_len = u16::from_le_bytes([token_data[offset], token_data[offset + 1]]) as usize;
457        offset += 2;
458        if token_data.len() < offset + user_id_len + 2 {
459            return None;
460        }
461        let user_id = String::from_utf8(token_data[offset..offset + user_id_len].to_vec()).ok()?;
462        offset += user_id_len;
463
464        let service_len = u16::from_le_bytes([token_data[offset], token_data[offset + 1]]) as usize;
465        offset += 2;
466        if token_data.len() < offset + service_len + 2 {
467            return None;
468        }
469        let token_service =
470            String::from_utf8(token_data[offset..offset + service_len].to_vec()).ok()?;
471        offset += service_len;
472
473        if token_service != service {
474            return None;
475        }
476
477        let perms_len = u16::from_le_bytes([token_data[offset], token_data[offset + 1]]) as usize;
478        offset += 2;
479        if token_data.len() < offset + perms_len + 8 {
480            return None;
481        }
482        let perms_json = String::from_utf8(token_data[offset..offset + perms_len].to_vec()).ok()?;
483        let permissions: Vec<String> = serde_json::from_str(&perms_json).ok()?;
484        offset += perms_len;
485
486        let expires_at = u64::from_le_bytes(token_data[offset..offset + 8].try_into().ok()?);
487
488        let session = Session {
489            user_id,
490            permissions,
491            expires_at: Some(expires_at),
492            attributes: HashMap::new(),
493        };
494
495        if session.is_expired() {
496            return None;
497        }
498
499        Some(session)
500    }
501
502    /// Refresh session token.
503    #[must_use]
504    pub fn refresh_token(&self, token: &str) -> Option<String> {
505        let session = self.validate_token(token)?;
506        self.create_token(&session.user_id, &session.permissions, self.token_ttl)
507            .ok()
508    }
509
510    /// Get user identity.
511    #[must_use]
512    pub fn get_identity(&self, user_id: &str) -> Option<&Identity> {
513        self.users.get(user_id).map(|u| &u.identity)
514    }
515
516    /// Revoke user.
517    pub fn revoke_user(&mut self, user_id: &str) {
518        if let Some(mut user) = self.users.remove(user_id) {
519            user.password_hash.zeroize();
520            user.salt.zeroize();
521        }
522    }
523}
524
525/// Derive separated encryption and MAC subkeys for session operations.
526fn derive_session_subkeys(key: &[u8; 32]) -> ([u8; 32], [u8; 32]) {
527    let hmac_key = hmac::Key::new(hmac::HMAC_SHA256, key);
528
529    let enc_tag = hmac::sign(&hmac_key, b"shield-session-encrypt");
530    let mut enc_key = [0u8; 32];
531    enc_key.copy_from_slice(&enc_tag.as_ref()[..32]);
532
533    let mac_tag = hmac::sign(&hmac_key, b"shield-session-authenticate");
534    let mut mac_key = [0u8; 32];
535    mac_key.copy_from_slice(&mac_tag.as_ref()[..32]);
536
537    (enc_key, mac_key)
538}
539
540/// Secure session with automatic key rotation.
541pub struct SecureSession {
542    master_key: [u8; 32],
543    rotation_interval: u64,
544    max_old_keys: usize,
545    key_version: u32,
546    keys: HashMap<u32, [u8; 32]>,
547    last_rotation: u64,
548}
549
550impl SecureSession {
551    /// Create new secure session.
552    #[must_use]
553    pub fn new(master_key: [u8; 32], rotation_interval: u64, max_old_keys: usize) -> Self {
554        let now = SystemTime::now()
555            .duration_since(UNIX_EPOCH)
556            .map_or(0, |d| d.as_secs());
557
558        let key = Self::derive_session_key(&master_key, 1);
559        let mut keys = HashMap::new();
560        keys.insert(1, key);
561
562        Self {
563            master_key,
564            rotation_interval: if rotation_interval == 0 {
565                3600
566            } else {
567                rotation_interval
568            },
569            max_old_keys: if max_old_keys == 0 { 3 } else { max_old_keys },
570            key_version: 1,
571            keys,
572            last_rotation: now,
573        }
574    }
575
576    fn derive_session_key(master_key: &[u8; 32], version: u32) -> [u8; 32] {
577        let hmac_key = hmac::Key::new(hmac::HMAC_SHA256, master_key);
578        let tag = hmac::sign(&hmac_key, format!("session:{version}").as_bytes());
579        let mut key = [0u8; 32];
580        key.copy_from_slice(&tag.as_ref()[..32]);
581        key
582    }
583
584    fn maybe_rotate(&mut self) {
585        let now = SystemTime::now()
586            .duration_since(UNIX_EPOCH)
587            .map_or(0, |d| d.as_secs());
588
589        if now - self.last_rotation >= self.rotation_interval {
590            self.key_version += 1;
591            let new_key = Self::derive_session_key(&self.master_key, self.key_version);
592            self.keys.insert(self.key_version, new_key);
593            self.last_rotation = now;
594
595            // Prune old keys (zeroize before removing)
596            let mut versions: Vec<u32> = self.keys.keys().copied().collect();
597            versions.sort_by(|a, b| b.cmp(a));
598            for v in versions.into_iter().skip(self.max_old_keys + 1) {
599                if let Some(key) = self.keys.get_mut(&v) {
600                    key.zeroize();
601                }
602                self.keys.remove(&v);
603            }
604        }
605    }
606
607    /// Encrypt session data.
608    pub fn encrypt(&mut self, data: &[u8]) -> Result<Vec<u8>> {
609        self.maybe_rotate();
610
611        let key = self
612            .keys
613            .get(&self.key_version)
614            .ok_or(ShieldError::UnknownVersion(self.key_version))?;
615        let (enc_key, mac_key) = derive_session_subkeys(key);
616        let nonce: [u8; 16] = crate::random::random_bytes()?;
617
618        let keystream = generate_keystream(&enc_key, &nonce, data.len())?;
619        let ciphertext: Vec<u8> = data
620            .iter()
621            .zip(keystream.iter())
622            .map(|(p, k)| p ^ k)
623            .collect();
624
625        let version_bytes = self.key_version.to_le_bytes();
626
627        let hmac_signing_key = hmac::Key::new(hmac::HMAC_SHA256, &mac_key);
628        let mut hmac_data = Vec::with_capacity(4 + 16 + ciphertext.len());
629        hmac_data.extend_from_slice(&version_bytes);
630        hmac_data.extend_from_slice(&nonce);
631        hmac_data.extend_from_slice(&ciphertext);
632        let tag = hmac::sign(&hmac_signing_key, &hmac_data);
633
634        let mut result = Vec::with_capacity(4 + 16 + ciphertext.len() + 16);
635        result.extend_from_slice(&version_bytes);
636        result.extend_from_slice(&nonce);
637        result.extend_from_slice(&ciphertext);
638        result.extend_from_slice(&tag.as_ref()[..16]);
639
640        Ok(result)
641    }
642
643    /// Decrypt session data.
644    pub fn decrypt(&mut self, encrypted: &[u8]) -> Option<Vec<u8>> {
645        self.maybe_rotate();
646
647        if encrypted.len() < 36 {
648            return None;
649        }
650
651        let version = u32::from_le_bytes(encrypted[..4].try_into().ok()?);
652        let nonce = &encrypted[4..20];
653        let ciphertext = &encrypted[20..encrypted.len() - 16];
654        let mac = &encrypted[encrypted.len() - 16..];
655
656        let key = self.keys.get(&version)?;
657        let (enc_key, mac_key) = derive_session_subkeys(key);
658
659        // Verify MAC with mac_key
660        let hmac_signing_key = hmac::Key::new(hmac::HMAC_SHA256, &mac_key);
661        let expected_tag = hmac::sign(&hmac_signing_key, &encrypted[..encrypted.len() - 16]);
662
663        if mac.ct_eq(&expected_tag.as_ref()[..16]).unwrap_u8() != 1 {
664            return None;
665        }
666
667        let keystream = generate_keystream(&enc_key, nonce, ciphertext.len()).ok()?;
668        let plaintext: Vec<u8> = ciphertext
669            .iter()
670            .zip(keystream.iter())
671            .map(|(c, k)| c ^ k)
672            .collect();
673
674        Some(plaintext)
675    }
676
677    /// Get current key version (for testing).
678    #[must_use]
679    pub fn key_version(&self) -> u32 {
680        self.key_version
681    }
682}
683
684impl Drop for IdentityProvider {
685    fn drop(&mut self) {
686        self.master_key.zeroize();
687        for user in self.users.values_mut() {
688            user.password_hash.zeroize();
689            user.salt.zeroize();
690        }
691    }
692}
693
694impl Drop for SecureSession {
695    fn drop(&mut self) {
696        self.master_key.zeroize();
697        for key in self.keys.values_mut() {
698            key.zeroize();
699        }
700    }
701}
702
703#[cfg(test)]
704mod tests {
705    use super::*;
706
707    #[test]
708    fn test_register_user() {
709        let mut provider = IdentityProvider::new([0u8; 32], 3600);
710        let identity = provider
711            .register("alice", "password123", Some("Alice Smith"), HashMap::new())
712            .unwrap();
713
714        assert_eq!(identity.user_id, "alice");
715        assert_eq!(identity.display_name, "Alice Smith");
716    }
717
718    #[test]
719    fn test_register_duplicate() {
720        let mut provider = IdentityProvider::new([0u8; 32], 3600);
721        provider
722            .register("alice", "password", None, HashMap::new())
723            .unwrap();
724        assert!(provider
725            .register("alice", "password2", None, HashMap::new())
726            .is_err());
727    }
728
729    #[test]
730    fn test_authenticate() {
731        let mut provider = IdentityProvider::new([0u8; 32], 3600);
732        provider
733            .register("alice", "password123", None, HashMap::new())
734            .unwrap();
735
736        let token = provider.authenticate("alice", "password123", &[], None);
737        assert!(token.is_some());
738    }
739
740    #[test]
741    fn test_authenticate_wrong_password() {
742        let mut provider = IdentityProvider::new([0u8; 32], 3600);
743        provider
744            .register("alice", "password123", None, HashMap::new())
745            .unwrap();
746
747        let token = provider.authenticate("alice", "wrongpassword", &[], None);
748        assert!(token.is_none());
749    }
750
751    #[test]
752    fn test_validate_token() {
753        let mut provider = IdentityProvider::new([0u8; 32], 3600);
754        provider
755            .register("alice", "password", None, HashMap::new())
756            .unwrap();
757        let token = provider
758            .authenticate("alice", "password", &[], None)
759            .unwrap();
760
761        let session = provider.validate_token(&token);
762        assert!(session.is_some());
763        assert_eq!(session.unwrap().user_id, "alice");
764    }
765
766    #[test]
767    fn test_service_token() {
768        let mut provider = IdentityProvider::new([0u8; 32], 3600);
769        provider
770            .register("alice", "password", None, HashMap::new())
771            .unwrap();
772        let session_token = provider
773            .authenticate("alice", "password", &[], None)
774            .unwrap();
775
776        let service_token = provider
777            .create_service_token(
778                &session_token,
779                "api.example.com",
780                &["read".to_string()],
781                300,
782            )
783            .unwrap();
784
785        let session = provider.validate_service_token(&service_token, "api.example.com");
786        assert!(session.is_some());
787        assert_eq!(session.as_ref().unwrap().user_id, "alice");
788        assert!(session.unwrap().has_permission("read"));
789    }
790
791    #[test]
792    fn test_service_token_wrong_service() {
793        let mut provider = IdentityProvider::new([0u8; 32], 3600);
794        provider
795            .register("alice", "password", None, HashMap::new())
796            .unwrap();
797        let session_token = provider
798            .authenticate("alice", "password", &[], None)
799            .unwrap();
800        let service_token = provider
801            .create_service_token(&session_token, "api.example.com", &[], 300)
802            .unwrap();
803
804        let session = provider.validate_service_token(&service_token, "other.example.com");
805        assert!(session.is_none());
806    }
807
808    #[test]
809    fn test_secure_session() {
810        let mut session = SecureSession::new([0u8; 32], 3600, 3);
811        let plaintext = b"session data";
812        let encrypted = session.encrypt(plaintext).unwrap();
813        let decrypted = session.decrypt(&encrypted).unwrap();
814        assert_eq!(plaintext.as_slice(), decrypted.as_slice());
815    }
816
817    #[test]
818    fn test_secure_session_tampered() {
819        let mut session = SecureSession::new([0u8; 32], 3600, 3);
820        let mut encrypted = session.encrypt(b"data").unwrap();
821        encrypted[20] ^= 0xFF;
822        assert!(session.decrypt(&encrypted).is_none());
823    }
824}