Add initial impl of org.gnome.Shell.Screenshot

Enough to make the portal all-outputs screenshot work. With this,
Flameshot kinda-works.

Cargo.lock

@@ -1483,6 +1483,7 @@ name = "niri"
 version = "0.1.0"
 dependencies = [
  "anyhow",
+ "async-channel",
  "async-io",
  "bitflags 2.4.0",
  "clap",

Cargo.toml

@@ -28,6 +28,7 @@ tracy-client = { version = "0.16.2", default-features = false }
 xcursor = "0.3.4"
 zbus = { version = "3.14.1" }
 async-io = "1.13.0"
+async-channel = "1.9.0"
 
 [dependencies.smithay]
 git = "https://github.com/Smithay/smithay.git"

src/config.rs

@@ -136,7 +136,7 @@ pub struct DebugConfig {
     #[knuffel(child, unwrap(argument), default = 1.)]
     pub animation_slowdown: f64,
     #[knuffel(child)]
-    pub screen_cast_in_non_session_instances: bool,
+    pub dbus_interfaces_in_non_session_instances: bool,
     #[knuffel(child)]
     pub wait_for_frame_completion_before_queueing: bool,
     #[knuffel(child)]
@@ -149,7 +149,7 @@ impl Default for DebugConfig {
     fn default() -> Self {
         Self {
             animation_slowdown: 1.,
-            screen_cast_in_non_session_instances: false,
+            dbus_interfaces_in_non_session_instances: false,
             wait_for_frame_completion_before_queueing: false,
             enable_color_transformations_capability: false,
             enable_overlay_planes: false,

src/dbus/gnome_shell_screenshot.rs

@@ -0,0 +1,57 @@
+use std::path::PathBuf;
+
+use smithay::reexports::calloop;
+use zbus::{dbus_interface, fdo};
+
+pub struct Screenshot {
+    to_niri: calloop::channel::Sender<ScreenshotToNiri>,
+    from_niri: async_channel::Receiver<NiriToScreenshot>,
+}
+
+pub enum ScreenshotToNiri {
+    TakeScreenshot { include_cursor: bool },
+}
+
+pub enum NiriToScreenshot {
+    ScreenshotResult(Option<PathBuf>),
+}
+
+#[dbus_interface(name = "org.gnome.Shell.Screenshot")]
+impl Screenshot {
+    async fn screenshot(
+        &self,
+        include_cursor: bool,
+        _flash: bool,
+        _filename: PathBuf,
+    ) -> fdo::Result<(bool, PathBuf)> {
+        if let Err(err) = self
+            .to_niri
+            .send(ScreenshotToNiri::TakeScreenshot { include_cursor })
+        {
+            warn!("error sending message to niri: {err:?}");
+            return Err(fdo::Error::Failed("internal error".to_owned()));
+        }
+
+        let filename = match self.from_niri.recv().await {
+            Ok(NiriToScreenshot::ScreenshotResult(Some(filename))) => filename,
+            Ok(NiriToScreenshot::ScreenshotResult(None)) => {
+                return Err(fdo::Error::Failed("internal error".to_owned()));
+            }
+            Err(err) => {
+                warn!("error receiving message from niri: {err:?}");
+                return Err(fdo::Error::Failed("internal error".to_owned()));
+            }
+        };
+
+        Ok((true, filename))
+    }
+}
+
+impl Screenshot {
+    pub fn new(
+        to_niri: calloop::channel::Sender<ScreenshotToNiri>,
+        from_niri: async_channel::Receiver<NiriToScreenshot>,
+    ) -> Self {
+        Self { to_niri, from_niri }
+    }
+}

src/dbus/mod.rs

@@ -1,3 +1,4 @@
+pub mod gnome_shell_screenshot;
 pub mod mutter_display_config;
 pub mod mutter_screen_cast;
 pub mod mutter_service_channel;

src/niri.rs

@@ -1,6 +1,8 @@
 use std::cell::RefCell;
+use std::cmp::max;
 use std::collections::HashMap;
 use std::os::unix::io::AsRawFd;
+use std::path::PathBuf;
 use std::process::Command;
 use std::rc::Rc;
 use std::sync::{Arc, Mutex};
@@ -15,6 +17,7 @@ use smithay::backend::renderer::element::surface::{
     render_elements_from_surface_tree, WaylandSurfaceRenderElement,
 };
 use smithay::backend::renderer::element::texture::{TextureBuffer, TextureRenderElement};
+use smithay::backend::renderer::element::utils::{Relocate, RelocateRenderElement};
 use smithay::backend::renderer::element::{
     render_elements, AsRenderElements, Element, Kind, RenderElement, RenderElementStates,
 };
@@ -59,6 +62,7 @@ use zbus::fdo::RequestNameFlags;
 
 use crate::backend::{Backend, Tty, Winit};
 use crate::config::Config;
+use crate::dbus::gnome_shell_screenshot::{self, NiriToScreenshot, ScreenshotToNiri};
 use crate::dbus::mutter_display_config::DisplayConfig;
 use crate::dbus::mutter_screen_cast::{self, ScreenCast, ToNiriMsg};
 use crate::dbus::mutter_service_channel::ServiceChannel;
@@ -351,6 +355,43 @@ impl Niri {
             .unwrap();
         let screen_cast = ScreenCast::new(backend.connectors(), to_niri);
 
+        let (to_niri, from_screenshot) = calloop::channel::channel();
+        let (to_screenshot, from_niri) = async_channel::unbounded();
+        event_loop
+            .insert_source(from_screenshot, move |event, _, data| match event {
+                calloop::channel::Event::Msg(ScreenshotToNiri::TakeScreenshot {
+                    include_cursor,
+                }) => {
+                    let renderer = data.state.backend.renderer();
+                    let on_done = {
+                        let to_screenshot = to_screenshot.clone();
+                        move |path| {
+                            let msg = NiriToScreenshot::ScreenshotResult(Some(path));
+                            if let Err(err) = to_screenshot.send_blocking(msg) {
+                                warn!("error sending path to screenshot: {err:?}");
+                            }
+                        }
+                    };
+
+                    let res =
+                        data.state
+                            .niri
+                            .screenshot_all_outputs(renderer, include_cursor, on_done);
+
+                    if let Err(err) = res {
+                        warn!("error taking a screenshot: {err:?}");
+
+                        let msg = NiriToScreenshot::ScreenshotResult(None);
+                        if let Err(err) = to_screenshot.send_blocking(msg) {
+                            warn!("error sending None to screenshot: {err:?}");
+                        }
+                    }
+                }
+                calloop::channel::Event::Closed => (),
+            })
+            .unwrap();
+        let screenshot = gnome_shell_screenshot::Screenshot::new(to_niri, from_niri);
+
         let mut zbus_conn = None;
         let mut inhibit_power_key_fd = None;
         if std::env::var_os("NOTIFY_SOCKET").is_some() {
@@ -395,25 +436,33 @@ impl Niri {
                 .build()
                 .unwrap();
 
-            if pipewire.is_some() {
+            {
                 let server = conn.object_server();
-                server
-                    .at("/org/gnome/Mutter/ScreenCast", screen_cast.clone())
-                    .unwrap();
-                server
-                    .at(
-                        "/org/gnome/Mutter/DisplayConfig",
-                        DisplayConfig::new(backend.connectors()),
-                    )
-                    .unwrap();
-
                 let flags = RequestNameFlags::AllowReplacement
                     | RequestNameFlags::ReplaceExisting
                     | RequestNameFlags::DoNotQueue;
-                conn.request_name_with_flags("org.gnome.Mutter.ScreenCast", flags)
+
+                server
+                    .at("/org/gnome/Shell/Screenshot", screenshot)
                     .unwrap();
-                conn.request_name_with_flags("org.gnome.Mutter.DisplayConfig", flags)
+                conn.request_name_with_flags("org.gnome.Shell.Screenshot", flags)
                     .unwrap();
+
+                if pipewire.is_some() {
+                    server
+                        .at("/org/gnome/Mutter/ScreenCast", screen_cast.clone())
+                        .unwrap();
+                    server
+                        .at(
+                            "/org/gnome/Mutter/DisplayConfig",
+                            DisplayConfig::new(backend.connectors()),
+                        )
+                        .unwrap();
+                    conn.request_name_with_flags("org.gnome.Mutter.ScreenCast", flags)
+                        .unwrap();
+                    conn.request_name_with_flags("org.gnome.Mutter.DisplayConfig", flags)
+                        .unwrap();
+                }
             }
 
             zbus_conn = Some(conn);
@@ -448,28 +497,37 @@ impl Niri {
                     warn!("error inhibiting power key: {err:?}");
                 }
             }
-        } else if pipewire.is_some() && config_.debug.screen_cast_in_non_session_instances {
+        } else if config_.debug.dbus_interfaces_in_non_session_instances {
             let conn = zbus::blocking::Connection::session().unwrap();
-            {
-                let server = conn.object_server();
-                server
-                    .at("/org/gnome/Mutter/ScreenCast", screen_cast.clone())
-                    .unwrap();
-                server
-                    .at(
-                        "/org/gnome/Mutter/DisplayConfig",
-                        DisplayConfig::new(backend.connectors()),
-                    )
-                    .unwrap();
-            }
-
             let flags = RequestNameFlags::AllowReplacement
                 | RequestNameFlags::ReplaceExisting
                 | RequestNameFlags::DoNotQueue;
-            conn.request_name_with_flags("org.gnome.Mutter.ScreenCast", flags)
-                .unwrap();
-            conn.request_name_with_flags("org.gnome.Mutter.DisplayConfig", flags)
-                .unwrap();
+
+            {
+                let server = conn.object_server();
+
+                server
+                    .at("/org/gnome/Shell/Screenshot", screenshot)
+                    .unwrap();
+                conn.request_name_with_flags("org.gnome.Shell.Screenshot", flags)
+                    .unwrap();
+
+                if pipewire.is_some() {
+                    server
+                        .at("/org/gnome/Mutter/ScreenCast", screen_cast.clone())
+                        .unwrap();
+                    server
+                        .at(
+                            "/org/gnome/Mutter/DisplayConfig",
+                            DisplayConfig::new(backend.connectors()),
+                        )
+                        .unwrap();
+                    conn.request_name_with_flags("org.gnome.Mutter.ScreenCast", flags)
+                        .unwrap();
+                    conn.request_name_with_flags("org.gnome.Mutter.DisplayConfig", flags)
+                        .unwrap();
+                }
+            }
 
             zbus_conn = Some(conn);
         }
@@ -1083,6 +1141,56 @@ impl Niri {
 
         Ok(())
     }
+
+    pub fn screenshot_all_outputs(
+        &mut self,
+        renderer: &mut GlesRenderer,
+        include_pointer: bool,
+        on_done: impl FnOnce(PathBuf) + Send + 'static,
+    ) -> anyhow::Result<()> {
+        let _span = tracy_client::span!("Niri::screenshot_all_outputs");
+
+        let mut elements = vec![];
+        let mut size = Size::from((0, 0));
+
+        let outputs: Vec<_> = self.global_space.outputs().cloned().collect();
+        for output in outputs {
+            let geom = self.global_space.output_geometry(&output).unwrap();
+            let geom = geom.to_physical(1);
+
+            size.w = max(size.w, geom.loc.x + geom.size.w);
+            size.h = max(size.h, geom.loc.y + geom.size.h);
+
+            let output_elements = self.render(renderer, &output, include_pointer);
+            elements.extend(output_elements.into_iter().map(|elem| {
+                RelocateRenderElement::from_element(elem, geom.loc, Relocate::Relative)
+            }));
+        }
+
+        let pixels = render_to_vec(renderer, size, &elements)?;
+
+        let path = make_screenshot_path().context("error making screenshot path")?;
+        debug!("saving screenshot to {path:?}");
+
+        thread::spawn(move || {
+            let res = image::save_buffer(
+                &path,
+                &pixels,
+                size.w as u32,
+                size.h as u32,
+                image::ColorType::Rgba8,
+            );
+
+            if let Err(err) = res {
+                warn!("error saving screenshot image: {err:?}");
+                return;
+            }
+
+            on_done(path);
+        });
+
+        Ok(())
+    }
 }
 
 render_elements! {