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

shadow_crypt_shell/
kdf.rs
