From 2164c1ea5a6793eeda72e4bbd85084b99e810eea Mon Sep 17 00:00:00 2001
From: Ian Douglas Scott <idscott@system76.com>
Date: Fri, 26 Jan 2024 14:13:58 -0800
Subject: [PATCH] Create shm buffer directly instead of using `RawPool`

Currently needed if we want an fd to pass to iced-sctk for subsurfaces.
Rather than a wl_buffer.

It would be good if this logic and the dmabuf type could be shared in
some crate.
---
 Cargo.lock            |  1 +
 Cargo.toml            |  1 +
 src/wayland/buffer.rs | 87 ++++++++++++++++++++++++++++++++++++++-----
 3 files changed, 79 insertions(+), 10 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 4f4e431..25ef548 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -938,6 +938,7 @@ dependencies = [
  "memmap2 0.9.4",
  "once_cell",
  "rust-embed",
+ "rustix 0.38.30",
  "tokio",
  "wayland-protocols 0.31.0",
  "zbus",
diff --git a/Cargo.toml b/Cargo.toml
index 2308507..64473c2 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -26,6 +26,7 @@ itertools = "0.12.0"
 log = "0.4.20"
 i18n-embed-fl = "0.7.0"
 rust-embed = "8.1.0"
+rustix = { version = "0.38.30", features  = ["fs"] }
 
 [dependencies.i18n-embed]
 version = "0.14.1"
diff --git a/src/wayland/buffer.rs b/src/wayland/buffer.rs
index 118aeb3..db8c615 100644
--- a/src/wayland/buffer.rs
+++ b/src/wayland/buffer.rs
@@ -3,23 +3,72 @@ use cctk::{
     screencopy::BufferInfo,
     sctk::shm::raw::RawPool,
     wayland_client::{
-        protocol::{wl_buffer, wl_shm},
+        protocol::{wl_buffer, wl_shm, wl_shm_pool},
         Connection, Dispatch, QueueHandle, WEnum,
     },
 };
 use cosmic::cctk;
 use cosmic::iced::widget::image;
+use memmap2::MmapMut;
+use rustix::{io::Errno, shm::ShmOFlags};
 use std::{
     os::fd::{AsFd, OwnedFd},
     path::{Path, PathBuf},
+    time::{SystemTime, UNIX_EPOCH},
 };
 use wayland_protocols::wp::linux_dmabuf::zv1::client::zwp_linux_buffer_params_v1;
 
 use super::AppData;
 
+#[cfg(target_os = "linux")]
+fn create_memfd() -> rustix::io::Result<OwnedFd> {
+    let fd = rustix::io::retry_on_intr(|| {
+        rustix::fs::memfd_create(
+            "cosmic-workspaces-shm",
+            rustix::fs::MemfdFlags::CLOEXEC | rustix::fs::MemfdFlags::ALLOW_SEALING,
+        )
+    })?;
+    let _ = rustix::fs::fcntl_add_seals(
+        &fd,
+        rustix::fs::SealFlags::SHRINK | rustix::fs::SealFlags::SEAL,
+    );
+    Ok(fd)
+}
+
+fn create_memfile() -> rustix::io::Result<OwnedFd> {
+    #[cfg(target_os = "linux")]
+    if let Ok(fd) = create_memfd() {
+        return Ok(fd);
+    }
+
+    loop {
+        let flags = ShmOFlags::CREATE | ShmOFlags::EXCL | ShmOFlags::RDWR;
+
+        let time = SystemTime::now();
+        let name = format!(
+            "/cosmic-workspaces-shm-{}",
+            time.duration_since(UNIX_EPOCH).unwrap().subsec_nanos()
+        );
+
+        match rustix::shm::shm_open(&name, flags, 0600.into()) {
+            Ok(fd) => match rustix::shm::shm_unlink(&name) {
+                Ok(_) => return Ok(fd),
+                Err(errno) => {
+                    return Err(errno.into());
+                }
+            },
+            Err(Errno::EXIST | Errno::EXIST) => {
+                continue;
+            }
+            Err(err) => return Err(err.into()),
+        }
+    }
+}
+
 enum BufferBacking {
     Shm {
-        pool: RawPool,
+        fd: OwnedFd,
+        mmap: MmapMut,
     },
     Dmabuf {
         fd: OwnedFd,
@@ -36,11 +85,15 @@ pub struct Buffer {
 
 impl AppData {
     fn create_shm_backing(&self, buffer_info: &BufferInfo) -> (BufferBacking, wl_buffer::WlBuffer) {
-        let mut pool = RawPool::new(
-            (buffer_info.stride * buffer_info.height) as usize,
-            &self.shm_state,
-        )
-        .unwrap();
+        let fd = create_memfile().unwrap(); // XXX?
+        rustix::fs::ftruncate(&fd, buffer_info.stride as u64 * buffer_info.height as u64);
+
+        let pool = self.shm_state.wl_shm().create_pool(
+            fd.as_fd(),
+            buffer_info.stride as i32 * buffer_info.height as i32,
+            &self.qh,
+            (),
+        );
 
         let format = wl_shm::Format::try_from(buffer_info.format).unwrap();
         let buffer = pool.create_buffer(
@@ -49,11 +102,13 @@ impl AppData {
             buffer_info.height as i32,
             buffer_info.stride as i32,
             format,
-            (),
             &self.qh,
+            (),
         );
 
-        (BufferBacking::Shm { pool }, buffer)
+        let mmap = unsafe { MmapMut::map_mut(&fd).unwrap() };
+
+        (BufferBacking::Shm { fd, mmap }, buffer)
     }
 
     #[allow(dead_code)]
@@ -178,7 +233,7 @@ impl Buffer {
     #[allow(clippy::wrong_self_convention)]
     pub unsafe fn to_image(&mut self) -> image::Handle {
         let pixels = match &mut self.backing {
-            BufferBacking::Shm { pool } => pool.mmap().to_vec(),
+            BufferBacking::Shm { mmap, .. } => mmap.to_vec(),
             // NOTE: Only will work with linear modifier
             BufferBacking::Dmabuf {
                 fd,
@@ -235,3 +290,15 @@ impl Dispatch<wl_buffer::WlBuffer, ()> for AppData {
         }
     }
 }
+
+impl Dispatch<wl_shm_pool::WlShmPool, ()> for AppData {
+    fn event(
+        _app_data: &mut Self,
+        _shm: &wl_shm_pool::WlShmPool,
+        _event: wl_shm_pool::Event,
+        _: &(),
+        _: &Connection,
+        _qh: &QueueHandle<Self>,
+    ) {
+    }
+}