From 3ee9b64cd9fc506883b1829ead32c242081c68c3 Mon Sep 17 00:00:00 2001
From: ShootingStarDragons <ShootingStarDragons@protonmail.com>
Date: Thu, 28 Mar 2024 16:24:28 +0800
Subject: [PATCH] fix: scale

and wrong use of slurp, slurp only accept args like `-d`, it do not
accept region, just run -d is enough
---
 Cargo.lock             |  1 +
 libwayshot/Cargo.toml  |  1 +
 libwayshot/src/lib.rs  | 45 +++++++++++++++++++++++++++++++++++++++++-
 wayshot/src/cli.rs     |  4 ++--
 wayshot/src/wayshot.rs |  8 ++------
 5 files changed, 50 insertions(+), 9 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 04fcbd3d..8d72667f 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -476,6 +476,7 @@ dependencies = [
  "image",
  "memmap2",
  "nix 0.27.1",
+ "tempfile",
  "thiserror",
  "tracing",
  "wayland-client",
diff --git a/libwayshot/Cargo.toml b/libwayshot/Cargo.toml
index 8a9719b0..3b22fc45 100644
--- a/libwayshot/Cargo.toml
+++ b/libwayshot/Cargo.toml
@@ -18,3 +18,4 @@ thiserror = "1"
 wayland-client = "0.31.1"
 wayland-protocols = { version = "0.31.0", features = ["client", "unstable"] }
 wayland-protocols-wlr = { version = "0.2.0", features = ["client"] }
+tempfile = "3.10.1"
diff --git a/libwayshot/src/lib.rs b/libwayshot/src/lib.rs
index 4e8bc98b..e614299f 100644
--- a/libwayshot/src/lib.rs
+++ b/libwayshot/src/lib.rs
@@ -14,6 +14,7 @@ mod screencopy;
 use std::{
     collections::HashSet,
     fs::File,
+    io::Write,
     os::fd::AsFd,
     sync::atomic::{AtomicBool, Ordering},
     thread,
@@ -446,6 +447,7 @@ impl WayshotConnection {
             }
         };
 
+        let shm = self.globals.bind::<WlShm, _, _>(&qh, 1..=1, ())?;
         for (frame_copy, frame_guard, output_info) in frames {
             tracing::span!(
                 tracing::Level::DEBUG,
@@ -453,6 +455,38 @@ impl WayshotConnection {
                 output = format!("{output_info}")
             )
             .in_scope(|| -> Result<()> {
+                let buffer = {
+                    let output_size = output_info.logical_region.inner.size;
+                    if output_size != frame_copy.frame_format.size {
+                        let file = tempfile::tempfile().unwrap();
+                        let image: DynamicImage = frame_copy.try_into()?;
+                        let image = image::imageops::resize(
+                            &image,
+                            output_size.width,
+                            output_size.height,
+                            image::imageops::FilterType::Triangle,
+                        );
+                        init_overlay(&image, &file);
+                        let pool = shm.create_pool(
+                            file.as_fd(),
+                            (output_size.width * output_size.height * 4) as i32,
+                            &qh,
+                            (),
+                        );
+
+                        pool.create_buffer(
+                            0,
+                            output_size.width as i32,
+                            output_size.height as i32,
+                            (output_size.width * 4) as i32,
+                            frame_copy.frame_format.format,
+                            &qh,
+                            (),
+                        )
+                    } else {
+                        frame_guard.buffer.clone()
+                    }
+                };
                 let surface = compositor.create_surface(&qh, ());
 
                 let layer_surface = layer_shell.get_layer_surface(
@@ -481,7 +515,7 @@ impl WayshotConnection {
 
                 surface.set_buffer_transform(output_info.transform);
                 // surface.set_buffer_scale(output_info.scale());
-                surface.attach(Some(&frame_guard.buffer), 0, 0);
+                surface.attach(Some(&buffer), 0, 0);
 
                 debug!("Committing surface with attached buffer.");
                 surface.commit();
@@ -675,3 +709,12 @@ impl WayshotConnection {
         self.screenshot_outputs(self.get_all_outputs(), cursor_overlay)
     }
 }
+
+fn init_overlay(origin_image: &image::ImageBuffer<image::Rgba<u8>, Vec<u8>>, tmp: &File) {
+    let mut buf = std::io::BufWriter::new(tmp);
+
+    for index in origin_image.pixels() {
+        buf.write_all(&index.0).unwrap();
+    }
+    buf.flush().unwrap();
+}
diff --git a/wayshot/src/cli.rs b/wayshot/src/cli.rs
index 0ad61c23..9b2e4073 100644
--- a/wayshot/src/cli.rs
+++ b/wayshot/src/cli.rs
@@ -28,8 +28,8 @@ pub struct Cli {
     pub log_level: tracing::Level,
 
     /// Arguments to call slurp with for selecting a region
-    #[arg(short, long, value_name = "SLURP_ARGS")]
-    pub slurp: Option<String>,
+    #[arg(short, long)]
+    pub slurp: bool,
 
     /// Enable cursor in screenshots
     #[arg(short, long)]
diff --git a/wayshot/src/wayshot.rs b/wayshot/src/wayshot.rs
index 2d25e098..1caf4932 100644
--- a/wayshot/src/wayshot.rs
+++ b/wayshot/src/wayshot.rs
@@ -66,15 +66,11 @@ fn main() -> Result<()> {
         return Ok(());
     }
 
-    let image_buffer = if let Some(slurp_region) = cli.slurp {
-        let slurp_region = slurp_region.clone();
+    let image_buffer = if cli.slurp {
         wayshot_conn.screenshot_freeze(
             Box::new(move || {
                 || -> Result<LogicalRegion> {
-                    let slurp_output = Command::new("slurp")
-                        .args(slurp_region.split(' '))
-                        .output()?
-                        .stdout;
+                    let slurp_output = Command::new("slurp").arg("-d").output()?.stdout;
 
                     utils::parse_geometry(&String::from_utf8(slurp_output)?)
                 }()