fix: cap clipboard image reads

3 files changed, 195 insertions, 12 deletions

diff --git a/src/platform/linux.rs b/src/platform/linux.rs
index 4d2ed049..8e5ad7a6 100644
--- a/src/platform/linux.rs
+++ b/src/platform/linux.rs
@@ -4,7 +4,10 @@ use std::{
     process::{Command, Stdio},
 };
 
-use super::{ClipboardCommand, ClipboardImage, ForegroundJob, ForegroundProcess, Signal};
+use super::{
+    read_limited_reader, ClipboardCommand, ClipboardImage, ForegroundJob, ForegroundProcess,
+    LimitedRead, Signal,
+};
 
 /// Collect the foreground terminal job for a given child PID.
 pub fn foreground_job(child_pid: u32) -> Option<ForegroundJob> {
@@ -221,17 +224,53 @@ fn run_notification_command(mut command: Command) -> std::io::Result<bool> {
 }
 
 fn read_clipboard_image_with_command(program: &str, args: &[&str]) -> Option<Vec<u8>> {
-    let output = Command::new(program)
-        .args(args)
+    let mut command = Command::new(program);
+    command.args(args);
+    read_clipboard_image_with_spawned_command(command)
+}
+
+fn read_clipboard_image_with_spawned_command(command: Command) -> Option<Vec<u8>> {
+    read_clipboard_image_with_spawned_command_max(
+        command,
+        crate::server::protocol::MAX_CLIPBOARD_IMAGE_PAYLOAD,
+    )
+}
+
+fn read_clipboard_image_with_spawned_command_max(
+    mut command: Command,
+    max_bytes: usize,
+) -> Option<Vec<u8>> {
+    let mut child = command
         .stdin(Stdio::null())
+        .stdout(Stdio::piped())
         .stderr(Stdio::null())
-        .output()
+        .spawn()
         .ok()?;
+    let stdout = child.stdout.take()?;
+
+    let read = match read_limited_reader(stdout, max_bytes) {
+        Ok(read) => read,
+        Err(_) => {
+            let _ = child.kill();
+            let _ = child.wait();
+            return None;
+        }
+    };
 
-    if output.status.success() && !output.stdout.is_empty() {
-        Some(output.stdout)
-    } else {
-        None
+    if read == LimitedRead::Oversized {
+        let _ = child.kill();
+        let _ = child.wait();
+        return None;
+    }
+
+    let status = child.wait().ok()?;
+    if !status.success() {
+        return None;
+    }
+
+    match read {
+        LimitedRead::Complete(bytes) => Some(bytes),
+        LimitedRead::Empty | LimitedRead::Oversized => None,
     }
 }
 
@@ -330,6 +369,28 @@ mod tests {
     }
 
     #[test]
+    fn read_clipboard_image_with_spawned_command_reads_under_limit() {
+        let mut command = Command::new("sh");
+        command.arg("-c").arg("printf image");
+
+        assert_eq!(
+            read_clipboard_image_with_spawned_command_max(command, 16),
+            Some(b"image".to_vec())
+        );
+    }
+
+    #[test]
+    fn read_clipboard_image_with_spawned_command_rejects_over_limit() {
+        let mut command = Command::new("sh");
+        command.arg("-c").arg("printf oversized");
+
+        assert_eq!(
+            read_clipboard_image_with_spawned_command_max(command, 4),
+            None
+        );
+    }
+
+    #[test]
     fn desktop_notification_separates_option_like_titles() {
         let _guard = env_lock().lock().unwrap();
         unsafe {
diff --git a/src/platform/macos.rs b/src/platform/macos.rs
index d0cf43c8..7bec8e26 100644
--- a/src/platform/macos.rs
+++ b/src/platform/macos.rs
@@ -5,7 +5,10 @@ use std::path::{Path, PathBuf};
 use std::process::{Command, Stdio};
 use std::sync::OnceLock;
 
-use super::{ClipboardCommand, ClipboardImage, ForegroundJob, ForegroundProcess, Signal};
+use super::{
+    read_limited_reader, ClipboardCommand, ClipboardImage, ForegroundJob, ForegroundProcess,
+    LimitedRead, Signal,
+};
 
 const PROC_PGRP_ONLY: u32 = 2;
 
@@ -227,10 +230,17 @@ pub fn read_clipboard_image() -> Option<ClipboardImage> {
         return None;
     }
 
-    let bytes = std::fs::read(&path).ok();
+    let bytes = match std::fs::File::open(&path).ok().and_then(|file| {
+        read_limited_reader(file, crate::server::protocol::MAX_CLIPBOARD_IMAGE_PAYLOAD).ok()
+    }) {
+        Some(LimitedRead::Complete(bytes)) => bytes,
+        Some(LimitedRead::Empty | LimitedRead::Oversized) | None => {
+            let _ = std::fs::remove_file(&path);
+            return None;
+        }
+    };
     let _ = std::fs::remove_file(&path);
-    let bytes = bytes?;
-    (!bytes.is_empty()).then_some(ClipboardImage {
+    Some(ClipboardImage {
         bytes,
         extension: "png",
     })
diff --git a/src/platform/mod.rs b/src/platform/mod.rs
index 00fa8833..8cf2b08d 100644
--- a/src/platform/mod.rs
+++ b/src/platform/mod.rs
@@ -36,6 +36,50 @@ pub struct ClipboardImage {
     pub extension: &'static str,
 }
 
+#[derive(Debug, PartialEq, Eq)]
+pub(crate) enum LimitedRead {
+    Empty,
+    Complete(Vec<u8>),
+    Oversized,
+}
+
+pub(crate) fn read_limited_reader(
+    mut reader: impl std::io::Read,
+    max_bytes: usize,
+) -> std::io::Result<LimitedRead> {
+    let mut bytes = Vec::new();
+    let mut buffer = [0_u8; 8192];
+
+    while bytes.len() < max_bytes {
+        let remaining = max_bytes - bytes.len();
+        let read_len = remaining.min(buffer.len());
+        let bytes_read = match reader.read(&mut buffer[..read_len]) {
+            Ok(bytes_read) => bytes_read,
+            Err(err) if err.kind() == std::io::ErrorKind::Interrupted => continue,
+            Err(err) => return Err(err),
+        };
+        if bytes_read == 0 {
+            return if bytes.is_empty() {
+                Ok(LimitedRead::Empty)
+            } else {
+                Ok(LimitedRead::Complete(bytes))
+            };
+        }
+        bytes.extend_from_slice(&buffer[..bytes_read]);
+    }
+
+    let mut sentinel = [0_u8; 1];
+    loop {
+        return match reader.read(&mut sentinel) {
+            Ok(0) if bytes.is_empty() => Ok(LimitedRead::Empty),
+            Ok(0) => Ok(LimitedRead::Complete(bytes)),
+            Ok(_) => Ok(LimitedRead::Oversized),
+            Err(err) if err.kind() == std::io::ErrorKind::Interrupted => continue,
+            Err(err) => Err(err),
+        };
+    }
+}
+
 #[cfg(target_os = "linux")]
 mod linux;
 #[cfg(target_os = "linux")]
@@ -50,3 +94,71 @@ pub use macos::*;
 mod fallback;
 #[cfg(not(any(target_os = "linux", target_os = "macos")))]
 pub use fallback::*;
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn read_limited_reader_returns_complete_data_under_limit() {
+        let input = std::io::Cursor::new(b"image".to_vec());
+        assert_eq!(
+            read_limited_reader(input, 16).expect("limited read"),
+            LimitedRead::Complete(b"image".to_vec())
+        );
+    }
+
+    #[test]
+    fn read_limited_reader_returns_empty_for_empty_input() {
+        let input = std::io::Cursor::new(Vec::<u8>::new());
+        assert_eq!(
+            read_limited_reader(input, 16).expect("limited read"),
+            LimitedRead::Empty
+        );
+    }
+
+    #[test]
+    fn read_limited_reader_accepts_data_exactly_at_limit() {
+        let input = std::io::Cursor::new(b"four".to_vec());
+        assert_eq!(
+            read_limited_reader(input, 4).expect("limited read"),
+            LimitedRead::Complete(b"four".to_vec())
+        );
+    }
+
+    #[test]
+    fn read_limited_reader_rejects_data_over_limit() {
+        let input = std::io::Cursor::new(b"oversized".to_vec());
+        assert_eq!(
+            read_limited_reader(input, 4).expect("limited read"),
+            LimitedRead::Oversized
+        );
+    }
+
+    #[test]
+    fn read_limited_reader_retries_interrupted_reads() {
+        struct InterruptedOnce {
+            interrupted: bool,
+            inner: std::io::Cursor<Vec<u8>>,
+        }
+
+        impl std::io::Read for InterruptedOnce {
+            fn read(&mut self, buffer: &mut [u8]) -> std::io::Result<usize> {
+                if !self.interrupted {
+                    self.interrupted = true;
+                    return Err(std::io::ErrorKind::Interrupted.into());
+                }
+                self.inner.read(buffer)
+            }
+        }
+
+        let input = InterruptedOnce {
+            interrupted: false,
+            inner: std::io::Cursor::new(b"image".to_vec()),
+        };
+        assert_eq!(
+            read_limited_reader(input, 16).expect("limited read"),
+            LimitedRead::Complete(b"image".to_vec())
+        );
+    }
+}