From ef2e50bb2df0d38f735dcba0a419494da04db642 Mon Sep 17 00:00:00 2001
From: Anna <git@anna.lgbt>
Date: Wed, 23 Aug 2023 00:50:59 -0400
Subject: [PATCH] fix: limit reading to the size specified

---
 src/ttmp_extractor.rs | 24 +++++++++++++++++-------
 1 file changed, 17 insertions(+), 7 deletions(-)

diff --git a/src/ttmp_extractor.rs b/src/ttmp_extractor.rs
index 3c44aca..5cb1445 100644
--- a/src/ttmp_extractor.rs
+++ b/src/ttmp_extractor.rs
@@ -71,16 +71,18 @@ impl<R: Read + Seek> TtmpExtractor<R> {
             let expected = file.mod_size;
 
             let info: SqPackFileInfoHeader = read_struct(&mut data_file, &mut buf)?;
+            let limit = expected - std::mem::size_of::<SqPackFileInfoHeader>();
+            let limited = (&mut data_file).take(limit as u64);
             match info.kind {
                 FileKind::Empty => todo!(),
                 FileKind::Standard => {
-                    Self::extract_standard_file(&info, &mut data_file, &mut writer, &mut buf)?;
+                    Self::extract_standard_file(&info, limited, &mut writer, &mut buf)?;
                 }
                 FileKind::Model => {
-                    Self::extract_model_file(&info, &mut data_file, &mut writer, &mut buf)?;
+                    Self::extract_model_file(&info, limited, &mut writer, &mut buf)?;
                 }
                 FileKind::Texture => {
-                    Self::extract_texture_file(&info, &mut data_file, &mut writer, &mut buf)?;
+                    Self::extract_texture_file(&info, limited, &mut writer, &mut buf)?;
                 }
             }
 
@@ -143,19 +145,22 @@ impl<R: Read> TtmpExtractor<R> {
         let expected = file.mod_size;
 
         let info: SqPackFileInfoHeader = read_struct(&mut reader, &mut buf)?;
+        let limit = expected - std::mem::size_of::<SqPackFileInfoHeader>();
+        let limited = (&mut reader).take(limit as u64);
         match info.kind {
             FileKind::Empty => todo!(),
             FileKind::Standard => {
-                Self::extract_standard_file(&info, &mut reader, &mut writer, &mut buf)?;
+                Self::extract_standard_file(&info, limited, &mut writer, &mut buf)?;
             }
             FileKind::Model => {
-                Self::extract_model_file(&info, &mut reader, &mut writer, &mut buf)?;
+                Self::extract_model_file(&info, limited, &mut writer, &mut buf)?;
             }
             FileKind::Texture => {
-                Self::extract_texture_file(&info, &mut reader, &mut writer, &mut buf)?;
+                Self::extract_texture_file(&info, limited, &mut writer, &mut buf)?;
             }
         }
 
+        debug_assert!(reader.read <= expected);
         if reader.read < expected {
             let to_skip = expected - reader.read;
             Self::skip(&mut reader, &mut buf, to_skip)?;
@@ -433,7 +438,12 @@ impl<R: Read> TtmpExtractor<R> {
         let mut left = amt;
         while left > 0 {
             let to_read = std::cmp::min(left, buf.len());
-            left -= reader.read(&mut buf[..to_read]).map_err(Error::Io)?;
+            let amt_read = reader.read(&mut buf[..to_read]).map_err(Error::Io)?;
+            if amt_read == 0 {
+                break;
+            }
+
+            left -= amt_read;
         }
 
         Ok(())