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::{thread_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").unwrap();
159
160        let params = Params::new(
161            memory_kib,
162            iterations,
163            parallelism,
164            Some(32), // output length
165        )
166        .map_err(|e| EncryptionError::key_derivation_failed(e.to_string()))?;
167
168        let argon2 = Argon2::new(Algorithm::Argon2id, Version::V0x13, params);
169        let password_hash = argon2
170            .hash_password(secret, &salt)
171            .map_err(|e| EncryptionError::key_derivation_failed(e.to_string()))?;
172
173        let hash_binding = password_hash.hash.unwrap();
174        let hash_bytes = hash_binding.as_bytes();
175        let key_bytes: Vec<u8> = hash_bytes.to_vec();
176
177        // Take only the required number of bytes for the algorithm
178        let key_len = match algorithm {
179            EncryptionAlgorithm::Aes256Gcm => 32,
180            EncryptionAlgorithm::ChaCha20Poly1305 => 32,
181        };
182
183        let final_key_bytes = if key_bytes.len() >= key_len {
184            key_bytes[..key_len].to_vec()
185        } else {
186            return Err(EncryptionError::key_derivation_failed(
187                "Derived key too short for algorithm",
188            ));
189        };
190
191        EncryptionKey::new(final_key_bytes, algorithm)
192    }
193
194    /// Derive key using PBKDF2
195    fn derive_key_pbkdf2(
196        &self,
197        secret: &[u8],
198        salt: &str,
199        iterations: u32,
200        algorithm: EncryptionAlgorithm,
201    ) -> EncryptionResult<EncryptionKey> {
202        let salt_bytes = salt.as_bytes();
203        let key_len = match algorithm {
204            EncryptionAlgorithm::Aes256Gcm => 32,
205            EncryptionAlgorithm::ChaCha20Poly1305 => 32,
206        };
207
208        let mut derived_key = vec![0u8; key_len];
209        pbkdf2_hmac::<Sha256>(secret, salt_bytes, iterations, &mut derived_key);
210
211        EncryptionKey::new(derived_key, algorithm)
212    }
213
214    /// Verify a password against a derived key (synchronous version)
215    ///
216    /// Note: This is CPU-intensive. Use `verify_password_async` when calling from async context.
217    #[allow(dead_code)]
218    pub fn verify_password(
219        &self,
220        password: &str,
221        expected_key: &EncryptionKey,
222    ) -> EncryptionResult<bool> {
223        let derived_key = self.derive_master_key(password)?;
224
225        Ok(derived_key.as_bytes() == expected_key.as_bytes())
226    }
227
228    /// Verify a password against a derived key (async version using spawn_blocking)
229    ///
230    /// This method offloads the CPU-intensive Argon2 computation to a blocking thread pool.
231    #[allow(dead_code)]
232    pub async fn verify_password_async(
233        &self,
234        password: String,
235        expected_key: EncryptionKey,
236    ) -> EncryptionResult<bool> {
237        let params = self.default_argon2_params.clone();
238        tokio::task::spawn_blocking(move || {
239            let manager = KeyDerivationManager {
240                default_argon2_params: params,
241            };
242            let derived_key = manager.derive_master_key(&password)?;
243            Ok(derived_key.as_bytes() == expected_key.as_bytes())
244        })
245        .await
246        .map_err(|e| EncryptionError::key_derivation_failed(format!("Task join error: {}", e)))?
247    }
248
249    /// Generate a secure random salt
250    #[allow(dead_code)]
251    pub fn generate_salt() -> String {
252        let mut salt = [0u8; 16];
253        let mut rng = thread_rng();
254        rng.fill(&mut salt);
255        base64::Engine::encode(&base64::engine::general_purpose::STANDARD, salt)
256    }
257
258    /// Validate key derivation parameters
259    #[allow(dead_code)]
260    pub fn validate_parameters(&self, method: &KeyDerivationMethod) -> EncryptionResult<()> {
261        match method {
262            KeyDerivationMethod::Argon2 {
263                memory_kib,
264                iterations,
265                parallelism,
266            } => {
267                if *memory_kib < 8 {
268                    return Err(EncryptionError::invalid_algorithm(
269                        "Argon2 memory must be at least 8 KiB",
270                    ));
271                }
272                if *iterations < 1 {
273                    return Err(EncryptionError::invalid_algorithm(
274                        "Argon2 iterations must be at least 1",
275                    ));
276                }
277                if *parallelism < 1 {
278                    return Err(EncryptionError::invalid_algorithm(
279                        "Argon2 parallelism must be at least 1",
280                    ));
281                }
282            }
283            KeyDerivationMethod::Pbkdf2 { iterations } => {
284                if *iterations < 1000 {
285                    return Err(EncryptionError::invalid_algorithm(
286                        "PBKDF2 iterations should be at least 1000 for security",
287                    ));
288                }
289            }
290        }
291        Ok(())
292    }
293}
294
295impl Default for KeyDerivationManager {
296    fn default() -> Self {
297        Self::new()
298    }
299}