Skip to main content

shadow_crypt_shell/
kdf.rs

1//! Guards for Argon2 parameters read from untrusted file headers.
2//!
3//! Every code path that derives a key from header-supplied parameters
4//! (decryption and listing alike) must validate them here first, so a
5//! crafted .shadow file cannot request an enormous allocation or an
6//! excessive amount of CPU time.
7
8use std::sync::{Condvar, Mutex};
9
10use shadow_crypt_core::{
11    memory::{SecureKey, SecureString},
12    vault::{KdfRequest, ParsedFile},
13};
14
15use crate::errors::{WorkflowError, WorkflowResult};
16
17/// Upper bounds for KDF parameters read from untrusted file headers.
18/// These prevent a crafted file from causing OOM or excessive CPU use.
19pub const MAX_KDF_MEMORY_KIB: u32 = 8 * 1024 * 1024; // 8 GiB
20pub const MAX_KDF_ITERATIONS: u32 = 1_000;
21pub const MAX_KDF_PARALLELISM: u32 = 256;
22pub const MAX_KDF_KEY_SIZE: u8 = 64;
23
24/// Validates the KDF parameters a file header requests, against the bounds
25/// above. The request is version-agnostic, so every format version passes
26/// through the same checks.
27pub fn validate_kdf_request(request: &KdfRequest) -> WorkflowResult<()> {
28    if request.memory_cost > MAX_KDF_MEMORY_KIB {
29        return Err(WorkflowError::UserInput(format!(
30            "KDF memory cost in file header is too large: {} KiB (max {} KiB)",
31            request.memory_cost, MAX_KDF_MEMORY_KIB
32        )));
33    }
34    if request.time_cost > MAX_KDF_ITERATIONS {
35        return Err(WorkflowError::UserInput(format!(
36            "KDF iteration count in file header is too large: {} (max {})",
37            request.time_cost, MAX_KDF_ITERATIONS
38        )));
39    }
40    if request.parallelism > MAX_KDF_PARALLELISM {
41        return Err(WorkflowError::UserInput(format!(
42            "KDF parallelism in file header is too large: {} (max {})",
43            request.parallelism, MAX_KDF_PARALLELISM
44        )));
45    }
46    if request.key_size > MAX_KDF_KEY_SIZE {
47        return Err(WorkflowError::UserInput(format!(
48            "KDF key size in file header is too large: {} bytes (max {})",
49            request.key_size, MAX_KDF_KEY_SIZE
50        )));
51    }
52    Ok(())
53}
54
55/// Derives a parsed file's key from a password, guarding the (untrusted)
56/// KDF parameters its header requests.
57///
58/// Validates the request against the bounds above and, only if it passes,
59/// runs the file's own key derivation while holding a reservation against
60/// the global memory budget. This is the single intended entry point for
61/// header-supplied parameters: fusing the two steps makes it impossible to
62/// run an unvalidated derivation from a crafted file.
63pub fn derive_untrusted_key(
64    parsed: &ParsedFile,
65    password: &SecureString,
66) -> WorkflowResult<SecureKey> {
67    let request = parsed.kdf_request();
68    validate_kdf_request(&request)?;
69    let (key, _) = with_kdf_memory_permit(request.memory_cost, || {
70        parsed.derive_key(password.as_str().as_bytes())
71    })?;
72    Ok(key)
73}
74
75/// Total Argon2 buffer memory allowed across all concurrent derivations.
76///
77/// Argon2 allocates its full memory cost per derivation, and files are
78/// processed in parallel (one rayon thread per core). Without a cap, ten
79/// production-profile files on a 10-core machine would hold ~10 GiB of
80/// Argon2 buffers at once. The budget throttles concurrency by declared
81/// memory cost instead of thread count, so cheap derivations still run
82/// fully parallel.
83pub const KDF_MEMORY_BUDGET_KIB: u64 = 4 * 1024 * 1024; // 4 GiB
84
85/// Global gate shared by all workflows in the process.
86static KDF_GATE: KdfGate = KdfGate::new(KDF_MEMORY_BUDGET_KIB);
87
88/// Runs `f` (a key derivation) while holding a reservation of `cost_kib`
89/// against the global memory budget, blocking until enough budget is free.
90///
91/// A cost larger than the whole budget is clamped to it, so an oversized
92/// (but validated) derivation waits for exclusive use of the budget and
93/// then runs alone rather than deadlocking.
94pub fn with_kdf_memory_permit<T>(cost_kib: u32, f: impl FnOnce() -> T) -> T {
95    let _permit = KDF_GATE.acquire(cost_kib as u64);
96    f()
97}
98
99struct KdfGate {
100    budget_kib: u64,
101    in_use_kib: Mutex<u64>,
102    released: Condvar,
103}
104
105impl KdfGate {
106    const fn new(budget_kib: u64) -> Self {
107        Self {
108            budget_kib,
109            in_use_kib: Mutex::new(0),
110            released: Condvar::new(),
111        }
112    }
113
114    fn acquire(&self, cost_kib: u64) -> KdfPermit<'_> {
115        // Reserve at least 1 KiB so a zero-cost caller still participates,
116        // and never more than the budget so acquisition always succeeds.
117        let cost_kib = cost_kib.clamp(1, self.budget_kib);
118        let mut in_use = self.in_use_kib.lock().unwrap_or_else(|e| e.into_inner());
119        while *in_use + cost_kib > self.budget_kib {
120            in_use = self
121                .released
122                .wait(in_use)
123                .unwrap_or_else(|e| e.into_inner());
124        }
125        *in_use += cost_kib;
126        KdfPermit {
127            gate: self,
128            cost_kib,
129        }
130    }
131}
132
133struct KdfPermit<'a> {
134    gate: &'a KdfGate,
135    cost_kib: u64,
136}
137
138impl Drop for KdfPermit<'_> {
139    fn drop(&mut self) {
140        let mut in_use = self
141            .gate
142            .in_use_kib
143            .lock()
144            .unwrap_or_else(|e| e.into_inner());
145        *in_use -= self.cost_kib;
146        drop(in_use);
147        self.gate.released.notify_all();
148    }
149}
150
151#[cfg(test)]
152mod tests {
153    use super::*;
154
155    fn request(memory_cost: u32, time_cost: u32, parallelism: u32, key_size: u8) -> KdfRequest {
156        KdfRequest {
157            salt: [0u8; 16],
158            memory_cost,
159            time_cost,
160            parallelism,
161            key_size,
162        }
163    }
164
165    #[test]
166    fn test_valid_kdf_params_accepted() {
167        assert!(validate_kdf_request(&request(1024, 1, 1, 32)).is_ok());
168    }
169
170    #[test]
171    fn test_memory_cost_too_large() {
172        assert!(validate_kdf_request(&request(MAX_KDF_MEMORY_KIB + 1, 1, 1, 32)).is_err());
173    }
174
175    #[test]
176    fn test_memory_cost_at_limit_accepted() {
177        assert!(validate_kdf_request(&request(MAX_KDF_MEMORY_KIB, 1, 1, 32)).is_ok());
178    }
179
180    #[test]
181    fn test_iterations_too_large() {
182        assert!(validate_kdf_request(&request(1024, MAX_KDF_ITERATIONS + 1, 1, 32)).is_err());
183    }
184
185    #[test]
186    fn test_parallelism_too_large() {
187        assert!(validate_kdf_request(&request(1024, 1, MAX_KDF_PARALLELISM + 1, 32)).is_err());
188    }
189
190    #[test]
191    fn test_key_size_too_large() {
192        assert!(validate_kdf_request(&request(1024, 1, 1, MAX_KDF_KEY_SIZE + 1)).is_err());
193    }
194
195    #[test]
196    fn test_gate_allows_cost_within_budget() {
197        let gate = KdfGate::new(100);
198        let permit = gate.acquire(60);
199        assert_eq!(*gate.in_use_kib.lock().unwrap(), 60);
200        drop(permit);
201        assert_eq!(*gate.in_use_kib.lock().unwrap(), 0);
202    }
203
204    #[test]
205    fn test_gate_clamps_oversized_cost_to_budget() {
206        let gate = KdfGate::new(100);
207        // A cost above the budget must not deadlock: it is clamped and runs alone.
208        let permit = gate.acquire(1_000_000);
209        assert_eq!(*gate.in_use_kib.lock().unwrap(), 100);
210        drop(permit);
211    }
212
213    #[test]
214    fn test_gate_zero_cost_still_reserves() {
215        let gate = KdfGate::new(100);
216        let permit = gate.acquire(0);
217        assert_eq!(*gate.in_use_kib.lock().unwrap(), 1);
218        drop(permit);
219    }
220
221    #[test]
222    fn test_gate_limits_concurrent_memory_use() {
223        use std::sync::{
224            Arc,
225            atomic::{AtomicU64, Ordering},
226        };
227
228        // Budget fits exactly two concurrent 50-KiB permits.
229        let gate = Arc::new(KdfGate::new(100));
230        let active = Arc::new(AtomicU64::new(0));
231        let max_active = Arc::new(AtomicU64::new(0));
232
233        let handles: Vec<_> = (0..8)
234            .map(|_| {
235                let gate = Arc::clone(&gate);
236                let active = Arc::clone(&active);
237                let max_active = Arc::clone(&max_active);
238                std::thread::spawn(move || {
239                    let _permit = gate.acquire(50);
240                    let now = active.fetch_add(1, Ordering::SeqCst) + 1;
241                    max_active.fetch_max(now, Ordering::SeqCst);
242                    std::thread::sleep(std::time::Duration::from_millis(10));
243                    active.fetch_sub(1, Ordering::SeqCst);
244                })
245            })
246            .collect();
247
248        for handle in handles {
249            handle.join().unwrap();
250        }
251
252        assert!(max_active.load(Ordering::SeqCst) <= 2);
253        assert_eq!(*gate.in_use_kib.lock().unwrap(), 0);
254    }
255
256    #[test]
257    fn test_with_kdf_memory_permit_returns_value() {
258        let result = with_kdf_memory_permit(1024, || 42);
259        assert_eq!(result, 42);
260    }
261
262    #[test]
263    fn test_paranoid_params_accepted() {
264        let params = shadow_crypt_core::v1::key::KeyDerivationParams::paranoid_defaults();
265        assert!(
266            validate_kdf_request(&request(
267                params.memory_cost,
268                params.time_cost,
269                params.parallelism,
270                params.key_size
271            ))
272            .is_ok()
273        );
274    }
275}