From 950f77d35c11f6646815e780c91511f22d7439a4 Mon Sep 17 00:00:00 2001
From: Charles Johnson <charles@charlesjohnson.name>
Date: Sat, 19 Aug 2023 14:39:54 +0100
Subject: [PATCH] support wlroots based compositors

---
 Cargo.toml                      |  1 +
 src/linux/wayland_screenshot.rs | 38 +++++++++++++++++++++++++++------
 2 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 45577da..e49a222 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -29,4 +29,5 @@ features = ["Win32_Foundation", "Win32_Graphics_Gdi"]
 
 [target.'cfg(target_os="linux")'.dependencies]
 dbus = { version = "0.9.7", features = ["vendored"] }
+libwayshot = "0.2"
 xcb = "1.2.1"
diff --git a/src/linux/wayland_screenshot.rs b/src/linux/wayland_screenshot.rs
index 82d0c8e..461e131 100644
--- a/src/linux/wayland_screenshot.rs
+++ b/src/linux/wayland_screenshot.rs
@@ -1,18 +1,20 @@
+use std::{
+    collections::HashMap,
+    env::temp_dir,
+    fs::{self, File},
+    sync::{Arc, Mutex},
+    time::{Duration, SystemTime, UNIX_EPOCH},
+};
+
 use anyhow::{anyhow, Result};
 use dbus::{
     arg::{AppendAll, Iter, IterAppend, PropMap, ReadAll, RefArg, TypeMismatchError, Variant},
     blocking::Connection,
     message::{MatchRule, SignalArgs},
 };
+use libwayshot::{CaptureRegion, WayshotConnection};
 use percent_encoding::percent_decode;
 use png::Decoder;
-use std::{
-    collections::HashMap,
-    env::temp_dir,
-    fs::{self, File},
-    sync::{Arc, Mutex},
-    time::{Duration, SystemTime, UNIX_EPOCH},
-};
 
 #[derive(Debug)]
 pub struct OrgFreedesktopPortalRequestResponse {
@@ -192,10 +194,32 @@ fn org_freedesktop_portal_screenshot(
     Ok(rgba)
 }
 
+fn wlr_screenshot(
+    x_coordinate: i32,
+    y_coordinate: i32,
+    width: i32,
+    height: i32,
+) -> Result<Vec<u8>> {
+    let wayshot_connection = WayshotConnection::new()?;
+    Ok(wayshot_connection
+        .screenshot(
+            CaptureRegion {
+                x_coordinate,
+                y_coordinate,
+                width,
+                height,
+            },
+            false,
+        )?
+        .into_raw())
+}
+
 // TODO: 失败后尝试删除文件
 pub fn wayland_screenshot(x: i32, y: i32, width: i32, height: i32) -> Result<Vec<u8>> {
     let conn = Connection::new_session()?;
 
+    // TODO: work out if compositor is wlroots before attempting anything else
     org_gnome_shell_screenshot(&conn, x, y, width, height)
         .or_else(|_| org_freedesktop_portal_screenshot(&conn, x, y, width, height))
+        .or_else(|_| wlr_screenshot(x, y, width, height))
 }