From 7381159578470306a105a2a1698a77e40784809c Mon Sep 17 00:00:00 2001
From: Sooraj S <94284954+CheerfulPianissimo@users.noreply.github.com>
Date: Tue, 26 Mar 2024 19:26:22 +0530
Subject: [PATCH] fix: Clipboard flag ignored if path is qualified to file
 (#108)

* feat(clipboard): fix issue #106 clipboard flag ignored if path is qualified to file

* feat(clipboard): clarify flag description of --clipboard

* feat(clipboard): enable multiline comment description for --clipboard feature flag

Signed-off-by: Shinyzenith <aakashsensharma@gmail.com>

* feat(clipboard): reduce buffer allocations for encoding image

* feat(clipboard): improve code quality

* feat(clipboard): code style change: perform match inside function call

---------

Signed-off-by: Shinyzenith <aakashsensharma@gmail.com>
Co-authored-by: Shinyzenith <aakashsensharma@gmail.com>
---
 wayshot/src/cli.rs     |  4 +--
 wayshot/src/wayshot.rs | 76 +++++++++++++++++++++++++-----------------
 2 files changed, 48 insertions(+), 32 deletions(-)

diff --git a/wayshot/src/cli.rs b/wayshot/src/cli.rs
index 7a5dc1f6..0ad61c23 100644
--- a/wayshot/src/cli.rs
+++ b/wayshot/src/cli.rs
@@ -18,8 +18,8 @@ pub struct Cli {
     #[arg(value_name = "OUTPUT", verbatim_doc_comment)]
     pub file: Option<PathBuf>,
 
-    /// Copy image to clipboard along with [OUTPUT] or stdout.
-    /// Wayshot persists in the background to offer the image till the clipboard is overwritten.
+    /// Copy image to clipboard. Can be used simultaneously with [OUTPUT] or stdout.
+    /// Wayshot persists in the background offering the image till the clipboard is overwritten.
     #[arg(long, verbatim_doc_comment)]
     pub clipboard: bool,
 
diff --git a/wayshot/src/wayshot.rs b/wayshot/src/wayshot.rs
index 4ba2955f..b8a9420b 100644
--- a/wayshot/src/wayshot.rs
+++ b/wayshot/src/wayshot.rs
@@ -1,4 +1,5 @@
 use std::{
+    fs::File,
     io::{stdout, BufWriter, Cursor, Write},
     process::Command,
 };
@@ -126,43 +127,58 @@ fn main() -> Result<()> {
         }
     };
 
+    let mut image_buf: Option<Cursor<Vec<u8>>> = None;
     if let Some(file) = file {
         image_buffer.save(file)?;
-    } else {
+    } else if stdout_print {
         let mut buffer = Cursor::new(Vec::new());
         image_buffer.write_to(&mut buffer, requested_encoding)?;
+        let stdout = stdout();
+        let mut writer = BufWriter::new(stdout.lock());
+        writer.write_all(buffer.get_ref())?;
+        image_buf = Some(buffer);
+    }
 
-        if stdout_print {
-            let stdout = stdout();
-            let mut writer = BufWriter::new(stdout.lock());
-            writer.write_all(buffer.get_ref())?;
-        }
-        if cli.clipboard {
-            let mut opts = Options::new();
-            match unsafe { fork() } {
-                // Having the image persistently available on the clipboard requires a wayshot process to be alive.
-                // Fork the process with a child detached from the main process and have the parent exit
-                Ok(ForkResult::Parent { .. }) => {
-                    return Ok(());
-                }
-                Ok(ForkResult::Child) => {
-                    opts.foreground(true); // Offer the image till something else is available on the clipboard
-                    opts.copy(
-                        Source::Bytes(buffer.into_inner().into()),
-                        MimeType::Autodetect,
-                    )?;
-                }
-                Err(e) => {
-                    tracing::warn!("Fork failed with error: {e}, couldn't offer image on the clipboard persistently.
-                     Use a clipboard manager to record screenshot.");
-                    opts.copy(
-                        Source::Bytes(buffer.into_inner().into()),
-                        MimeType::Autodetect,
-                    )?;
-                }
+    if cli.clipboard {
+        clipboard_daemonize(match image_buf {
+            Some(buf) => buf,
+            None => {
+                let mut buffer = Cursor::new(Vec::new());
+                image_buffer.write_to(&mut buffer, requested_encoding)?;
+                buffer
             }
-        }
+        })?;
     }
 
     Ok(())
 }
+
+/// Daemonize and copy the given buffer containing the encoded image to the clipboard
+fn clipboard_daemonize(buffer: Cursor<Vec<u8>>) -> Result<()> {
+    let mut opts = Options::new();
+    match unsafe { fork() } {
+        // Having the image persistently available on the clipboard requires a wayshot process to be alive.
+        // Fork the process with a child detached from the main process and have the parent exit
+        Ok(ForkResult::Parent { .. }) => {
+            return Ok(());
+        }
+        Ok(ForkResult::Child) => {
+            opts.foreground(true); // Offer the image till something else is available on the clipboard
+            opts.copy(
+                Source::Bytes(buffer.into_inner().into()),
+                MimeType::Autodetect,
+            )?;
+        }
+        Err(e) => {
+            tracing::warn!(
+                "Fork failed with error: {e}, couldn't offer image on the clipboard persistently.
+                 Use a clipboard manager to record screenshot."
+            );
+            opts.copy(
+                Source::Bytes(buffer.into_inner().into()),
+                MimeType::Autodetect,
+            )?;
+        }
+    }
+    Ok(())
+}