From 3e2f5263bc847be2d6b6a21a2d38fa65c5b8ffbd Mon Sep 17 00:00:00 2001
From: Wez Furlong <wez@wezfurlong.org>
Date: Fri, 4 Feb 2022 08:58:23 -0700
Subject: [PATCH] centralize image data size sanity checking, add to kitty

refs: https://github.com/wez/wezterm/issues/1031#issuecomment-1030056063
---
 term/src/terminalstate/image.rs | 34 +++++++++++++++++++++++++++++++++
 term/src/terminalstate/iterm.rs | 32 ++-----------------------------
 term/src/terminalstate/kitty.rs |  4 ++++
 term/src/terminalstate/sixel.rs | 13 +++----------
 4 files changed, 43 insertions(+), 40 deletions(-)

diff --git a/term/src/terminalstate/image.rs b/term/src/terminalstate/image.rs
index a7d17fbb1..2b6f0e64e 100644
--- a/term/src/terminalstate/image.rs
+++ b/term/src/terminalstate/image.rs
@@ -238,3 +238,37 @@ impl TerminalState {
         }
     }
 }
+
+pub(crate) fn check_image_dimensions(width: u32, height: u32) -> anyhow::Result<()> {
+    const MAX_IMAGE_SIZE: u32 = 100_000_000;
+    let size = width.saturating_mul(height).saturating_mul(4);
+    if size > MAX_IMAGE_SIZE {
+        anyhow::bail!(
+            "Ignoring image data {}x{} because {} bytes > max allowed {}",
+            width,
+            height,
+            size,
+            MAX_IMAGE_SIZE
+        );
+    }
+    Ok(())
+}
+
+pub(crate) struct ImageInfo {
+    pub width: u32,
+    pub height: u32,
+    pub format: image::ImageFormat,
+}
+
+pub(crate) fn dimensions(data: &[u8]) -> anyhow::Result<ImageInfo> {
+    let reader = image::io::Reader::new(std::io::Cursor::new(data)).with_guessed_format()?;
+    let format = reader
+        .format()
+        .ok_or_else(|| anyhow::anyhow!("unknown format!?"))?;
+    let (width, height) = reader.into_dimensions()?;
+    Ok(ImageInfo {
+        width,
+        height,
+        format,
+    })
+}
diff --git a/term/src/terminalstate/iterm.rs b/term/src/terminalstate/iterm.rs
index 705a396fc..31689bfbb 100644
--- a/term/src/terminalstate/iterm.rs
+++ b/term/src/terminalstate/iterm.rs
@@ -21,26 +21,6 @@ impl TerminalState {
             return;
         }
 
-        struct Info {
-            width: u32,
-            height: u32,
-            format: ImageFormat,
-        }
-
-        fn dimensions(data: &[u8]) -> anyhow::Result<Info> {
-            let reader =
-                image::io::Reader::new(std::io::Cursor::new(data)).with_guessed_format()?;
-            let format = reader
-                .format()
-                .ok_or_else(|| anyhow::anyhow!("unknown format!?"))?;
-            let (width, height) = reader.into_dimensions()?;
-            Ok(Info {
-                width,
-                height,
-                format,
-            })
-        }
-
         let info = match dimensions(&image.data) {
             Ok(dims) => dims,
             Err(e) => {
@@ -54,16 +34,8 @@ impl TerminalState {
             }
         };
 
-        const MAX_IMAGE_SIZE: u32 = 100_000_000;
-        let size = info.width.saturating_mul(info.height).saturating_mul(4);
-        if size > MAX_IMAGE_SIZE {
-            log::error!(
-                "Ignoring iterm image data {}x{} because {} bytes > max allowed {}",
-                info.width,
-                info.height,
-                size,
-                MAX_IMAGE_SIZE
-            );
+        if let Err(err) = check_image_dimensions(info.width, info.height) {
+            log::error!("{}", err);
             return;
         }
 
diff --git a/term/src/terminalstate/kitty.rs b/term/src/terminalstate/kitty.rs
index 4d09c8fd6..110ab4c1c 100644
--- a/term/src/terminalstate/kitty.rs
+++ b/term/src/terminalstate/kitty.rs
@@ -784,6 +784,8 @@ impl TerminalState {
                     }
                 };
 
+                check_image_dimensions(width, height)?;
+
                 let data = match transmit.format {
                     Some(KittyImageFormat::Rgb) => {
                         let img = DynamicImage::ImageRgb8(
@@ -808,6 +810,8 @@ impl TerminalState {
                 ImageDataType::new_single_frame(width, height, data)
             }
             Some(KittyImageFormat::Png) => {
+                let info = dimensions(&data)?;
+                check_image_dimensions(info.width, info.height)?;
                 let decoded = image::load_from_memory(&data).context("decode png")?;
                 let (width, height) = decoded.dimensions();
                 let data = decoded.into_rgba8().into_vec();
diff --git a/term/src/terminalstate/sixel.rs b/term/src/terminalstate/sixel.rs
index 7172960a3..9e2c7c12e 100644
--- a/term/src/terminalstate/sixel.rs
+++ b/term/src/terminalstate/sixel.rs
@@ -9,16 +9,9 @@ use termwiz::image::ImageDataType;
 impl TerminalState {
     pub(crate) fn sixel(&mut self, sixel: Box<Sixel>) {
         let (width, height) = sixel.dimensions();
-        const MAX_IMAGE_SIZE: u32 = 100_000_000;
-        let size = width.saturating_mul(height).saturating_mul(4);
-        if size > MAX_IMAGE_SIZE {
-            log::error!(
-                "Ignoring sixel image data {}x{} because {} bytes > max allowed {}",
-                width,
-                height,
-                size,
-                MAX_IMAGE_SIZE
-            );
+
+        if let Err(err) = check_image_dimensions(width, height) {
+            log::error!("{}", err);
             return;
         }