From 6c08800089b79717ff33c0091ac445609501c552 Mon Sep 17 00:00:00 2001
From: KAMADA Ken'ichi <kamada@nanohz.org>
Date: Sun, 24 Jan 2021 19:38:31 +0900
Subject: [PATCH] Cleanup buffer length calculation.

---
 src/jpeg.rs |  8 +++-----
 src/png.rs  |  3 ++-
 src/webp.rs | 13 +++++++++++++
 3 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/src/jpeg.rs b/src/jpeg.rs
index ad6d03d..854b06c 100644
--- a/src/jpeg.rs
+++ b/src/jpeg.rs
@@ -85,12 +85,10 @@ fn get_exif_attr_sub<R>(reader: &mut R)
             _ => {},
         }
         // Read marker segments.
-        let seglen = read16(reader)?;
-        if seglen < 2 {
-            return Err(Error::InvalidFormat("Invalid segment length"));
-        }
+        let len = read16(reader)?.checked_sub(2)
+            .ok_or(Error::InvalidFormat("Invalid segment length"))?;
         let mut seg = Vec::new();
-        reader.by_ref().take(seglen as u64 - 2).read_to_end(&mut seg)?;
+        reader.by_ref().take(len.into()).read_to_end(&mut seg)?;
         if code == marker::APP1 && seg.starts_with(&EXIF_ID) {
             seg.drain(..EXIF_ID.len());
             return Ok(seg);
diff --git a/src/png.rs b/src/png.rs
index 20237a2..e4b4afd 100644
--- a/src/png.rs
+++ b/src/png.rs
@@ -77,7 +77,8 @@ fn get_exif_attr_sub<R>(reader: &mut R)
             return Ok(data);
         }
         // Chunk data and CRC.
-        reader.discard_exact(len + 4)?;
+        reader.discard_exact(len.checked_add(4).ok_or(
+            Error::InvalidFormat("Invalid chunk length"))?)?;
     }
 }
 
diff --git a/src/webp.rs b/src/webp.rs
index 580dacb..6cea825 100644
--- a/src/webp.rs
+++ b/src/webp.rs
@@ -131,6 +131,19 @@ mod tests {
         assert_err_pat!(get_exif_attr(&mut &data[..]), Error::NotFound(_));
     }
 
+    #[test]
+    fn overflowing_parent() {
+        let mut data = b"RIFF\x10\0\0\0WEBPEXIF\x04\0\0\0Exif".to_vec();
+        assert_eq!(get_exif_attr(&mut &data[..]).unwrap(), b"Exif");
+        for x in 0x05..=0x0f {
+            data[4] = x;
+            assert_err_pat!(get_exif_attr(&mut &data[..]),
+                            Error::InvalidFormat(_));
+        }
+        data[4] = 0x04;
+        assert_err_pat!(get_exif_attr(&mut &data[..]), Error::NotFound(_));
+    }
+
     #[test]
     fn odd_at_last_without_padding() {
         let data = b"RIFF\x17\0\0\0WEBPwhat\0\0\0\0EXIF\x03\0\0\0abc";