Implement safety around zip-bomb-style attacks in replay files.

It will still be up to the library user to decide what reasonable limits
are for their usecase, although the defaults, while conservative, should
catch the really malicious cases.

Author: tec27

.gitattributes

@@ -1,3 +1,3 @@
 * text=auto eol=lf
-
 *.rep filter=lfs diff=lfs merge=lfs -text
+broodrep/testdata/** filter=lfs diff=lfs merge=lfs -text

broodrep/src/compression.rs

@@ -0,0 +1,209 @@
+use std::{
+    io::{Read, Take},
+    time::{Duration, Instant},
+};
+
+use thiserror::Error;
+
+#[derive(Debug, Copy, Clone)]
+pub struct DecompressionConfig {
+    /// Maximum bytes to decompress (default: 100MB)
+    pub max_decompressed_size: u64,
+    /// Maximum compression ratio allowed (default: 1000:1)
+    pub max_compression_ratio: f64,
+    /// Maximum time to spend decompressing (default: 30 seconds)
+    pub max_decompression_time: Duration,
+}
+
+impl Default for DecompressionConfig {
+    fn default() -> Self {
+        Self {
+            max_decompressed_size: 100 * 1024 * 1024, // 100MB
+            max_compression_ratio: 1000.0,
+            max_decompression_time: Duration::from_secs(30),
+        }
+    }
+}
+
+#[derive(Debug, Error)]
+pub enum DecompressionError {
+    #[error("Decompressed size limit exceeded")]
+    SizeLimitExceeded,
+    #[error("Compression ratio too high")]
+    CompressionRatioExceeded,
+    #[error("Decompression timeout")]
+    TimeoutExceeded,
+    #[error("IO error: {0}")]
+    Io(#[from] std::io::Error),
+}
+
+/// A wrapper around decompression implementations that implement [Read], providing various
+/// mechanisms for limiting decompression to avoid things like zip bombs. For best results, a config
+/// should be used that takes into account the characteristics of the data being decompressed.
+pub struct SafeDecompressor<R: Read> {
+    inner: Take<R>,
+    max_decompressed_size: u64,
+    max_ratio: f64,
+    max_time: Duration,
+    input_size: Option<u64>,
+
+    start_time: Option<Instant>,
+    bytes_read: u64,
+}
+
+impl<R: Read> SafeDecompressor<R> {
+    /// Constructs a new SafeDecompressor wrapping the given [Read] implementation. `input_size` is
+    /// the size of the compressed input data in bytes, if known. If not specified, compression
+    /// ratio limits will not apply.
+    pub fn new(reader: R, config: DecompressionConfig, input_size: Option<u64>) -> Self {
+        Self {
+            inner: reader.take(config.max_decompressed_size),
+            max_decompressed_size: config.max_decompressed_size,
+            max_ratio: config.max_compression_ratio,
+            max_time: config.max_decompression_time,
+            input_size,
+
+            start_time: None,
+            bytes_read: 0,
+        }
+    }
+}
+
+impl<R: Read> Read for SafeDecompressor<R> {
+    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
+        if self.start_time.is_none() {
+            self.start_time = Some(Instant::now());
+        }
+
+        if self.start_time.map(|t| t.elapsed()).unwrap_or_default() > self.max_time {
+            return Err(std::io::Error::new(
+                std::io::ErrorKind::TimedOut,
+                DecompressionError::TimeoutExceeded,
+            ));
+        }
+
+        let bytes_read = self.inner.read(buf)?;
+        self.bytes_read = self.bytes_read.saturating_add(bytes_read as u64);
+
+        if bytes_read == 0 && self.bytes_read == self.max_decompressed_size {
+            // EOF and we've reached the max the Take will allow, try to read 1 more byte to see if
+            // there was more data
+            self.inner.set_limit(1);
+            let mut buf = [0; 1];
+            if let Ok(1) = self.read(&mut buf) {
+                return Err(std::io::Error::new(
+                    std::io::ErrorKind::InvalidData,
+                    DecompressionError::SizeLimitExceeded,
+                ));
+            }
+        }
+
+        if let Some(input_size) = self.input_size {
+            let ratio = self.bytes_read as f64 / input_size as f64;
+            if ratio > self.max_ratio {
+                return Err(std::io::Error::new(
+                    std::io::ErrorKind::InvalidData,
+                    DecompressionError::CompressionRatioExceeded,
+                ));
+            }
+        }
+
+        Ok(bytes_read)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::io::Write as _;
+
+    use explode::ExplodeReader;
+    use flate2::bufread::ZlibDecoder;
+
+    use super::*;
+
+    const IMPLODE_BOMB: &[u8] = include_bytes!("../testdata/all_zeroes_1MB.impode");
+
+    #[test]
+    fn implode_bomb_size() {
+        let config = DecompressionConfig {
+            max_decompressed_size: 1000 * 1024, // slightly less than 1MB
+            ..Default::default()
+        };
+        let mut safe_reader = SafeDecompressor::new(
+            ExplodeReader::new(IMPLODE_BOMB),
+            config,
+            Some(IMPLODE_BOMB.len() as u64),
+        );
+        let mut out = Vec::new();
+        let result = safe_reader.read_to_end(&mut out);
+        assert!(result.is_err());
+        let err = result.unwrap_err();
+        assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
+        let err = err.downcast::<DecompressionError>().unwrap();
+        assert!(matches!(err, DecompressionError::SizeLimitExceeded));
+    }
+
+    #[test]
+    fn implode_bomb_ratio() {
+        let config = DecompressionConfig {
+            max_compression_ratio: 100.0,
+            ..Default::default()
+        };
+        let mut safe_reader = SafeDecompressor::new(
+            ExplodeReader::new(IMPLODE_BOMB),
+            config,
+            Some(IMPLODE_BOMB.len() as u64),
+        );
+        let mut out = Vec::new();
+        let result = safe_reader.read_to_end(&mut out);
+        assert!(result.is_err());
+        let err = result.unwrap_err();
+        assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
+        let err = err.downcast::<DecompressionError>().unwrap();
+        assert!(matches!(err, DecompressionError::CompressionRatioExceeded));
+    }
+
+    fn create_zlib_bomb() -> Vec<u8> {
+        let mut encoder =
+            flate2::write::ZlibEncoder::new(Vec::new(), flate2::Compression::default());
+        let data = vec![0u8; 1024 * 1024];
+        encoder.write_all(&data).unwrap();
+        encoder.finish().unwrap()
+    }
+
+    #[test]
+    fn zlib_bomb_size() {
+        let config = DecompressionConfig {
+            max_decompressed_size: 1000 * 1024, // slightly less than 1MB
+            ..Default::default()
+        };
+        let data = create_zlib_bomb();
+        let mut safe_reader =
+            SafeDecompressor::new(ZlibDecoder::new(&data[..]), config, Some(data.len() as u64));
+        let mut out = Vec::new();
+        let result = safe_reader.read_to_end(&mut out);
+        assert!(result.is_err());
+        let err = result.unwrap_err();
+        assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
+        let err = err.downcast::<DecompressionError>().unwrap();
+        assert!(matches!(err, DecompressionError::SizeLimitExceeded));
+    }
+
+    #[test]
+    fn zlib_bomb_ratio() {
+        let config = DecompressionConfig {
+            max_compression_ratio: 1000.0,
+            ..Default::default()
+        };
+        let data = create_zlib_bomb();
+        let mut safe_reader =
+            SafeDecompressor::new(ZlibDecoder::new(&data[..]), config, Some(data.len() as u64));
+        let mut out = Vec::new();
+        let result = safe_reader.read_to_end(&mut out);
+        assert!(result.is_err());
+        let err = result.unwrap_err();
+        assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
+        let err = err.downcast::<DecompressionError>().unwrap();
+        assert!(matches!(err, DecompressionError::CompressionRatioExceeded));
+    }
+}

broodrep/src/lib.rs

@@ -4,17 +4,23 @@ use std::{
 };
 
 use byteorder::{LittleEndian as LE, ReadBytesExt as _};
+use explode::ExplodeReader;
 use flate2::bufread::ZlibDecoder;
 use thiserror::Error;
 
+use crate::compression::SafeDecompressor;
+pub use crate::compression::{DecompressionConfig, DecompressionError};
+
+mod compression;
+
 #[derive(Error, Debug)]
 pub enum BroodrepError {
     #[error(transparent)]
     IoError(#[from] std::io::Error),
     #[error("malformed header: {0}")]
     MalformedHeader(&'static str),
-    #[error("problem decompressing legacy compressed data: {0}")]
-    LegacyCompressionError(explode::Error),
+    #[error("problem decompressing data: {0}")]
+    Decompression(#[from] DecompressionError),
 }
 
 pub struct Replay<R: Read + Seek> {
@@ -24,7 +30,21 @@ pub struct Replay<R: Read + Seek> {
 }
 
 impl<R: Read + Seek> Replay<R> {
-    pub fn new(mut reader: R) -> Result<Self, BroodrepError> {
+    /// Creates a new Replay by parsing data from a [Read] implementation with default settings for
+    /// reading.
+    pub fn new(reader: R) -> Result<Self, BroodrepError> {
+        Self::new_with_decompression_config(reader, DecompressionConfig::default())
+    }
+
+    // TODO(tec27): Would probably be nice to be able to specify limits for the file as a whole as
+    // well
+    /// Creates a new Replay by parsing data from a [Read] implementation with specified settings
+    /// for reading. Note that the limits specified will apply to each chunk individually, rather
+    /// than to the entire replay collectively.
+    pub fn new_with_decompression_config(
+        mut reader: R,
+        config: DecompressionConfig,
+    ) -> Result<Self, BroodrepError> {
         let format = Self::detect_format(&mut reader)?;
 
         reader.seek(SeekFrom::Start(0))?;
@@ -49,7 +69,7 @@ impl<R: Read + Seek> Replay<R> {
         }
 
         // Replay header section
-        let replay_header = Self::read_legacy_section(&mut reader, format)?;
+        let replay_header = Self::read_legacy_section(&mut reader, format, config)?;
         let replay_header = Self::parse_replay_header(&replay_header)?;
 
         Ok(Replay {
@@ -100,12 +120,17 @@ impl<R: Read + Seek> Replay<R> {
         })
     }
 
-    fn read_legacy_section(reader: &mut R, format: ReplayFormat) -> Result<Vec<u8>, BroodrepError> {
+    fn read_legacy_section(
+        reader: &mut R,
+        format: ReplayFormat,
+        config: DecompressionConfig,
+    ) -> Result<Vec<u8>, BroodrepError> {
         let header = Self::read_section_header(reader)?;
         // TODO(tec27): Pass a size hint for known sections to avoid reallocations?
-        let mut data = vec![];
+        let mut data = Vec::new();
         for _ in 0..header.num_chunks {
             let size = reader.read_u32::<LE>()?;
+            data.reserve(size as usize);
             // TODO(tec27): Keep a working buffer around to avoid needing to reallocate buffers
             // frequently? Peek the first byte and seek back to avoid needing this allocation at
             // all?
@@ -114,20 +139,23 @@ impl<R: Read + Seek> Replay<R> {
 
             match format {
                 ReplayFormat::Legacy => {
-                    // TODO(tec27): Often our data will be much smaller than the default buffer
-                    // size, so maybe we could use whatever size hint based on the section to
-                    // provide a smaller buffer (and use explode_with_buffer)
-                    let decompressed = explode::explode(&compressed[..])
-                        .map_err(BroodrepError::LegacyCompressionError)?;
-                    data.extend(decompressed);
+                    let mut decoder = SafeDecompressor::new(
+                        ExplodeReader::new(&compressed[..]),
+                        config,
+                        Some(size as u64),
+                    );
+                    decoder.read_to_end(&mut data)?;
                 }
                 ReplayFormat::Modern | ReplayFormat::Modern121 => {
                     if size <= 4 || compressed[0] != 0x78 {
-                        // FIXME: implement implode decoding
                         // Not compressed, we can return it directly
                         data.extend(compressed);
                     } else {
-                        let mut decoder = ZlibDecoder::new(&compressed[..]);
+                        let mut decoder = SafeDecompressor::new(
+                            ZlibDecoder::new(&compressed[..]),
+                            config,
+                            Some(size as u64),
+                        );
                         decoder.read_to_end(&mut data)?;
                     }
                 }
@@ -167,7 +195,6 @@ impl<R: Read + Seek> Replay<R> {
         let game_sub_type = cursor.read_u16::<LE>()?;
 
         cursor.seek(SeekFrom::Current(8))?; // unknown
-        eprint!("unknown cursor position: {:#x}", cursor.position());
 
         let mut host_name = vec![0u8; 25];
         cursor.read_exact(&mut host_name[..24])?;

broodrep/testdata/all_zeroes_1MB.impode

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:232b7b9b284f0b4f3eb3c30f74cfc279b0baf4975c22d0499bc32901ed87b3d1
+size 7121