Add image dimension and data validation to editor

Author: richiemcilroy

apps/desktop/src-tauri/src/screenshot_editor.rs

@@ -23,6 +23,8 @@ use tauri::{
 use tokio::sync::{RwLock, watch};
 use tokio_util::sync::CancellationToken;
 
+const MAX_DIMENSION: u32 = 16_384;
+
 pub struct ScreenshotEditorInstance {
     pub ws_port: u16,
     pub ws_shutdown_token: CancellationToken,
@@ -105,14 +107,51 @@ impl ScreenshotEditorInstances {
                     let pending_frame = pending.and_then(|p| p.remove(&key));
 
                     if let Some(frame) = pending_frame {
-                        let rgb_img =
-                            RgbImage::from_raw(frame.width, frame.height, frame.data).unwrap();
+                        let width = frame.width;
+                        let height = frame.height;
+
+                        if width > MAX_DIMENSION || height > MAX_DIMENSION {
+                            return Err(format!(
+                                "Image dimensions exceed maximum: {width}x{height}"
+                            ));
+                        }
+
+                        let expected_len = width
+                            .checked_mul(height)
+                            .and_then(|p| p.checked_mul(3))
+                            .ok_or_else(|| {
+                                format!("Image dimensions overflow: {width}x{height}")
+                            })?;
+                        let expected_len = usize::try_from(expected_len)
+                            .map_err(|_| format!("Image size too large: {width}x{height}"))?;
+
+                        let data = frame.data;
+
+                        if data.len() != expected_len {
+                            return Err(format!(
+                                "Image data length mismatch: expected {expected_len} bytes for {width}x{height} frame, got {}",
+                                data.len()
+                            ));
+                        }
+
+                        let rgb_img = RgbImage::from_raw(width, height, data).ok_or_else(|| {
+                            format!("Invalid RGB data for {width}x{height} frame")
+                        })?;
                         let rgba_img: image::RgbaImage = rgb_img.convert();
-                        (rgba_img.into_raw(), frame.width, frame.height)
+                        (rgba_img.into_raw(), width, height)
                     } else {
                         let img =
                             image::open(&path).map_err(|e| format!("Failed to open image: {e}"))?;
                         let (w, h) = img.dimensions();
+
+                        if w > MAX_DIMENSION || h > MAX_DIMENSION {
+                            return Err(format!("Image dimensions exceed maximum: {w}x{h}"));
+                        }
+
+                        w.checked_mul(h)
+                            .and_then(|p| p.checked_mul(4))
+                            .ok_or_else(|| format!("Image dimensions overflow: {w}x{h}"))?;
+
                         (img.to_rgba8().into_raw(), w, h)
                     }
                 };