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