Skip to main content

mockforge_security_core/encryption/
derivation.rs

1//! Key derivation functions (Argon2, PBKDF2)
2//!
3//! This module provides secure key derivation functions for generating
4//! encryption keys from passwords and other secret material.
5
6use crate::encryption::algorithms::{EncryptionAlgorithm, EncryptionKey};
7use crate::encryption::errors::{EncryptionError, EncryptionResult};
8use argon2::{
9    password_hash::{PasswordHasher, SaltString},
10    Algorithm, Argon2, Params, Version,
11};
12use pbkdf2::pbkdf2_hmac;
13use rand::{rng, Rng};
14use sha2::Sha256;
15
16/// Key derivation method
17#[derive(Debug, Clone)]
18pub enum KeyDerivationMethod {
19    /// Argon2id (recommended for passwords)
20    Argon2 {
21        memory_kib: u32,
22        iterations: u32,
23        parallelism: u32,
24    },
25    /// PBKDF2-HMAC-SHA256
26    #[allow(dead_code)]
27    Pbkdf2 { iterations: u32 },
28}
29
30/// Key derivation manager
31#[derive(Debug, Clone)]
32pub struct KeyDerivationManager {
33    /// Default Argon2 parameters
34    default_argon2_params: Argon2Params,
35}
36
37/// Argon2 parameters
38#[derive(Debug, Clone)]
39pub struct Argon2Params {
40    pub memory_kib: u32,
41    pub iterations: u32,
42    pub parallelism: u32,
43}
44
45impl Default for Argon2Params {
46    fn default() -> Self {
47        Self {
48            memory_kib: 19456, // 19 MiB
49            iterations: 2,
50            parallelism: 1,
51        }
52    }
53}
54
55impl KeyDerivationManager {
56    /// Create a new key derivation manager
57    pub fn new() -> Self {
58        Self {
59            default_argon2_params: Argon2Params::default(),
60        }
61    }
62
63    /// Derive a master key from a password (synchronous version)
64    ///
65    /// Note: This is CPU-intensive. Use `derive_master_key_async` when calling from async context.
66    pub fn derive_master_key(&self, password: &str) -> EncryptionResult<EncryptionKey> {
67        self.derive_key(
68            password.as_bytes(),
69            KeyDerivationMethod::Argon2 {
70                memory_kib: self.default_argon2_params.memory_kib,
71                iterations: self.default_argon2_params.iterations,
72                parallelism: self.default_argon2_params.parallelism,
73            },
74            "master_key_salt",
75            EncryptionAlgorithm::Aes256Gcm,
76        )
77    }
78
79    /// Derive a master key from a password (async version using spawn_blocking)
80    ///
81    /// This method offloads the CPU-intensive Argon2 computation to a blocking thread pool.
82    #[allow(dead_code)]
83    pub async fn derive_master_key_async(
84        &self,
85        password: String,
86    ) -> EncryptionResult<EncryptionKey> {
87        let params = self.default_argon2_params.clone();
88        tokio::task::spawn_blocking(move || {
89            let manager = Self::new();
90            manager.derive_key(
91                password.as_bytes(),
92                KeyDerivationMethod::Argon2 {
93                    memory_kib: params.memory_kib,
94                    iterations: params.iterations,
95                    parallelism: params.parallelism,
96                },
97                "master_key_salt",
98                EncryptionAlgorithm::Aes256Gcm,
99            )
100        })
101        .await
102        .map_err(|e| EncryptionError::key_derivation_failed(format!("Task join error: {}", e)))?
103    }
104
105    /// Derive a workspace key from workspace ID and master key
106    #[allow(dead_code)]
107    pub fn derive_workspace_key(
108        &self,
109        master_key: &EncryptionKey,
110        workspace_id: &str,
111    ) -> EncryptionResult<EncryptionKey> {
112        let master_bytes = master_key.as_bytes();
113        let workspace_bytes = workspace_id.as_bytes();
114
115        let mut derived_key = vec![0u8; 32];
116        pbkdf2_hmac::<Sha256>(
117            master_bytes,
118            workspace_bytes,
119            10000, // iterations
120            &mut derived_key,
121        );
122
123        EncryptionKey::new(derived_key, EncryptionAlgorithm::Aes256Gcm)
124    }
125
126    /// Derive a key using the specified method
127    pub fn derive_key(
128        &self,
129        secret: &[u8],
130        method: KeyDerivationMethod,
131        salt: &str,
132        algorithm: EncryptionAlgorithm,
133    ) -> EncryptionResult<EncryptionKey> {
134        match method {
135            KeyDerivationMethod::Argon2 {
136                memory_kib,
137                iterations,
138                parallelism,
139            } => {
140                self.derive_key_argon2(secret, salt, memory_kib, iterations, parallelism, algorithm)
141            }
142            KeyDerivationMethod::Pbkdf2 { iterations } => {
143                self.derive_key_pbkdf2(secret, salt, iterations, algorithm)
144            }
145        }
146    }
147
148    /// Derive key using Argon2
149    fn derive_key_argon2(
150        &self,
151        secret: &[u8],
152        _salt: &str,
153        memory_kib: u32,
154        iterations: u32,
155        parallelism: u32,
156        algorithm: EncryptionAlgorithm,
157    ) -> EncryptionResult<EncryptionKey> {
158        let salt = SaltString::encode_b64(b"randomsalt12345678901234567890123456789012")
159            .expect("encoding a fixed-length literal salt into base64 cannot fail");
160
161        let params = Params::new(
162            memory_kib,
163            iterations,
164            parallelism,
165            Some(32), // output length
166        )
167        .map_err(|e| EncryptionError::key_derivation_failed(e.to_string()))?;
168
169        let argon2 = Argon2::new(Algorithm::Argon2id, Version::V0x13, params);
170        let password_hash = argon2
171            .hash_password(secret, &salt)
172            .map_err(|e| EncryptionError::key_derivation_failed(e.to_string()))?;
173
174        let hash_binding = password_hash
175            .hash
176            .expect("argon2 hash_password always populates the hash output on success");
177        let hash_bytes = hash_binding.as_bytes();
178        let key_bytes: Vec<u8> = hash_bytes.to_vec();
179
180        // Take only the required number of bytes for the algorithm
181        let key_len = match algorithm {
182            EncryptionAlgorithm::Aes256Gcm => 32,
183            EncryptionAlgorithm::ChaCha20Poly1305 => 32,
184        };
185
186        let final_key_bytes = if key_bytes.len() >= key_len {
187            key_bytes[..key_len].to_vec()
188        } else {
189            return Err(EncryptionError::key_derivation_failed(
190                "Derived key too short for algorithm",
191            ));
192        };
193
194        EncryptionKey::new(final_key_bytes, algorithm)
195    }
196
197    /// Derive key using PBKDF2
198    fn derive_key_pbkdf2(
199        &self,
200        secret: &[u8],
201        salt: &str,
202        iterations: u32,
203        algorithm: EncryptionAlgorithm,
204    ) -> EncryptionResult<EncryptionKey> {
205        let salt_bytes = salt.as_bytes();
206        let key_len = match algorithm {
207            EncryptionAlgorithm::Aes256Gcm => 32,
208            EncryptionAlgorithm::ChaCha20Poly1305 => 32,
209        };
210
211        let mut derived_key = vec![0u8; key_len];
212        pbkdf2_hmac::<Sha256>(secret, salt_bytes, iterations, &mut derived_key);
213
214        EncryptionKey::new(derived_key, algorithm)
215    }
216
217    /// Verify a password against a derived key (synchronous version)
218    ///
219    /// Note: This is CPU-intensive. Use `verify_password_async` when calling from async context.
220    #[allow(dead_code)]
221    pub fn verify_password(
222        &self,
223        password: &str,
224        expected_key: &EncryptionKey,
225    ) -> EncryptionResult<bool> {
226        let derived_key = self.derive_master_key(password)?;
227
228        Ok(derived_key.as_bytes() == expected_key.as_bytes())
229    }
230
231    /// Verify a password against a derived key (async version using spawn_blocking)
232    ///
233    /// This method offloads the CPU-intensive Argon2 computation to a blocking thread pool.
234    #[allow(dead_code)]
235    pub async fn verify_password_async(
236        &self,
237        password: String,
238        expected_key: EncryptionKey,
239    ) -> EncryptionResult<bool> {
240        let params = self.default_argon2_params.clone();
241        tokio::task::spawn_blocking(move || {
242            let manager = KeyDerivationManager {
243                default_argon2_params: params,
244            };
245            let derived_key = manager.derive_master_key(&password)?;
246            Ok(derived_key.as_bytes() == expected_key.as_bytes())
247        })
248        .await
249        .map_err(|e| EncryptionError::key_derivation_failed(format!("Task join error: {}", e)))?
250    }
251
252    /// Generate a secure random salt
253    #[allow(dead_code)]
254    pub fn generate_salt() -> String {
255        let mut salt = [0u8; 16];
256        let mut rng = rng();
257        rng.fill(&mut salt);
258        base64::Engine::encode(&base64::engine::general_purpose::STANDARD, salt)
259    }
260
261    /// Validate key derivation parameters
262    #[allow(dead_code)]
263    pub fn validate_parameters(&self, method: &KeyDerivationMethod) -> EncryptionResult<()> {
264        match method {
265            KeyDerivationMethod::Argon2 {
266                memory_kib,
267                iterations,
268                parallelism,
269            } => {
270                if *memory_kib < 8 {
271                    return Err(EncryptionError::invalid_algorithm(
272                        "Argon2 memory must be at least 8 KiB",
273                    ));
274                }
275                if *iterations < 1 {
276                    return Err(EncryptionError::invalid_algorithm(
277                        "Argon2 iterations must be at least 1",
278                    ));
279                }
280                if *parallelism < 1 {
281                    return Err(EncryptionError::invalid_algorithm(
282                        "Argon2 parallelism must be at least 1",
283                    ));
284                }
285            }
286            KeyDerivationMethod::Pbkdf2 { iterations } => {
287                if *iterations < 1000 {
288                    return Err(EncryptionError::invalid_algorithm(
289                        "PBKDF2 iterations should be at least 1000 for security",
290                    ));
291                }
292            }
293        }
294        Ok(())
295    }
296}
297
298impl Default for KeyDerivationManager {
299    fn default() -> Self {
300        Self::new()
301    }
302}