Fix integer overflows

Related to #9
vstroebel · Oct 5, 2024 · 9bc1db3 · 9bc1db3
1 parent 999f9ed
commit 9bc1db3
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 9 deletions.
diff --git a/jfifdump/src/error.rs b/jfifdump/src/error.rs
@@ -6,6 +6,11 @@ pub enum JfifError {
     JfifMarkerNotFound,
     InvalidMarker(u8),
     InvalidMarkerLength(usize),
+    InvalidDhtSegmentLength(usize),
+    InvalidDqtSegmentLength(usize),
+    InvalidFrameSegmentLength(usize),
+    InvalidDriLength(usize),
+    InvalidScanHeaderLength(usize),
     IoError(std::io::Error),
 }
 
@@ -22,6 +27,11 @@ impl Display for JfifError {
             JfifMarkerNotFound => write!(f, "Not a JFIF file"),
             InvalidMarker(value) => write!(f, "Invalid marker: 0x{:X}", value),
             InvalidMarkerLength(length) => write!(f, "Invalid length for marker: {}", length),
+            InvalidDhtSegmentLength(length) => write!(f, "Invalid dht segment length: {}", length),
+            InvalidDqtSegmentLength(length) => write!(f, "Invalid dqt segment length: {}", length),
+            InvalidFrameSegmentLength(length) => write!(f, "Invalid dqt segment length: {}", length),
+            InvalidDriLength(length) => write!(f, "Invalid dri length: {}", length),
+            InvalidScanHeaderLength(length) => write!(f, "Invalid scan header length: {}", length),
             IoError(err) => err.fmt(f),
         }
     }

diff --git a/jfifdump/src/reader.rs b/jfifdump/src/reader.rs
@@ -160,6 +160,11 @@ impl<R: Read> Reader<R> {
 
         let num_tables = length / 65;
 
+        let remaining = match length.checked_sub(num_tables * 65) {
+            Some(length) => length,
+            None => return Err(JfifError::InvalidDqtSegmentLength(length)),
+        };
+
         let mut tables = vec![];
 
         for _ in 0..num_tables {
@@ -175,7 +180,6 @@ impl<R: Read> Reader<R> {
             });
         }
 
-        let remaining = length - num_tables * 65;
         if remaining > 0 {
             self.skip(remaining)?;
         }
@@ -184,11 +188,13 @@ impl<R: Read> Reader<R> {
     }
 
     fn read_dht(&mut self) -> Result<Vec<Dht>, JfifError> {
-        let mut length = self.read_length()?;
+        let length = self.read_length()?;
 
         let mut tables = vec![];
 
-        while length > 17 {
+        let mut remaining = length;
+
+        while remaining > 17 {
             let (class, destination) = self.read_u4_tuple()?;
             let mut code_lengths = [0u8; 16];
             self.read_exact(&mut code_lengths)?;
@@ -204,11 +210,14 @@ impl<R: Read> Reader<R> {
                 values,
             });
 
-            length -= 17 + num_codes;
+            remaining = match remaining.checked_sub(17 + num_codes) {
+                Some(length) => length,
+                None => return Err(JfifError::InvalidDhtSegmentLength(length)),
+            }
         }
 
-        if length > 0 {
-            self.skip(length)?;
+        if remaining > 0 {
+            self.skip(remaining)?;
         }
 
         Ok(tables)
@@ -250,7 +259,10 @@ impl<R: Read> Reader<R> {
         let selection_end = self.read_u8()?;
         let (approximation_low, approximation_high) = self.read_u4_tuple()?;
 
-        let remaining = length - 1 - num_components as usize * 2 - 3;
+        let remaining = match length.checked_sub(1 + num_components as usize * 2 + 3) {
+            Some(length) => length,
+            None => return Err(JfifError::InvalidScanHeaderLength(length)),
+        };
 
         if remaining > 0 {
             self.skip(remaining)?;
@@ -308,7 +320,10 @@ impl<R: Read> Reader<R> {
         let length = self.read_length()?;
         let restart = self.read_u16()?;
 
-        let remaining = length - 2;
+        let remaining = match length.checked_sub(2) {
+            Some(length) => length,
+            None => return Err(JfifError::InvalidDriLength(length)),
+        };
 
         if remaining > 0 {
             self.skip(remaining)?;
@@ -341,7 +356,10 @@ impl<R: Read> Reader<R> {
             })
         }
 
-        let remaining = length - 6 - num_components as usize * 3;
+        let remaining = match length.checked_sub(6 + num_components as usize * 3) {
+            Some(length) => length,
+            None => return Err(JfifError::InvalidFrameSegmentLength(length)),
+        };
 
         if remaining > 0 {
             self.skip(remaining)?;