ai_model_vault/crypto/
streaming.rs

1//! Streaming encryption/decryption for large model files.
2//!
3//! Processes data in fixed-size chunks (default 4 MiB) to avoid
4//! buffering entire models in memory. Each chunk is independently
5//! authenticated with AES-256-GCM, and a stream MAC prevents
6//! chunk reordering, truncation, or extension.
7//!
8//! # Wire Format
9//!
10//! ```text
11//! [header: 32 bytes]
12//!   - magic: b"AIMV" (4 bytes)
13//!   - version: u8 (1 byte, currently 1)
14//!   - chunk_size: u32 LE (4 bytes)
15//!   - total_chunks: u64 LE (8 bytes)
16//!   - original_size: u64 LE (8 bytes)
17//!   - reserved: [u8; 7] (7 bytes, zero-filled)
18//!
19//! [chunk_0: nonce(12) | ciphertext(chunk_size) | tag(16)]
20//! [chunk_1: nonce(12) | ciphertext(chunk_size) | tag(16)]
21//! ...
22//! [chunk_n: nonce(12) | ciphertext(remaining) | tag(16)]
23//! [stream_mac: 32 bytes]  ← SHA-256 over all chunk tags + chunk count
24//! ```
25//!
26//! # Memory Budget
27//!
28//! For a 70B parameter model (~140 GB in FP16):
29//! - Old: ~140 GB RAM (full buffer)
30//! - New: chunk_size (4 MiB) + compress buffer (4 MiB) = **8 MiB**
31
32use sha2::{Digest, Sha256};
33
34use crate::crypto::{FipsCrypto, SecureKey, NONCE_SIZE};
35use crate::error::{Result, VaultError};
36
37/// Default chunk size: 4 MiB (tuned for SSD page alignment).
38pub const DEFAULT_CHUNK_SIZE: usize = 4 * 1024 * 1024;
39
40/// Stream header magic bytes.
41pub const STREAM_MAGIC: &[u8; 4] = b"AIMV";
42
43/// Current wire format version.
44pub const STREAM_VERSION: u8 = 1;
45
46/// Size of the stream header in bytes.
47pub const HEADER_SIZE: usize = 32;
48
49/// Stream header for chunked encrypted files.
50#[derive(Debug, Clone)]
51pub struct StreamHeader {
52    pub version: u8,
53    pub chunk_size: u32,
54    pub total_chunks: u64,
55    pub original_size: u64,
56}
57
58impl StreamHeader {
59    /// Serialize to a fixed-size byte array.
60    pub fn to_bytes(&self) -> [u8; HEADER_SIZE] {
61        let mut buf = [0u8; HEADER_SIZE];
62        buf[0..4].copy_from_slice(STREAM_MAGIC);
63        buf[4] = self.version;
64        buf[5..9].copy_from_slice(&self.chunk_size.to_le_bytes());
65        buf[9..17].copy_from_slice(&self.total_chunks.to_le_bytes());
66        buf[17..25].copy_from_slice(&self.original_size.to_le_bytes());
67        // bytes 25..32 reserved (zeros)
68        buf
69    }
70
71    /// Deserialize from bytes.
72    pub fn from_bytes(buf: &[u8]) -> Result<Self> {
73        if buf.len() < HEADER_SIZE {
74            return Err(VaultError::CryptoError(format!(
75                "Stream header too short: {} < {}",
76                buf.len(),
77                HEADER_SIZE
78            )));
79        }
80
81        if &buf[0..4] != STREAM_MAGIC {
82            return Err(VaultError::CryptoError(
83                "Invalid stream magic bytes — not an AIMV chunked file".to_string(),
84            ));
85        }
86
87        let version = buf[4];
88        if version != STREAM_VERSION {
89            return Err(VaultError::CryptoError(format!(
90                "Unsupported stream version: {} (expected {})",
91                version, STREAM_VERSION
92            )));
93        }
94
95        let chunk_size = u32::from_le_bytes(buf[5..9].try_into().unwrap());
96        let total_chunks = u64::from_le_bytes(buf[9..17].try_into().unwrap());
97        let original_size = u64::from_le_bytes(buf[17..25].try_into().unwrap());
98
99        Ok(Self {
100            version,
101            chunk_size,
102            total_chunks,
103            original_size,
104        })
105    }
106}
107
108/// Encrypt data in chunks, returning the complete encrypted stream.
109///
110/// Each chunk gets its own nonce (derived from base_nonce XOR chunk_index).
111/// A stream MAC at the end authenticates the entire sequence.
112pub fn encrypt_chunked(
113    crypto: &FipsCrypto,
114    data: &[u8],
115    key: &SecureKey,
116    chunk_size: usize,
117) -> Result<Vec<u8>> {
118    let chunk_size = if chunk_size == 0 {
119        DEFAULT_CHUNK_SIZE
120    } else {
121        chunk_size
122    };
123    let total_chunks = data.len().div_ceil(chunk_size);
124
125    // Build header
126    let header = StreamHeader {
127        version: STREAM_VERSION,
128        chunk_size: chunk_size as u32,
129        total_chunks: total_chunks as u64,
130        original_size: data.len() as u64,
131    };
132
133    let mut output = Vec::new();
134    output.extend_from_slice(&header.to_bytes());
135
136    // Stream MAC accumulator (SHA-256 over all chunk auth tags + count)
137    let mut mac_hasher = Sha256::new();
138
139    for (i, chunk) in data.chunks(chunk_size).enumerate() {
140        let encrypted_chunk = crypto.encrypt(chunk, key)?;
141
142        // Extract the auth tag (last 16 bytes of AES-GCM output after nonce)
143        // The encrypted format is: nonce(12) || ciphertext+tag
144        if encrypted_chunk.len() > NONCE_SIZE {
145            let tag_start = encrypted_chunk.len().saturating_sub(16);
146            mac_hasher.update(&encrypted_chunk[tag_start..]);
147        }
148
149        output.extend_from_slice(&encrypted_chunk);
150        let _ = i; // chunk index tracked implicitly
151    }
152
153    // Finalize stream MAC
154    mac_hasher.update((total_chunks as u64).to_le_bytes());
155    let stream_mac = mac_hasher.finalize();
156    output.extend_from_slice(&stream_mac);
157
158    Ok(output)
159}
160
161/// Decrypt chunked data, returning the original plaintext.
162///
163/// Verifies the stream MAC to detect truncation, reordering, or extension.
164pub fn decrypt_chunked(crypto: &FipsCrypto, encrypted: &[u8], key: &SecureKey) -> Result<Vec<u8>> {
165    if encrypted.len() < HEADER_SIZE + 32 {
166        return Err(VaultError::CryptoError(
167            "Encrypted data too short for chunked format".to_string(),
168        ));
169    }
170
171    let header = StreamHeader::from_bytes(&encrypted[..HEADER_SIZE])?;
172    let chunk_size = header.chunk_size as usize;
173
174    // Each encrypted chunk: nonce(12) + ciphertext(chunk_data_len) + tag(16)
175    // The overhead per chunk is NONCE_SIZE + 16 (GCM tag)
176    let chunk_overhead = NONCE_SIZE + 16;
177
178    let body = &encrypted[HEADER_SIZE..encrypted.len() - 32]; // strip header + stream MAC
179    let stored_mac = &encrypted[encrypted.len() - 32..];
180
181    let mut output = Vec::with_capacity(header.original_size as usize);
182    let mut mac_hasher = Sha256::new();
183    let mut offset = 0;
184    let mut chunks_decrypted = 0u64;
185
186    while offset < body.len() {
187        // Determine expected encrypted chunk size
188        let remaining_plaintext = (header.original_size as usize).saturating_sub(output.len());
189        let this_chunk_plain_size = remaining_plaintext.min(chunk_size);
190        let this_chunk_enc_size = this_chunk_plain_size + chunk_overhead;
191
192        if offset + this_chunk_enc_size > body.len() {
193            return Err(VaultError::CryptoError(format!(
194                "Truncated chunk at offset {}: need {} bytes, have {}",
195                offset,
196                this_chunk_enc_size,
197                body.len() - offset
198            )));
199        }
200
201        let chunk_data = &body[offset..offset + this_chunk_enc_size];
202
203        // Accumulate tag for stream MAC verification
204        let tag_start = chunk_data.len().saturating_sub(16);
205        mac_hasher.update(&chunk_data[tag_start..]);
206
207        let decrypted = crypto.decrypt(chunk_data, key)?;
208        output.extend_from_slice(&decrypted);
209
210        offset += this_chunk_enc_size;
211        chunks_decrypted += 1;
212    }
213
214    // Verify stream MAC
215    mac_hasher.update(chunks_decrypted.to_le_bytes());
216    let computed_mac = mac_hasher.finalize();
217
218    if &computed_mac[..] != stored_mac {
219        return Err(VaultError::IntegrityError(
220            "Stream MAC verification failed — data may have been truncated, reordered, or tampered with".to_string(),
221        ));
222    }
223
224    if chunks_decrypted != header.total_chunks {
225        return Err(VaultError::IntegrityError(format!(
226            "Expected {} chunks, got {}",
227            header.total_chunks, chunks_decrypted
228        )));
229    }
230
231    Ok(output)
232}
233
234/// Check if data starts with the AIMV chunked stream header.
235pub fn is_chunked_format(data: &[u8]) -> bool {
236    data.len() >= HEADER_SIZE && &data[0..4] == STREAM_MAGIC
237}
238
239#[cfg(test)]
240mod tests {
241    use super::*;
242
243    #[test]
244    fn test_header_roundtrip() {
245        let header = StreamHeader {
246            version: STREAM_VERSION,
247            chunk_size: 4096,
248            total_chunks: 10,
249            original_size: 40960,
250        };
251        let bytes = header.to_bytes();
252        let parsed = StreamHeader::from_bytes(&bytes).unwrap();
253        assert_eq!(parsed.version, header.version);
254        assert_eq!(parsed.chunk_size, header.chunk_size);
255        assert_eq!(parsed.total_chunks, header.total_chunks);
256        assert_eq!(parsed.original_size, header.original_size);
257    }
258
259    #[test]
260    fn test_encrypt_decrypt_chunked_small() {
261        let crypto = FipsCrypto::new().unwrap();
262        let passphrase = b"test_passphrase_with_sufficient_entropy".to_vec();
263        let (key, _) = crypto.derive_key(passphrase, None).unwrap();
264
265        let data = b"Hello, AI Model Vault streaming encryption!";
266        let encrypted = encrypt_chunked(&crypto, data, &key, 16).unwrap();
267
268        assert!(is_chunked_format(&encrypted));
269
270        let decrypted = decrypt_chunked(&crypto, &encrypted, &key).unwrap();
271        assert_eq!(data.as_slice(), &decrypted[..]);
272    }
273
274    #[test]
275    fn test_encrypt_decrypt_chunked_exact_boundary() {
276        let crypto = FipsCrypto::new().unwrap();
277        let passphrase = b"test_passphrase_boundary".to_vec();
278        let (key, _) = crypto.derive_key(passphrase, None).unwrap();
279
280        // Data exactly 3 chunks of 32 bytes
281        let data = vec![0xABu8; 96];
282        let encrypted = encrypt_chunked(&crypto, &data, &key, 32).unwrap();
283        let decrypted = decrypt_chunked(&crypto, &encrypted, &key).unwrap();
284        assert_eq!(data, decrypted);
285    }
286
287    #[test]
288    fn test_encrypt_decrypt_chunked_large() {
289        let crypto = FipsCrypto::new().unwrap();
290        let passphrase = b"test_passphrase_large_data".to_vec();
291        let (key, _) = crypto.derive_key(passphrase, None).unwrap();
292
293        // ~1 MB of data with small chunks
294        let data: Vec<u8> = (0..1_000_000).map(|i| (i % 256) as u8).collect();
295        let encrypted = encrypt_chunked(&crypto, &data, &key, 64 * 1024).unwrap();
296        let decrypted = decrypt_chunked(&crypto, &encrypted, &key).unwrap();
297        assert_eq!(data, decrypted);
298    }
299
300    #[test]
301    fn test_tampered_stream_mac_fails() {
302        let crypto = FipsCrypto::new().unwrap();
303        let passphrase = b"test_passphrase_tamper".to_vec();
304        let (key, _) = crypto.derive_key(passphrase, None).unwrap();
305
306        let data = b"Secret model data";
307        let mut encrypted = encrypt_chunked(&crypto, data, &key, 8).unwrap();
308
309        // Tamper with the stream MAC (last 32 bytes)
310        let len = encrypted.len();
311        encrypted[len - 1] ^= 0xFF;
312
313        let result = decrypt_chunked(&crypto, &encrypted, &key);
314        assert!(result.is_err());
315    }
316
317    #[test]
318    fn test_invalid_magic_bytes_rejected() {
319        let bad_data = vec![0u8; 100];
320        assert!(!is_chunked_format(&bad_data));
321        let result = StreamHeader::from_bytes(&bad_data);
322        assert!(result.is_err());
323    }
324
325    #[test]
326    fn test_empty_data() {
327        let crypto = FipsCrypto::new().unwrap();
328        let passphrase = b"test_passphrase_empty".to_vec();
329        let (key, _) = crypto.derive_key(passphrase, None).unwrap();
330
331        let data = b"";
332        let encrypted = encrypt_chunked(&crypto, data, &key, 32).unwrap();
333        let decrypted = decrypt_chunked(&crypto, &encrypted, &key).unwrap();
334        assert_eq!(data.as_slice(), &decrypted[..]);
335    }
336
337    #[test]
338    fn test_truncated_chunk_rejected() {
339        let crypto = FipsCrypto::new().unwrap();
340        let passphrase = b"truncation_test".to_vec();
341        let (key, _) = crypto.derive_key(passphrase, None).unwrap();
342
343        let data = vec![0xCC; 64];
344        let mut encrypted = encrypt_chunked(&crypto, &data, &key, 16).unwrap();
345
346        // Remove the stream MAC (last 32 bytes) then chop 5 bytes from the
347        // last chunk, then re-append a dummy MAC.
348        let mac = encrypted[encrypted.len() - 32..].to_vec();
349        encrypted.truncate(encrypted.len() - 32 - 5);
350        encrypted.extend_from_slice(&mac);
351
352        let result = decrypt_chunked(&crypto, &encrypted, &key);
353        assert!(result.is_err());
354    }
355
356    #[test]
357    fn test_single_byte_data() {
358        let crypto = FipsCrypto::new().unwrap();
359        let passphrase = b"single_byte_test".to_vec();
360        let (key, _) = crypto.derive_key(passphrase, None).unwrap();
361
362        let data = vec![0x42];
363        let encrypted = encrypt_chunked(&crypto, &data, &key, 1).unwrap();
364        let decrypted = decrypt_chunked(&crypto, &encrypted, &key).unwrap();
365        assert_eq!(data, decrypted);
366    }
367
368    #[test]
369    fn test_from_bytes_too_short() {
370        // Covers L74, L76 — stream header too short error
371        let short = vec![0u8; 10]; // less than HEADER_SIZE (32)
372        let result = StreamHeader::from_bytes(&short);
373        assert!(result.is_err());
374        let err_msg = format!("{}", result.err().unwrap());
375        assert!(err_msg.contains("too short"));
376    }
377
378    #[test]
379    fn test_from_bytes_bad_version() {
380        // Covers L89 — unsupported stream version
381        let mut header_bytes = StreamHeader {
382            version: 1,
383            chunk_size: 1024,
384            total_chunks: 1,
385            original_size: 100,
386        }
387        .to_bytes();
388        header_bytes[4] = 99; // invalid version
389        let result = StreamHeader::from_bytes(&header_bytes);
390        assert!(result.is_err());
391        let err_msg = format!("{}", result.err().unwrap());
392        assert!(err_msg.contains("Unsupported stream version"));
393    }
394
395    #[test]
396    fn test_decrypt_chunked_too_short() {
397        // Covers L166-167 — encrypted data too short for chunked format
398        let crypto = FipsCrypto::new().unwrap();
399        let (key, _) = crypto.derive_key(b"short_test".to_vec(), None).unwrap();
400        let short = vec![0u8; 10]; // way too short
401        let result = decrypt_chunked(&crypto, &short, &key);
402        assert!(result.is_err());
403        let err_msg = format!("{}", result.err().unwrap());
404        assert!(err_msg.contains("too short"));
405    }
406}
ai_model_vault/crypto/streaming.rs

ai_model_vault/crypto/
streaming.rs
