From 6884b5abf890661cf843eca9fadd2e27591f99a9 Mon Sep 17 00:00:00 2001
From: ShootingStarDragons <ShootingStarDragons@protonmail.com>
Date: Sat, 7 Oct 2023 10:34:21 +0800
Subject: [PATCH] chore: remove unsafe from rawfd

---
 libwayshot/src/lib.rs | 139 +++++++++++++++++++++++++++++++++++++-----
 1 file changed, 124 insertions(+), 15 deletions(-)

diff --git a/libwayshot/src/lib.rs b/libwayshot/src/lib.rs
index 2d2a4200..7507ae9a 100644
--- a/libwayshot/src/lib.rs
+++ b/libwayshot/src/lib.rs
@@ -13,10 +13,7 @@ mod screencopy;
 use std::{
     cmp,
     fs::File,
-    os::{
-        fd::{AsFd, AsRawFd},
-        unix::prelude::FromRawFd,
-    },
+    os::fd::AsFd,
     process::exit,
     sync::atomic::{AtomicBool, Ordering},
 };
@@ -175,17 +172,132 @@ impl WayshotConnection {
         fd: T,
         capture_region: Option<CaptureRegion>,
     ) -> Result<FrameFormat> {
-        self.capture_output_frame_shm_fd_inner(cursor_overlay, output, fd, None, capture_region)
+        // Connecting to wayland environment.
+        let mut state = CaptureFrameState {
+            formats: Vec::new(),
+            state: None,
+            buffer_done: AtomicBool::new(false),
+        };
+        let mut event_queue = self.conn.new_event_queue::<CaptureFrameState>();
+        let qh = event_queue.handle();
+
+        // Instantiating screencopy manager.
+        let screencopy_manager = match self.globals.bind::<ZwlrScreencopyManagerV1, _, _>(
+            &qh,
+            3..=3,
+            (),
+        ) {
+            Ok(x) => x,
+            Err(e) => {
+                log::error!("Failed to create screencopy manager. Does your compositor implement ZwlrScreencopy?");
+                log::error!("err: {e}");
+                return Err(Error::ProtocolNotFound(
+                    "ZwlrScreencopy Manager not found".to_string(),
+                ));
+            }
+        };
+
+        // Capture output.
+        let frame: ZwlrScreencopyFrameV1 = if let Some(region) = capture_region {
+            screencopy_manager.capture_output_region(
+                cursor_overlay,
+                output,
+                region.x_coordinate,
+                region.y_coordinate,
+                region.width,
+                region.height,
+                &qh,
+                (),
+            )
+        } else {
+            screencopy_manager.capture_output(cursor_overlay, output, &qh, ())
+        };
+
+        // Empty internal event buffer until buffer_done is set to true which is when the Buffer done
+        // event is fired, aka the capture from the compositor is succesful.
+        while !state.buffer_done.load(Ordering::SeqCst) {
+            event_queue.blocking_dispatch(&mut state)?;
+        }
+
+        log::debug!(
+            "Received compositor frame buffer formats: {:#?}",
+            state.formats
+        );
+        // Filter advertised wl_shm formats and select the first one that matches.
+        let frame_format = state
+            .formats
+            .iter()
+            .find(|frame| {
+                matches!(
+                    frame.format,
+                    wl_shm::Format::Xbgr2101010
+                        | wl_shm::Format::Abgr2101010
+                        | wl_shm::Format::Argb8888
+                        | wl_shm::Format::Xrgb8888
+                        | wl_shm::Format::Xbgr8888
+                )
+            })
+            .copied();
+        log::debug!("Selected frame buffer format: {:#?}", frame_format);
+
+        // Check if frame format exists.
+        let frame_format = match frame_format {
+            Some(format) => format,
+            None => {
+                log::error!("No suitable frame format found");
+                return Err(Error::NoSupportedBufferFormat);
+            }
+        };
+
+        // Bytes of data in the frame = stride * height.
+        let frame_bytes = frame_format.stride * frame_format.height;
+
+        // Create an in memory file and return it's file descriptor.
+
+        // Instantiate shm global.
+        let shm = self.globals.bind::<WlShm, _, _>(&qh, 1..=1, ()).unwrap();
+        let shm_pool = shm.create_pool(fd.as_fd(), frame_bytes as i32, &qh, ());
+        let buffer = shm_pool.create_buffer(
+            0,
+            frame_format.width as i32,
+            frame_format.height as i32,
+            frame_format.stride as i32,
+            frame_format.format,
+            &qh,
+            (),
+        );
+
+        // Copy the pixel data advertised by the compositor into the buffer we just created.
+        frame.copy(&buffer);
+        // On copy the Ready / Failed events are fired by the frame object, so here we check for them.
+        loop {
+            // Basically reads, if frame state is not None then...
+            if let Some(state) = state.state {
+                match state {
+                    FrameState::Failed => {
+                        log::error!("Frame copy failed");
+                        return Err(Error::FramecopyFailed);
+                    }
+                    FrameState::Finished => {
+                        buffer.destroy();
+                        shm_pool.destroy();
+                        return Ok(frame_format);
+                    }
+                }
+            }
+
+            event_queue.blocking_dispatch(&mut state)?;
+        }
     }
 
-    fn capture_output_frame_shm_fd_inner<T: AsFd>(
+    fn capture_output_frame_shm_from_file(
         &self,
         cursor_overlay: i32,
         output: &WlOutput,
-        fd: T,
-        file: Option<&File>,
+        file: &File,
         capture_region: Option<CaptureRegion>,
     ) -> Result<FrameFormat> {
+        let fd = file.as_fd();
         // Connecting to wayland environment.
         let mut state = CaptureFrameState {
             formats: Vec::new(),
@@ -265,9 +377,7 @@ impl WayshotConnection {
 
         // Bytes of data in the frame = stride * height.
         let frame_bytes = frame_format.stride * frame_format.height;
-        if let Some(file) = file {
-            file.set_len(frame_bytes as u64)?;
-        }
+        file.set_len(frame_bytes as u64)?;
         // Create an in memory file and return it's file descriptor.
 
         // Instantiate shm global.
@@ -317,13 +427,12 @@ impl WayshotConnection {
         // Create an in memory file and return it's file descriptor.
         let fd = create_shm_fd()?;
         // Create a writeable memory map backed by a mem_file.
-        let mem_file = unsafe { File::from_raw_fd(fd.as_raw_fd()) };
+        let mem_file = File::from(fd);
 
-        let frame_format = self.capture_output_frame_shm_fd_inner(
+        let frame_format = self.capture_output_frame_shm_from_file(
             cursor_overlay,
             output,
-            fd,
-            Some(&mem_file),
+            &mem_file,
             capture_region,
         )?;