1use crate::metrics::OperationMetrics;
9#[cfg(feature = "compression")]
10use lz4_flex;
11use serde::{Deserialize, Serialize};
12use std::sync::{Arc, Mutex};
13#[cfg(not(target_arch = "wasm32"))]
14use std::time::Instant;
15use thiserror::Error;
16#[cfg(feature = "checksum")]
17use xxhash_rust::xxh3::xxh3_64;
18
19#[derive(Debug, Error, Clone, PartialEq)]
21pub enum ByteStorageError {
22 #[error("input exceeds maximum size")]
23 InputTooLarge,
24
25 #[error("decompression ratio exceeds safety limit")]
26 DecompressionBomb,
27
28 #[error("integrity check failed")]
29 ChecksumMismatch,
30
31 #[error("compression failed")]
32 CompressionFailed,
33
34 #[error("decompression failed")]
35 DecompressionFailed,
36
37 #[error("size validation failed")]
38 SizeValidationFailed,
39
40 #[error("serialization failed: {0}")]
41 SerializationFailed(String),
42
43 #[error("deserialization failed: {0}")]
44 DeserializationFailed(String),
45}
46
47const MAX_UNCOMPRESSED_SIZE: usize = 512 * 1024 * 1024; const MAX_COMPRESSED_SIZE: usize = 512 * 1024 * 1024; const MAX_COMPRESSION_RATIO: u64 = 1000;
53
54#[derive(Serialize, Deserialize)]
57pub struct StorageEnvelope {
58 pub compressed_data: Vec<u8>,
60 pub checksum: [u8; 8],
62 pub original_size: u32,
64 pub format: String,
66}
67
68impl StorageEnvelope {
69 #[cfg(all(feature = "compression", feature = "checksum"))]
71 pub fn new(data: Vec<u8>, format: String) -> Result<Self, ByteStorageError> {
72 if data.len() > MAX_UNCOMPRESSED_SIZE {
74 return Err(ByteStorageError::InputTooLarge);
75 }
76
77 let original_size = data.len() as u32;
78
79 let compressed_data = lz4_flex::compress(&data);
81
82 if compressed_data.len() > MAX_COMPRESSED_SIZE {
84 return Err(ByteStorageError::InputTooLarge);
85 }
86
87 let checksum = xxh3_64(&data).to_be_bytes();
89
90 Ok(StorageEnvelope {
91 compressed_data,
92 checksum,
93 original_size,
94 format,
95 })
96 }
97
98 #[cfg(all(feature = "compression", feature = "checksum"))]
100 pub fn extract(&self) -> Result<Vec<u8>, ByteStorageError> {
101 if self.compressed_data.len() > MAX_COMPRESSED_SIZE {
103 return Err(ByteStorageError::InputTooLarge);
104 }
105
106 if self.original_size as usize > MAX_UNCOMPRESSED_SIZE {
107 return Err(ByteStorageError::InputTooLarge);
108 }
109
110 let compressed_size = self.compressed_data.len() as u64;
113
114 if compressed_size == 0 {
116 return Err(ByteStorageError::DecompressionBomb);
117 }
118
119 let max_allowed_original = MAX_COMPRESSION_RATIO
121 .checked_mul(compressed_size)
122 .ok_or(ByteStorageError::DecompressionBomb)?;
123
124 if (self.original_size as u64) > max_allowed_original {
126 return Err(ByteStorageError::DecompressionBomb);
127 }
128
129 let decompressed = lz4_flex::decompress(&self.compressed_data, self.original_size as usize)
131 .map_err(|_| ByteStorageError::DecompressionFailed)?;
132
133 let computed_checksum = xxh3_64(&decompressed).to_be_bytes();
137 if computed_checksum != self.checksum {
138 return Err(ByteStorageError::ChecksumMismatch);
139 }
140
141 if decompressed.len() != self.original_size as usize {
143 return Err(ByteStorageError::SizeValidationFailed);
144 }
145
146 Ok(decompressed)
147 }
148}
149
150pub struct ByteStorage {
153 default_format: String,
154 last_metrics: Arc<Mutex<OperationMetrics>>,
156}
157
158impl ByteStorage {
159 pub fn new(default_format: Option<String>) -> Self {
161 ByteStorage {
162 default_format: default_format.unwrap_or_else(|| "msgpack".to_string()),
163 last_metrics: Arc::new(Mutex::new(OperationMetrics::new())),
164 }
165 }
166
167 #[cfg(all(feature = "compression", feature = "checksum", feature = "messagepack"))]
171 pub fn store(&self, data: &[u8], format: Option<String>) -> Result<Vec<u8>, ByteStorageError> {
172 if data.len() > MAX_UNCOMPRESSED_SIZE {
174 return Err(ByteStorageError::InputTooLarge);
175 }
176
177 let format = format.unwrap_or_else(|| self.default_format.clone());
178
179 #[cfg(not(target_arch = "wasm32"))]
181 let compression_start = Instant::now();
182 let original_size = data.len();
183
184 let envelope = StorageEnvelope::new(data.to_vec(), format)?;
185
186 #[cfg(not(target_arch = "wasm32"))]
187 let compression_micros = compression_start.elapsed().as_micros() as u64;
188 #[cfg(target_arch = "wasm32")]
189 let compression_micros = 0u64;
190 let compressed_size = envelope.compressed_data.len();
191
192 let envelope_bytes = rmp_serde::to_vec(&envelope)
194 .map_err(|e| ByteStorageError::SerializationFailed(e.to_string()))?;
195
196 if envelope_bytes.len() > MAX_COMPRESSED_SIZE {
198 return Err(ByteStorageError::InputTooLarge);
199 }
200
201 if let Ok(mut metrics) = self.last_metrics.lock() {
203 *metrics = OperationMetrics::new().with_compression(
204 compression_micros,
205 original_size,
206 compressed_size,
207 );
208 }
209
210 Ok(envelope_bytes)
211 }
212
213 #[cfg(all(feature = "compression", feature = "checksum", feature = "messagepack"))]
217 pub fn retrieve(&self, envelope_bytes: &[u8]) -> Result<(Vec<u8>, String), ByteStorageError> {
218 if envelope_bytes.len() > MAX_COMPRESSED_SIZE {
220 return Err(ByteStorageError::InputTooLarge);
221 }
222
223 let envelope: StorageEnvelope = rmp_serde::from_slice(envelope_bytes)
225 .map_err(|e| ByteStorageError::DeserializationFailed(e.to_string()))?;
226
227 #[cfg(not(target_arch = "wasm32"))]
229 let decompress_start = Instant::now();
230
231 let data = envelope.extract()?;
233
234 #[cfg(not(target_arch = "wasm32"))]
235 let decompress_micros = decompress_start.elapsed().as_micros() as u64;
236 #[cfg(target_arch = "wasm32")]
237 let decompress_micros = 0u64;
238
239 let compressed_size = envelope.compressed_data.len();
241 let original_size = envelope.original_size as usize;
242
243 if let Ok(mut metrics) = self.last_metrics.lock() {
245 *metrics = OperationMetrics::new().with_compression(
246 decompress_micros,
247 original_size,
248 compressed_size,
249 );
250 }
251
252 Ok((data, envelope.format))
253 }
254
255 #[cfg(feature = "compression")]
257 pub fn estimate_compression(&self, data: &[u8]) -> Result<f64, ByteStorageError> {
258 if data.len() > MAX_UNCOMPRESSED_SIZE {
260 return Err(ByteStorageError::InputTooLarge);
261 }
262
263 let compressed = lz4_flex::compress(data);
264
265 Ok(data.len() as f64 / compressed.len() as f64)
266 }
267
268 #[cfg(all(feature = "compression", feature = "checksum", feature = "messagepack"))]
270 pub fn validate(&self, envelope_bytes: &[u8]) -> bool {
271 if envelope_bytes.len() > MAX_COMPRESSED_SIZE {
273 return false; }
275
276 match rmp_serde::from_slice::<StorageEnvelope>(envelope_bytes) {
277 Ok(envelope) => envelope.extract().is_ok(),
278 Err(_) => false,
279 }
280 }
281
282 pub fn get_last_metrics(&self) -> OperationMetrics {
286 self.last_metrics
287 .lock()
288 .map(|metrics| metrics.clone())
289 .unwrap_or_else(|_| OperationMetrics::new())
290 }
291
292 pub fn max_uncompressed_size(&self) -> usize {
294 MAX_UNCOMPRESSED_SIZE
295 }
296
297 pub fn max_compressed_size(&self) -> usize {
298 MAX_COMPRESSED_SIZE
299 }
300
301 pub fn max_compression_ratio(&self) -> u64 {
302 MAX_COMPRESSION_RATIO
303 }
304}
305
306impl Default for ByteStorage {
307 fn default() -> Self {
308 Self::new(None)
309 }
310}
311
312#[cfg(all(
313 test,
314 feature = "compression",
315 feature = "checksum",
316 feature = "messagepack"
317))]
318mod tests {
319 use super::*;
320
321 #[test]
322 fn test_storage_envelope_roundtrip() {
323 let data = b"Hello, World! This is test data for compression.".to_vec();
324 let envelope = StorageEnvelope::new(data.clone(), "test".to_string()).unwrap();
325 let extracted = envelope.extract().unwrap();
326 assert_eq!(data, extracted);
327 }
328
329 #[test]
330 fn test_compression_works() {
331 let data = b"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa".to_vec(); let envelope = StorageEnvelope::new(data.clone(), "test".to_string()).unwrap();
333 assert!(envelope.compressed_data.len() < data.len());
334 }
335
336 #[test]
337 fn test_checksum_validation() {
338 let mut envelope = StorageEnvelope::new(b"test".to_vec(), "test".to_string()).unwrap();
339 envelope.checksum[0] = !envelope.checksum[0];
341 assert!(envelope.extract().is_err());
342 }
343
344 #[test]
345 fn test_raw_persistence_roundtrip() {
346 let storage = ByteStorage::new(None);
347 let test_data = b"test data for persistence";
348
349 let stored = storage.store(test_data, None).unwrap();
350 let (retrieved_data, format) = storage.retrieve(&stored).unwrap();
351 assert_eq!(test_data, retrieved_data.as_slice());
352 assert_eq!("msgpack", format);
353 }
354
355 #[test]
356 fn test_size_limits_input() {
357 let storage = ByteStorage::new(None);
358
359 let large_data = vec![0u8; MAX_UNCOMPRESSED_SIZE + 1];
361
362 let result = storage.store(&large_data, None);
363 assert!(matches!(result, Err(ByteStorageError::InputTooLarge)));
364 }
365
366 #[test]
367 fn test_size_limits_envelope() {
368 let max_data = vec![0u8; MAX_UNCOMPRESSED_SIZE];
370 let envelope_result = StorageEnvelope::new(max_data, "test".to_string());
371
372 assert!(envelope_result.is_ok());
374 }
375
376 #[test]
377 fn test_compression_ratio_bomb_protection() {
378 let malicious_envelope = StorageEnvelope {
380 compressed_data: vec![0u8; 1000], checksum: [0u8; 8], original_size: 200 * 1024 * 1024, format: "test".to_string(),
384 };
385
386 let result = malicious_envelope.extract();
387 assert!(matches!(result, Err(ByteStorageError::DecompressionBomb)));
388 }
389
390 #[test]
395 fn test_decompression_bomb_zero_compressed_size() {
396 let malicious_envelope = StorageEnvelope {
398 compressed_data: vec![], checksum: [0u8; 8],
400 original_size: 1000, format: "test".to_string(),
402 };
403
404 let result = malicious_envelope.extract();
405 assert!(
406 matches!(result, Err(ByteStorageError::DecompressionBomb)),
407 "Zero compressed size should be rejected as decompression bomb"
408 );
409 }
410
411 #[test]
412 fn test_decompression_bomb_extreme_ratio() {
413 let malicious_envelope = StorageEnvelope {
417 compressed_data: vec![0u8; 1], checksum: [0u8; 8],
419 original_size: 2000, format: "test".to_string(),
421 };
422
423 let result = malicious_envelope.extract();
424 assert!(
425 matches!(result, Err(ByteStorageError::DecompressionBomb)),
426 "Extreme ratio should be rejected as bomb: {:?}",
427 result
428 );
429 }
430
431 #[test]
432 fn test_decompression_u32_max_original_size() {
433 let malicious_envelope = StorageEnvelope {
437 compressed_data: vec![0u8; 1000],
438 checksum: [0u8; 8],
439 original_size: u32::MAX, format: "test".to_string(),
441 };
442
443 let result = malicious_envelope.extract();
444 assert!(
445 matches!(result, Err(ByteStorageError::InputTooLarge)),
446 "u32::MAX should be rejected as InputTooLarge (exceeds 512MB limit): {:?}",
447 result
448 );
449 }
450
451 #[test]
452 fn test_decompression_exactly_at_threshold() {
453 let envelope = StorageEnvelope {
457 compressed_data: vec![0u8; 100], checksum: [0u8; 8],
459 original_size: 100_000, format: "test".to_string(),
461 };
462
463 let result = envelope.extract();
464 assert!(
467 !matches!(result, Err(ByteStorageError::DecompressionBomb)),
468 "Exactly 1000:1 ratio should pass bomb check: {:?}",
469 result
470 );
471 assert!(
472 result.is_err(),
473 "Invalid data should still fail after ratio check"
474 );
475 }
476
477 #[test]
478 fn test_decompression_just_over_threshold() {
479 let malicious_envelope = StorageEnvelope {
481 compressed_data: vec![0u8; 100], checksum: [0u8; 8],
483 original_size: 100_001, format: "test".to_string(),
485 };
486
487 let result = malicious_envelope.extract();
488 assert!(
489 matches!(result, Err(ByteStorageError::DecompressionBomb)),
490 "Just over 1000:1 ratio should be rejected as bomb"
491 );
492 }
493
494 #[test]
495 fn test_decompression_bomb_integer_boundary() {
496 let envelope = StorageEnvelope {
503 compressed_data: vec![0u8; 1_000_000], checksum: [0u8; 8],
505 original_size: 1_000_000_000, format: "test".to_string(),
507 };
508
509 let result = envelope.extract();
511 assert!(
512 matches!(result, Err(ByteStorageError::InputTooLarge)),
513 "Should fail size check before ratio check: {:?}",
514 result
515 );
516 }
517
518 #[test]
519 fn test_envelope_size_validation() {
520 let storage = ByteStorage::new(None);
521
522 let oversized_envelope = vec![0u8; MAX_COMPRESSED_SIZE + 1];
524
525 let result = storage.retrieve(&oversized_envelope);
526 assert!(matches!(result, Err(ByteStorageError::InputTooLarge)));
527 }
528
529 #[test]
530 fn test_security_limits_getters() {
531 let storage = ByteStorage::new(None);
532
533 assert_eq!(storage.max_uncompressed_size(), MAX_UNCOMPRESSED_SIZE);
534 assert_eq!(storage.max_compressed_size(), MAX_COMPRESSED_SIZE);
535 assert_eq!(storage.max_compression_ratio(), 1000u64);
536 }
537
538 #[test]
539 fn test_compression_estimate_security() {
540 let storage = ByteStorage::new(None);
541
542 let large_data = vec![0u8; MAX_UNCOMPRESSED_SIZE + 1];
544 let result = storage.estimate_compression(&large_data);
545 assert!(matches!(result, Err(ByteStorageError::InputTooLarge)));
546 }
547
548 #[test]
549 fn test_validate_security() {
550 let storage = ByteStorage::new(None);
551
552 let large_envelope = vec![0u8; MAX_COMPRESSED_SIZE + 1];
554 let result = storage.validate(&large_envelope);
555 assert!(!result); }
557
558 #[test]
559 fn test_edge_case_exactly_at_limits() {
560 let storage = ByteStorage::new(None);
562 let max_size_data = vec![1u8; MAX_UNCOMPRESSED_SIZE]; let result = storage.store(&max_size_data, None);
566 assert!(result.is_ok());
567 }
568
569 #[test]
570 fn test_zero_size_edge_case() {
571 let storage = ByteStorage::new(None);
572 let empty_data = vec![];
573
574 let stored = storage.store(&empty_data, None).unwrap();
575 let (retrieved_data, format) = storage.retrieve(&stored).unwrap();
576 assert_eq!(empty_data, retrieved_data);
577 assert_eq!("msgpack", format);
578 }
579
580 #[test]
581 fn test_metrics_collection_on_store() {
582 let storage = ByteStorage::new(None);
583 let test_data = b"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa".to_vec(); storage.store(&test_data, None).unwrap();
587 let metrics = storage.get_last_metrics();
588
589 assert!(metrics.compression_ratio > 0.0); }
592
593 #[test]
594 fn test_metrics_collection_on_retrieve() {
595 let storage = ByteStorage::new(None);
596 let test_data = b"test data for retrieval metrics";
597
598 let stored = storage.store(test_data, None).unwrap();
600 storage.retrieve(&stored).unwrap();
601 let metrics = storage.get_last_metrics();
602
603 assert!(metrics.compression_ratio > 0.0); }
606}
607
608#[cfg(kani)]
612mod kani_proofs {
613 use super::*;
614
615 #[kani::proof]
618 #[kani::unwind(10)] fn verify_checksum_detects_corruption() {
620 let checksum_a: [u8; 8] = kani::any();
622 let mut checksum_b = checksum_a;
623
624 let byte_index: usize = kani::any();
626 let bit_index: usize = kani::any();
627 kani::assume(byte_index < 8);
628 kani::assume(bit_index < 8);
629
630 checksum_b[byte_index] ^= 1 << bit_index;
631
632 assert_ne!(checksum_a, checksum_b);
634 }
635
636 #[kani::proof]
640 #[kani::unwind(3)]
641 fn verify_decompression_bomb_protection() {
642 let compressed_size: u64 = kani::any();
644 let original_size: u64 = kani::any();
645
646 kani::assume(compressed_size > 0 && compressed_size <= 1000);
648 kani::assume(original_size > 0);
649
650 let max_allowed = MAX_COMPRESSION_RATIO.checked_mul(compressed_size);
654
655 if let Some(max) = max_allowed {
657 let would_reject = original_size > max;
658 let exceeds_ratio = original_size > MAX_COMPRESSION_RATIO * compressed_size;
659 assert_eq!(would_reject, exceeds_ratio);
660 } else {
661 assert!(true); }
664 }
665
666 #[kani::proof]
669 #[kani::unwind(3)]
670 fn verify_input_size_limits() {
671 let size: usize = kani::any();
672
673 kani::assume(size <= MAX_UNCOMPRESSED_SIZE + 100);
675
676 let exceeds_limit = size > MAX_UNCOMPRESSED_SIZE;
678 let should_reject = size > MAX_UNCOMPRESSED_SIZE;
679
680 assert_eq!(exceeds_limit, should_reject);
681 }
682
683 #[kani::proof]
686 #[kani::unwind(3)]
687 fn verify_compressed_size_limits() {
688 let compressed_size: usize = kani::any();
689
690 kani::assume(compressed_size <= MAX_COMPRESSED_SIZE + 100);
692
693 let exceeds_limit = compressed_size > MAX_COMPRESSED_SIZE;
695 let should_reject = compressed_size > MAX_COMPRESSED_SIZE;
696
697 assert_eq!(exceeds_limit, should_reject);
698 }
699
700 #[kani::proof]
703 #[kani::unwind(3)]
704 fn verify_compression_ratio_calculation_safety() {
705 let original_size: u64 = kani::any();
706 let compressed_size: u64 = kani::any();
707
708 kani::assume(compressed_size > 0);
710 kani::assume(compressed_size <= 10000);
711 kani::assume(original_size <= 100_000_000); let result = MAX_COMPRESSION_RATIO.checked_mul(compressed_size);
715
716 if let Some(max_allowed) = result {
718 let is_bomb = original_size > max_allowed;
720
721 if original_size <= max_allowed {
724 assert!(!is_bomb);
725 } else {
726 assert!(is_bomb);
727 }
728 }
729
730 }
732}