1use crate::{Error, Result};
4use std::io::Read;
5use thiserror::Error;
6
7#[derive(Error, Debug, Clone)]
9pub enum CompressionBombError {
10 #[error("Compression ratio exceeded: {ratio:.2}x > {max_ratio:.2}x")]
12 RatioExceeded {
13 ratio: f64,
15 max_ratio: f64,
17 },
18
19 #[error("Decompressed size exceeded: {size} bytes > {max_size} bytes")]
21 SizeExceeded {
22 size: usize,
24 max_size: usize,
26 },
27
28 #[error("Compression depth exceeded: {depth} > {max_depth}")]
30 DepthExceeded {
31 depth: usize,
33 max_depth: usize,
35 },
36}
37
38#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
45pub struct CompressionBombConfig {
46 pub max_ratio: f64,
49 pub max_decompressed_size: usize,
51 pub max_compressed_size: usize,
55 pub max_compression_depth: usize,
57 pub check_interval_bytes: usize,
59}
60
61impl Default for CompressionBombConfig {
62 fn default() -> Self {
63 Self {
64 max_ratio: 300.0,
67 max_decompressed_size: 100 * 1024 * 1024, max_compressed_size: 100 * 1024 * 1024, max_compression_depth: 3,
70 check_interval_bytes: 64 * 1024, }
72 }
73}
74
75impl CompressionBombConfig {
76 pub fn high_security() -> Self {
78 Self {
79 max_ratio: 20.0,
80 max_decompressed_size: 10 * 1024 * 1024, max_compressed_size: 10 * 1024 * 1024, max_compression_depth: 2,
83 check_interval_bytes: 32 * 1024, }
85 }
86
87 pub fn low_memory() -> Self {
89 Self {
90 max_ratio: 50.0,
91 max_decompressed_size: 5 * 1024 * 1024, max_compressed_size: 5 * 1024 * 1024, max_compression_depth: 2,
94 check_interval_bytes: 16 * 1024, }
96 }
97
98 pub fn high_throughput() -> Self {
100 Self {
101 max_ratio: 1000.0,
102 max_decompressed_size: 500 * 1024 * 1024, max_compressed_size: 500 * 1024 * 1024, max_compression_depth: 5,
105 check_interval_bytes: 128 * 1024, }
107 }
108}
109
110#[derive(Debug)]
112pub struct CompressionBombProtector<R: Read> {
113 inner: R,
114 config: CompressionBombConfig,
115 compressed_size: usize,
116 decompressed_size: usize,
117 bytes_since_check: usize,
118 compression_depth: usize,
119}
120
121impl<R: Read> CompressionBombProtector<R> {
122 pub fn new(inner: R, config: CompressionBombConfig, compressed_size: usize) -> Self {
124 Self {
125 inner,
126 config,
127 compressed_size,
128 decompressed_size: 0,
129 bytes_since_check: 0,
130 compression_depth: 0,
131 }
132 }
133
134 pub fn with_depth(
136 inner: R,
137 config: CompressionBombConfig,
138 compressed_size: usize,
139 depth: usize,
140 ) -> Result<Self> {
141 if depth > config.max_compression_depth {
142 return Err(Error::SecurityError(
143 CompressionBombError::DepthExceeded {
144 depth,
145 max_depth: config.max_compression_depth,
146 }
147 .to_string(),
148 ));
149 }
150
151 Ok(Self {
152 inner,
153 config,
154 compressed_size,
155 decompressed_size: 0,
156 bytes_since_check: 0,
157 compression_depth: depth,
158 })
159 }
160
161 fn check_limits(&self) -> Result<()> {
163 if self.decompressed_size > self.config.max_decompressed_size {
165 return Err(Error::SecurityError(
166 CompressionBombError::SizeExceeded {
167 size: self.decompressed_size,
168 max_size: self.config.max_decompressed_size,
169 }
170 .to_string(),
171 ));
172 }
173
174 if self.compressed_size > 0 && self.decompressed_size > 0 {
176 let ratio = self.decompressed_size as f64 / self.compressed_size as f64;
177 if ratio > self.config.max_ratio {
178 return Err(Error::SecurityError(
179 CompressionBombError::RatioExceeded {
180 ratio,
181 max_ratio: self.config.max_ratio,
182 }
183 .to_string(),
184 ));
185 }
186 }
187
188 Ok(())
189 }
190
191 pub fn stats(&self) -> CompressionStats {
193 let ratio = if self.compressed_size > 0 {
194 self.decompressed_size as f64 / self.compressed_size as f64
195 } else {
196 0.0
197 };
198
199 CompressionStats {
200 compressed_size: self.compressed_size,
201 decompressed_size: self.decompressed_size,
202 ratio,
203 compression_depth: self.compression_depth,
204 }
205 }
206
207 pub fn into_inner(self) -> R {
209 self.inner
210 }
211}
212
213impl<R: Read> Read for CompressionBombProtector<R> {
214 fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
215 let bytes_read = self.inner.read(buf)?;
216
217 self.decompressed_size += bytes_read;
218 self.bytes_since_check += bytes_read;
219
220 if self.bytes_since_check >= self.config.check_interval_bytes {
222 if let Err(e) = self.check_limits() {
223 return Err(std::io::Error::new(
224 std::io::ErrorKind::InvalidData,
225 e.to_string(),
226 ));
227 }
228 self.bytes_since_check = 0;
229 }
230
231 Ok(bytes_read)
232 }
233}
234
235#[derive(Debug, Clone)]
237pub struct CompressionStats {
238 pub compressed_size: usize,
240 pub decompressed_size: usize,
242 pub ratio: f64,
244 pub compression_depth: usize,
246}
247
248pub struct CompressionBombDetector {
250 config: CompressionBombConfig,
251}
252
253impl Default for CompressionBombDetector {
254 fn default() -> Self {
255 Self::new(CompressionBombConfig::default())
256 }
257}
258
259impl CompressionBombDetector {
260 pub fn new(config: CompressionBombConfig) -> Self {
262 Self { config }
263 }
264
265 pub fn validate_pre_decompression(&self, compressed_size: usize) -> Result<()> {
272 if compressed_size > self.config.max_compressed_size {
273 return Err(Error::SecurityError(format!(
274 "Compressed data size {} exceeds maximum allowed compressed size {}",
275 compressed_size, self.config.max_compressed_size
276 )));
277 }
278 Ok(())
279 }
280
281 pub fn protect_reader<R: Read>(
283 &self,
284 reader: R,
285 compressed_size: usize,
286 ) -> CompressionBombProtector<R> {
287 CompressionBombProtector::new(reader, self.config.clone(), compressed_size)
288 }
289
290 pub fn protect_nested_reader<R: Read>(
292 &self,
293 reader: R,
294 compressed_size: usize,
295 depth: usize,
296 ) -> Result<CompressionBombProtector<R>> {
297 CompressionBombProtector::with_depth(reader, self.config.clone(), compressed_size, depth)
298 }
299
300 pub fn validate_result(&self, compressed_size: usize, decompressed_size: usize) -> Result<()> {
302 if decompressed_size > self.config.max_decompressed_size {
303 return Err(Error::SecurityError(
304 CompressionBombError::SizeExceeded {
305 size: decompressed_size,
306 max_size: self.config.max_decompressed_size,
307 }
308 .to_string(),
309 ));
310 }
311
312 if compressed_size > 0 {
313 let ratio = decompressed_size as f64 / compressed_size as f64;
314 if ratio > self.config.max_ratio {
315 return Err(Error::SecurityError(
316 CompressionBombError::RatioExceeded {
317 ratio,
318 max_ratio: self.config.max_ratio,
319 }
320 .to_string(),
321 ));
322 }
323 }
324
325 Ok(())
326 }
327}
328
329#[cfg(test)]
330mod tests {
331 use super::*;
332 use std::io::Cursor;
333
334 #[test]
335 fn test_compression_bomb_config() {
336 let config = CompressionBombConfig::default();
337 assert!(config.max_ratio > 0.0);
338 assert!(config.max_decompressed_size > 0);
339
340 let high_sec = CompressionBombConfig::high_security();
341 assert!(high_sec.max_ratio < config.max_ratio);
342
343 let low_mem = CompressionBombConfig::low_memory();
344 assert!(low_mem.max_decompressed_size < config.max_decompressed_size);
345
346 let high_throughput = CompressionBombConfig::high_throughput();
347 assert!(high_throughput.max_decompressed_size > config.max_decompressed_size);
348 }
349
350 #[test]
351 fn test_compression_bomb_detector() {
352 let detector = CompressionBombDetector::default();
353
354 assert!(detector.validate_pre_decompression(1024).is_ok());
356 assert!(detector.validate_result(1024, 10 * 1024).is_ok());
357 }
358
359 #[test]
360 fn test_size_limit_exceeded() {
361 let config = CompressionBombConfig {
362 max_decompressed_size: 1024,
363 ..Default::default()
364 };
365 let detector = CompressionBombDetector::new(config);
366
367 let result = detector.validate_result(100, 2048);
369 assert!(result.is_err());
370 let error_msg = result.unwrap_err().to_string();
371 assert!(error_msg.contains("Size exceeded") || error_msg.contains("Security error"));
372 }
373
374 #[test]
375 fn test_ratio_limit_exceeded() {
376 let config = CompressionBombConfig {
377 max_ratio: 10.0,
378 ..Default::default()
379 };
380 let detector = CompressionBombDetector::new(config);
381
382 let result = detector.validate_result(100, 2000);
384 assert!(result.is_err());
385 assert!(
386 result
387 .unwrap_err()
388 .to_string()
389 .contains("Compression ratio exceeded")
390 );
391 }
392
393 #[test]
394 fn test_protected_reader() {
395 let data = b"Hello, world! This is test data for compression testing.";
396 let cursor = Cursor::new(data.as_slice());
397
398 let config = CompressionBombConfig::default();
399 let mut protector = CompressionBombProtector::new(cursor, config, data.len());
400
401 let mut buffer = Vec::new();
402 let bytes_read = protector.read_to_end(&mut buffer).unwrap();
403
404 assert_eq!(bytes_read, data.len());
405 assert_eq!(buffer.as_slice(), data);
406
407 let stats = protector.stats();
408 assert_eq!(stats.compressed_size, data.len());
409 assert_eq!(stats.decompressed_size, data.len());
410 assert!((stats.ratio - 1.0).abs() < 0.01); }
412
413 #[test]
414 fn test_protected_reader_size_limit() {
415 let data = vec![0u8; 2048]; let cursor = Cursor::new(data);
417
418 let config = CompressionBombConfig {
419 max_decompressed_size: 1024, check_interval_bytes: 512, ..Default::default()
422 };
423
424 let mut protector = CompressionBombProtector::new(cursor, config, 100); let mut buffer = vec![0u8; 2048];
427 let result = protector.read(&mut buffer);
428
429 if result.is_ok() {
431 let result2 = protector.read(&mut buffer[512..]);
433 assert!(result2.is_err());
434 } else {
435 assert!(result.is_err());
437 }
438 }
439
440 #[test]
441 fn test_compression_depth_limit() {
442 let data = b"test data";
443 let cursor = Cursor::new(data.as_slice());
444
445 let config = CompressionBombConfig {
446 max_compression_depth: 2,
447 ..Default::default()
448 };
449
450 let protector = CompressionBombProtector::with_depth(cursor, config.clone(), data.len(), 2);
452 assert!(protector.is_ok());
453
454 let cursor2 = Cursor::new(data.as_slice());
456 let result = CompressionBombProtector::with_depth(cursor2, config, data.len(), 3);
457 assert!(result.is_err());
458 }
459
460 #[test]
461 fn test_zero_compressed_size_handling() {
462 let detector = CompressionBombDetector::default();
463
464 assert!(detector.validate_result(0, 1024).is_ok());
466 }
467
468 #[test]
469 fn test_stats_calculation() {
470 let data = b"test";
471 let cursor = Cursor::new(data.as_slice());
472
473 let protector = CompressionBombProtector::new(cursor, CompressionBombConfig::default(), 2);
474 let stats = protector.stats();
475
476 assert_eq!(stats.compressed_size, 2);
477 assert_eq!(stats.decompressed_size, 0); assert_eq!(stats.ratio, 0.0);
479 assert_eq!(stats.compression_depth, 0);
480 }
481
482 #[test]
483 fn test_stats_with_zero_compressed_size() {
484 let data = b"test";
485 let cursor = Cursor::new(data.as_slice());
486
487 let protector = CompressionBombProtector::new(cursor, CompressionBombConfig::default(), 0);
489 let stats = protector.stats();
490
491 assert_eq!(stats.compressed_size, 0);
492 assert_eq!(stats.ratio, 0.0); }
494
495 #[test]
496 fn test_into_inner() {
497 let data = b"test data";
498 let cursor = Cursor::new(data.as_slice());
499 let original_position = cursor.position();
500
501 let protector =
502 CompressionBombProtector::new(cursor, CompressionBombConfig::default(), data.len());
503
504 let inner = protector.into_inner();
506 assert_eq!(inner.position(), original_position);
507 }
508
509 #[test]
510 fn test_protect_nested_reader_success() {
511 let detector = CompressionBombDetector::new(CompressionBombConfig {
512 max_compression_depth: 3,
513 ..Default::default()
514 });
515
516 let data = b"nested compression test";
517 let cursor = Cursor::new(data.as_slice());
518
519 let result = detector.protect_nested_reader(cursor, data.len(), 1);
521 assert!(result.is_ok());
522
523 let protector = result.unwrap();
524 let stats = protector.stats();
525 assert_eq!(stats.compression_depth, 1);
526 }
527
528 #[test]
529 fn test_protect_nested_reader_depth_exceeded() {
530 let detector = CompressionBombDetector::new(CompressionBombConfig {
531 max_compression_depth: 2,
532 ..Default::default()
533 });
534
535 let data = b"nested compression test";
536 let cursor = Cursor::new(data.as_slice());
537
538 let result = detector.protect_nested_reader(cursor, data.len(), 3);
540 assert!(result.is_err());
541
542 let error_msg = result.unwrap_err().to_string();
543 assert!(
544 error_msg.contains("Compression depth exceeded")
545 || error_msg.contains("Security error")
546 );
547 }
548
549 #[test]
550 fn test_validate_pre_decompression_size_exceeded() {
551 let config = CompressionBombConfig {
552 max_compressed_size: 1024,
553 ..Default::default()
554 };
555 let detector = CompressionBombDetector::new(config);
556
557 let result = detector.validate_pre_decompression(2048);
559 assert!(result.is_err());
560
561 let error_msg = result.unwrap_err().to_string();
562 assert!(error_msg.contains("exceeds maximum allowed"));
563 }
564
565 #[test]
566 fn test_validate_pre_decompression_success() {
567 let detector = CompressionBombDetector::default();
568
569 let result = detector.validate_pre_decompression(1024);
571 assert!(result.is_ok());
572 }
573
574 #[test]
575 fn test_protected_reader_stats_after_read() {
576 let data = b"Hello, world!";
577 let cursor = Cursor::new(data.as_slice());
578
579 let compressed_size = 5; let mut protector = CompressionBombProtector::new(
581 cursor,
582 CompressionBombConfig::default(),
583 compressed_size,
584 );
585
586 let mut buffer = Vec::new();
587 protector.read_to_end(&mut buffer).unwrap();
588
589 let stats = protector.stats();
590 assert_eq!(stats.compressed_size, compressed_size);
591 assert_eq!(stats.decompressed_size, data.len());
592
593 let expected_ratio = data.len() as f64 / compressed_size as f64;
594 assert!((stats.ratio - expected_ratio).abs() < 0.01);
595 }
596
597 #[test]
598 fn test_compression_bomb_error_display() {
599 let ratio_err = CompressionBombError::RatioExceeded {
600 ratio: 150.5,
601 max_ratio: 100.0,
602 };
603 assert!(ratio_err.to_string().contains("150.5"));
604 assert!(ratio_err.to_string().contains("100.0"));
605
606 let size_err = CompressionBombError::SizeExceeded {
607 size: 2048,
608 max_size: 1024,
609 };
610 assert!(size_err.to_string().contains("2048"));
611 assert!(size_err.to_string().contains("1024"));
612
613 let depth_err = CompressionBombError::DepthExceeded {
614 depth: 5,
615 max_depth: 3,
616 };
617 assert!(depth_err.to_string().contains("5"));
618 assert!(depth_err.to_string().contains("3"));
619 }
620
621 #[test]
622 fn test_detector_default() {
623 let detector1 = CompressionBombDetector::default();
624 let detector2 = CompressionBombDetector::new(CompressionBombConfig::default());
625
626 assert_eq!(detector1.config.max_ratio, detector2.config.max_ratio);
628 assert_eq!(
629 detector1.config.max_decompressed_size,
630 detector2.config.max_decompressed_size
631 );
632 }
633
634 #[test]
635 fn test_slow_drip_decompression_bomb() {
636 let config = CompressionBombConfig {
638 max_decompressed_size: 10_000,
639 check_interval_bytes: 1000, ..Default::default()
641 };
642
643 let data = vec![0u8; 15_000];
645 let cursor = Cursor::new(data);
646
647 let mut protector = CompressionBombProtector::new(cursor, config, 100);
648
649 let mut buffer = [0u8; 1024];
650 let mut total_read = 0;
651 let mut detected = false;
652
653 loop {
655 match protector.read(&mut buffer) {
656 Ok(0) => break, Ok(n) => {
658 total_read += n;
659 }
660 Err(e) => {
661 let err_str = e.to_string();
664 assert!(
665 err_str.contains("Size exceeded") || err_str.contains("Security"),
666 "Expected size limit error, got: {}",
667 err_str
668 );
669 detected = true;
670 break;
671 }
672 }
673 }
674
675 assert!(detected, "Slow-drip bomb should be detected");
676 assert!(total_read < 15_000, "Should not read all data");
677 }
678
679 #[test]
680 fn test_integer_overflow_protection_in_ratio() {
681 let detector = CompressionBombDetector::default();
682
683 let result = detector.validate_result(1, usize::MAX);
685 assert!(result.is_err());
686 }
687
688 #[test]
689 fn test_integer_overflow_protection_in_size() {
690 let config = CompressionBombConfig {
691 max_decompressed_size: usize::MAX - 1,
692 ..Default::default()
693 };
694 let detector = CompressionBombDetector::new(config);
695
696 let result = detector.validate_result(100, usize::MAX);
698 assert!(result.is_err());
699 }
700
701 #[test]
702 fn test_boundary_max_decompressed_size() {
703 let max_size = 10_000;
704 let config = CompressionBombConfig {
705 max_decompressed_size: max_size,
706 ..Default::default()
707 };
708 let detector = CompressionBombDetector::new(config);
709
710 assert!(detector.validate_result(100, max_size).is_ok());
712
713 assert!(detector.validate_result(100, max_size + 1).is_err());
715 }
716
717 #[test]
718 fn test_boundary_max_ratio() {
719 let max_ratio = 50.0;
720 let config = CompressionBombConfig {
721 max_ratio,
722 ..Default::default()
723 };
724 let detector = CompressionBombDetector::new(config);
725
726 let compressed = 100;
727 let at_limit = (compressed as f64 * max_ratio) as usize;
728
729 assert!(detector.validate_result(compressed, at_limit).is_ok());
731
732 assert!(
734 detector
735 .validate_result(compressed, at_limit + 100)
736 .is_err()
737 );
738 }
739
740 #[test]
741 fn test_boundary_max_compression_depth() {
742 let max_depth = 5;
743 let config = CompressionBombConfig {
744 max_compression_depth: max_depth,
745 ..Default::default()
746 };
747
748 let data = b"test";
749 let cursor = Cursor::new(data.as_slice());
750
751 let result =
753 CompressionBombProtector::with_depth(cursor, config.clone(), data.len(), max_depth);
754 assert!(result.is_ok());
755
756 let cursor2 = Cursor::new(data.as_slice());
758 let result2 =
759 CompressionBombProtector::with_depth(cursor2, config, data.len(), max_depth + 1);
760 assert!(result2.is_err());
761 }
762
763 #[test]
764 fn test_nested_compression_attack_simulation() {
765 let detector = CompressionBombDetector::new(CompressionBombConfig {
767 max_compression_depth: 2,
768 max_decompressed_size: 10_000,
769 ..Default::default()
770 });
771
772 let layer1_data = vec![0u8; 1000]; let cursor1 = Cursor::new(layer1_data.clone());
775
776 let protector1 = detector.protect_nested_reader(cursor1, 100, 1);
777 assert!(protector1.is_ok());
778
779 let cursor2 = Cursor::new(layer1_data.clone());
781 let protector2 = detector.protect_nested_reader(cursor2, 100, 2);
782 assert!(protector2.is_ok());
783
784 let cursor3 = Cursor::new(layer1_data);
786 let protector3 = detector.protect_nested_reader(cursor3, 100, 3);
787 assert!(protector3.is_err());
788 }
789
790 #[test]
791 fn test_check_limits_called_at_intervals() {
792 let check_interval = 100;
793 let config = CompressionBombConfig {
794 max_decompressed_size: 500,
795 check_interval_bytes: check_interval,
796 max_ratio: 10.0,
797 ..Default::default()
798 };
799
800 let data = vec![0u8; 600];
802 let cursor = Cursor::new(data);
803
804 let mut protector = CompressionBombProtector::new(cursor, config, 10); let mut buffer = [0u8; 50]; let mut total_read = 0;
808 let mut error_occurred = false;
809
810 loop {
811 match protector.read(&mut buffer) {
812 Ok(0) => break,
813 Ok(n) => {
814 total_read += n;
815 if total_read > 500 {
817 break;
819 }
820 }
821 Err(_) => {
822 error_occurred = true;
823 break;
824 }
825 }
826 }
827
828 assert!(error_occurred, "Should detect bomb during periodic checks");
829 }
830
831 #[test]
832 fn test_ratio_calculation_with_large_numbers() {
833 let detector = CompressionBombDetector::new(CompressionBombConfig {
834 max_ratio: 100.0,
835 ..Default::default()
836 });
837
838 let compressed = 1_000_000;
840 let decompressed = 50_000_000; assert!(detector.validate_result(compressed, decompressed).is_ok());
843
844 let decompressed_bad = 150_000_000;
846 assert!(
847 detector
848 .validate_result(compressed, decompressed_bad)
849 .is_err()
850 );
851 }
852
853 #[test]
854 fn test_protected_reader_multiple_small_reads() {
855 let data = vec![1u8; 5000];
857 let cursor = Cursor::new(data);
858
859 let config = CompressionBombConfig {
860 max_decompressed_size: 10_000,
861 check_interval_bytes: 1000,
862 ..Default::default()
863 };
864
865 let mut protector = CompressionBombProtector::new(cursor, config, 5000);
866
867 let mut buffer = [0u8; 10];
869 let mut total = 0;
870
871 while let Ok(n) = protector.read(&mut buffer) {
872 if n == 0 {
873 break;
874 }
875 total += n;
876 }
877
878 assert_eq!(total, 5000);
879 let stats = protector.stats();
880 assert_eq!(stats.decompressed_size, 5000);
881 }
882
883 #[test]
884 fn test_error_on_exact_check_interval_boundary() {
885 let check_interval = 1000;
886 let config = CompressionBombConfig {
887 max_decompressed_size: 1500,
888 check_interval_bytes: check_interval,
889 ..Default::default()
890 };
891
892 let data = vec![0u8; 2000];
894 let cursor = Cursor::new(data);
895
896 let mut protector = CompressionBombProtector::new(cursor, config, 100);
897
898 let mut buffer = [0u8; 1000]; let mut detected = false;
900
901 loop {
902 match protector.read(&mut buffer) {
903 Ok(0) => break,
904 Ok(_) => {}
905 Err(_) => {
906 detected = true;
907 break;
908 }
909 }
910 }
911
912 assert!(detected);
913 }
914
915 #[test]
916 fn test_config_serialization_roundtrip() {
917 let config = CompressionBombConfig {
918 max_ratio: 123.45,
919 max_decompressed_size: 999_888,
920 max_compressed_size: 512_000,
921 max_compression_depth: 7,
922 check_interval_bytes: 16_384,
923 };
924
925 let json = serde_json::to_string(&config).unwrap();
927
928 let deserialized: CompressionBombConfig = serde_json::from_str(&json).unwrap();
930
931 assert_eq!(config.max_ratio, deserialized.max_ratio);
932 assert_eq!(
933 config.max_decompressed_size,
934 deserialized.max_decompressed_size
935 );
936 assert_eq!(
937 config.max_compression_depth,
938 deserialized.max_compression_depth
939 );
940 assert_eq!(
941 config.check_interval_bytes,
942 deserialized.check_interval_bytes
943 );
944 }
945
946 #[test]
947 fn test_all_preset_configs() {
948 let default_cfg = CompressionBombConfig::default();
950 let high_sec = CompressionBombConfig::high_security();
951 let low_mem = CompressionBombConfig::low_memory();
952 let high_throughput = CompressionBombConfig::high_throughput();
953
954 assert!(high_sec.max_ratio < default_cfg.max_ratio);
956 assert!(high_sec.max_decompressed_size < default_cfg.max_decompressed_size);
957
958 assert!(low_mem.max_decompressed_size < default_cfg.max_decompressed_size);
960
961 assert!(high_throughput.max_ratio > default_cfg.max_ratio);
963 assert!(high_throughput.max_decompressed_size > default_cfg.max_decompressed_size);
964 }
965
966 #[test]
967 fn test_protect_reader_basic_usage() {
968 let detector = CompressionBombDetector::default();
969 let data = b"test data for protect_reader";
970 let cursor = Cursor::new(data.as_slice());
971
972 let mut protector = detector.protect_reader(cursor, data.len());
973
974 let mut buffer = Vec::new();
975 let bytes_read = protector.read_to_end(&mut buffer).unwrap();
976
977 assert_eq!(bytes_read, data.len());
978 assert_eq!(buffer.as_slice(), data);
979
980 let stats = protector.stats();
981 assert_eq!(stats.compressed_size, data.len());
982 assert_eq!(stats.decompressed_size, data.len());
983 }
984
985 #[test]
986 fn test_protect_reader_with_size_limit() {
987 let config = CompressionBombConfig {
988 max_decompressed_size: 500,
989 check_interval_bytes: 100,
990 ..Default::default()
991 };
992 let detector = CompressionBombDetector::new(config);
993
994 let data = vec![0u8; 1000];
995 let cursor = Cursor::new(data);
996
997 let mut protector = detector.protect_reader(cursor, 50);
998
999 let mut buffer = [0u8; 200];
1000 let mut error_occurred = false;
1001
1002 loop {
1003 match protector.read(&mut buffer) {
1004 Ok(0) => break,
1005 Ok(_) => {}
1006 Err(_) => {
1007 error_occurred = true;
1008 break;
1009 }
1010 }
1011 }
1012
1013 assert!(error_occurred, "protect_reader should detect size limit");
1014 }
1015
1016 struct FailingReader {
1017 fail_after: usize,
1018 bytes_read: usize,
1019 }
1020
1021 impl FailingReader {
1022 fn new(fail_after: usize) -> Self {
1023 Self {
1024 fail_after,
1025 bytes_read: 0,
1026 }
1027 }
1028 }
1029
1030 impl Read for FailingReader {
1031 fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
1032 if self.bytes_read >= self.fail_after {
1033 return Err(std::io::Error::new(
1034 std::io::ErrorKind::BrokenPipe,
1035 "simulated read failure",
1036 ));
1037 }
1038 let to_read = std::cmp::min(buf.len(), self.fail_after - self.bytes_read);
1039 for b in buf.iter_mut().take(to_read) {
1040 *b = 0;
1041 }
1042 self.bytes_read += to_read;
1043 Ok(to_read)
1044 }
1045 }
1046
1047 #[test]
1048 fn test_inner_reader_error_propagation() {
1049 let failing_reader = FailingReader::new(50);
1050 let config = CompressionBombConfig::default();
1051 let mut protector = CompressionBombProtector::new(failing_reader, config, 100);
1052
1053 let mut buffer = [0u8; 100];
1054
1055 let result1 = protector.read(&mut buffer);
1056 assert!(result1.is_ok());
1057 assert_eq!(result1.unwrap(), 50);
1058
1059 let result2 = protector.read(&mut buffer);
1060 assert!(result2.is_err());
1061 let err = result2.unwrap_err();
1062 assert_eq!(err.kind(), std::io::ErrorKind::BrokenPipe);
1063 assert!(err.to_string().contains("simulated read failure"));
1064 }
1065
1066 #[test]
1067 fn test_check_limits_with_zero_compressed_size_and_data_read() {
1068 let config = CompressionBombConfig {
1069 max_decompressed_size: 1000,
1070 check_interval_bytes: 50,
1071 ..Default::default()
1072 };
1073
1074 let data = vec![0u8; 100];
1075 let cursor = Cursor::new(data);
1076
1077 let mut protector = CompressionBombProtector::new(cursor, config, 0);
1078
1079 let mut buffer = [0u8; 60];
1080
1081 let result = protector.read(&mut buffer);
1082 assert!(result.is_ok());
1083 assert_eq!(result.unwrap(), 60);
1084
1085 let stats = protector.stats();
1086 assert_eq!(stats.compressed_size, 0);
1087 assert_eq!(stats.decompressed_size, 60);
1088 assert_eq!(stats.ratio, 0.0);
1089 }
1090
1091 #[test]
1092 fn test_check_limits_ratio_ok_branch() {
1093 let config = CompressionBombConfig {
1094 max_ratio: 100.0,
1095 max_decompressed_size: 10_000,
1096 check_interval_bytes: 50,
1097 ..Default::default()
1098 };
1099
1100 let data = vec![0u8; 100];
1101 let cursor = Cursor::new(data);
1102
1103 let mut protector = CompressionBombProtector::new(cursor, config, 50);
1104
1105 let mut buffer = [0u8; 60];
1106
1107 let result = protector.read(&mut buffer);
1108 assert!(result.is_ok());
1109 assert_eq!(result.unwrap(), 60);
1110
1111 let stats = protector.stats();
1112 assert_eq!(stats.decompressed_size, 60);
1113 assert!((stats.ratio - 1.2).abs() < 0.01);
1114 }
1115}