check bounds

LaihoE · Mar 21, 2024 · 14bd0c0 · 14bd0c0
1 parent bef9227
commit 14bd0c0
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 14 deletions.
diff --git a/src/parser/src/first_pass/parser.rs b/src/parser/src/first_pass/parser.rs
@@ -65,7 +65,7 @@ pub struct FirstPassOutput<'a> {
     pub wanted_players: AHashSet<u64>,
     pub header: AHashMap<String, String>,
 }
-struct Frame {
+pub struct Frame {
     pub size: usize,
     pub frame_starts_at: usize,
     pub is_compressed: bool,
@@ -140,13 +140,14 @@ impl<'a> FirstPassParser<'a> {
         possibly_uncompressed_bytes: &'b [u8],
         frame: &Frame,
     ) -> Result<&'b [u8], DemoParserError> {
-        FirstPassParser::resize_if_needed(buf, decompress_len(possibly_uncompressed_bytes))?;
-
         match frame.is_compressed {
-            true => match SnapDecoder::new().decompress(possibly_uncompressed_bytes, buf) {
-                Ok(idx) => Ok(&buf[..idx]),
-                Err(e) => return Err(DemoParserError::DecompressionFailure(format!("{}", e))),
-            },
+            true => {
+                FirstPassParser::resize_if_needed(buf, decompress_len(possibly_uncompressed_bytes))?;
+                match SnapDecoder::new().decompress(possibly_uncompressed_bytes, buf) {
+                    Ok(idx) => Ok(&buf[..idx]),
+                    Err(e) => return Err(DemoParserError::DecompressionFailure(format!("{}", e))),
+                }
+            }
             false => Ok(possibly_uncompressed_bytes),
         }
     }
@@ -311,6 +312,9 @@ impl<'a> FirstPassParser<'a> {
         Ok(())
     }
     fn handle_short_header(&mut self, file_len: usize, bytes: &[u8]) -> Result<(), DemoParserError> {
+        if bytes.len() < 16 {
+            return Err(DemoParserError::OutOfBytesError);
+        }
         match std::str::from_utf8(&bytes[..8]) {
             Ok(magic) => match magic {
                 "PBDEMS2\0" => {}

diff --git a/src/parser/src/second_pass/parser.rs b/src/parser/src/second_pass/parser.rs
@@ -67,14 +67,15 @@ impl<'a> SecondPassParser<'a> {
             }
 
             let input = &demo_bytes[self.ptr..self.ptr + size as usize];
-            FirstPassParser::resize_if_needed(&mut buf2, decompress_len(input))?;
             self.ptr += size as usize;
-
             let bytes = match is_compressed {
-                true => match SnapDecoder::new().decompress(input, &mut buf2) {
-                    Ok(idx) => &buf2[..idx],
-                    Err(e) => return Err(DemoParserError::DecompressionFailure(format!("{}", e))),
-                },
+                true => {
+                    FirstPassParser::resize_if_needed(&mut buf2, decompress_len(input))?;
+                    match SnapDecoder::new().decompress(input, &mut buf2) {
+                        Ok(idx) => &buf2[..idx],
+                        Err(e) => return Err(DemoParserError::DecompressionFailure(format!("{}", e))),
+                    }
+                }
                 false => input,
             };
 
@@ -120,7 +121,6 @@ impl<'a> SecondPassParser<'a> {
             if buf.len() < size as usize {
                 buf.resize(size as usize, 0)
             }
-
             bitreader.read_n_bytes_mut(size as usize, buf)?;
             let msg_bytes = &buf[..size as usize];
             let ok = match netmessage_type_from_int(msg_type as i32) {