Implement taking a monitor screenshot

4 files changed, 141 insertions, 13 deletions

diff --git a/src/input.rs b/src/input.rs
index c5420623..9ffe0613 100644
--- a/src/input.rs
+++ b/src/input.rs
@@ -20,6 +20,7 @@ enum Action {
     Quit,
     ChangeVt(i32),
     Spawn(String),
+    Screenshot,
     CloseWindow,
     ToggleFullscreen,
     FocusLeft,
@@ -48,6 +49,12 @@ enum Action {
     ToggleFullWidth,
 }
 
+pub enum BackendAction {
+    None,
+    ChangeVt(i32),
+    Screenshot,
+}
+
 pub enum CompositorMod {
     Super,
     Alt,
@@ -88,6 +95,8 @@ fn action(comp_mod: CompositorMod, keysym: KeysymHandle, mods: ModifiersState) -
         KEY_t => Action::Spawn("alacritty".to_owned()),
         KEY_d => Action::Spawn("fuzzel".to_owned()),
         KEY_n => Action::Spawn("nautilus".to_owned()),
+        // Alt + PrtSc = SysRq
+        KEY_Sys_Req | KEY_Print => Action::Screenshot,
         KEY_q => Action::CloseWindow,
         KEY_F => Action::ToggleFullscreen,
         KEY_comma => Action::ConsumeIntoColumn,
@@ -126,10 +135,9 @@ fn action(comp_mod: CompositorMod, keysym: KeysymHandle, mods: ModifiersState) -
 impl Niri {
     pub fn process_input_event<I: InputBackend>(
         &mut self,
-        change_vt: &mut dyn FnMut(i32),
         comp_mod: CompositorMod,
         event: InputEvent<I>,
-    ) {
+    ) -> BackendAction {
         let _span = tracy_client::span!("process_input_event");
         trace!("process_input_event");
 
@@ -167,13 +175,16 @@ impl Niri {
                             self.stop_signal.stop()
                         }
                         Action::ChangeVt(vt) => {
-                            (*change_vt)(vt);
+                            return BackendAction::ChangeVt(vt);
                         }
                         Action::Spawn(command) => {
                             if let Err(err) = Command::new(command).spawn() {
                                 warn!("error spawning alacritty: {err}");
                             }
                         }
+                        Action::Screenshot => {
+                            return BackendAction::Screenshot;
+                        }
                         Action::CloseWindow => {
                             if let Some(window) = self.monitor_set.focus() {
                                 window.toplevel().send_close();
@@ -602,6 +613,8 @@ impl Niri {
             }
             _ => {}
         }
+
+        BackendAction::None
     }
 
     pub fn process_libinput_event(&mut self, event: &mut InputEvent<LibinputInputBackend>) {
diff --git a/src/niri.rs b/src/niri.rs
index d8f53e00..b7361e64 100644
--- a/src/niri.rs
+++ b/src/niri.rs
@@ -2,16 +2,22 @@ use std::collections::HashMap;
 use std::os::unix::io::AsRawFd;
 use std::process::Command;
 use std::sync::{Arc, Mutex};
+use std::thread;
 use std::time::Duration;
 
+use anyhow::Context;
+use directories::UserDirs;
 use sd_notify::NotifyState;
+use smithay::backend::allocator::Fourcc;
 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::{render_elements, AsRenderElements, RenderElementStates};
+use smithay::backend::renderer::element::{
+    render_elements, AsRenderElements, Element, RenderElement, RenderElementStates,
+};
 use smithay::backend::renderer::gles::{GlesRenderer, GlesTexture};
-use smithay::backend::renderer::{ImportAll, Renderer};
+use smithay::backend::renderer::{Bind, ExportMem, Frame, ImportAll, Offscreen, Renderer};
 use smithay::desktop::utils::{
     send_frames_surface_tree, surface_presentation_feedback_flags_from_states,
     take_presentation_feedback_surface_tree, OutputPresentationFeedback,
@@ -32,7 +38,9 @@ use smithay::reexports::wayland_server::backend::{
 };
 use smithay::reexports::wayland_server::protocol::wl_surface::WlSurface;
 use smithay::reexports::wayland_server::{Display, DisplayHandle};
-use smithay::utils::{IsAlive, Logical, Physical, Point, Rectangle, Scale, SERIAL_COUNTER};
+use smithay::utils::{
+    IsAlive, Logical, Physical, Point, Rectangle, Scale, Transform, SERIAL_COUNTER,
+};
 use smithay::wayland::compositor::{with_states, CompositorClientState, CompositorState};
 use smithay::wayland::data_device::DataDeviceState;
 use smithay::wayland::output::OutputManagerState;
@@ -42,6 +50,7 @@ use smithay::wayland::shell::xdg::XdgShellState;
 use smithay::wayland::shm::ShmState;
 use smithay::wayland::socket::ListeningSocketSource;
 use smithay::wayland::tablet_manager::TabletManagerState;
+use time::OffsetDateTime;
 
 use crate::backend::Backend;
 use crate::dbus::mutter_service_channel::ServiceChannel;
@@ -712,6 +721,89 @@ impl Niri {
 
         feedback
     }
+
+    pub fn screenshot(
+        &mut self,
+        renderer: &mut GlesRenderer,
+        output: &Output,
+    ) -> anyhow::Result<()> {
+        let _span = tracy_client::span!("Niri::screenshot");
+
+        let size = output.current_mode().unwrap().size;
+        let output_rect = Rectangle::from_loc_and_size((0, 0), size);
+        let buffer_size = size.to_logical(1).to_buffer(1, Transform::Normal);
+        let fourcc = Fourcc::Abgr8888;
+
+        let texture: GlesTexture = renderer
+            .create_buffer(fourcc, buffer_size)
+            .context("error creating texture")?;
+
+        let elements = self.render(renderer, output);
+
+        renderer.bind(texture).context("error binding texture")?;
+        let mut frame = renderer
+            .render(size, Transform::Normal)
+            .context("error starting frame")?;
+
+        frame
+            .clear([0.1, 0.1, 0.1, 1.], &[output_rect])
+            .context("error clearing")?;
+
+        for element in elements.into_iter().rev() {
+            let src = element.src();
+            let dst = element.geometry(Scale::from(1.));
+            element
+                .draw(&mut frame, src, dst, &[output_rect])
+                .context("error drawing element")?;
+        }
+
+        let sync_point = frame.finish().context("error finishing frame")?;
+        sync_point.wait();
+
+        let mapping = renderer
+            .copy_framebuffer(Rectangle::from_loc_and_size((0, 0), buffer_size), fourcc)
+            .context("error copying framebuffer")?;
+        let copy = renderer
+            .map_texture(&mapping)
+            .context("error mapping texture")?;
+        let pixels = copy.to_vec();
+
+        let dirs = UserDirs::new().context("error retrieving home directory")?;
+        let mut path = dirs.picture_dir().map(|p| p.to_owned()).unwrap_or_else(|| {
+            let mut dir = dirs.home_dir().to_owned();
+            dir.push("Pictures");
+            dir
+        });
+        path.push("Screenshots");
+
+        unsafe {
+            // are you kidding me
+            time::util::local_offset::set_soundness(time::util::local_offset::Soundness::Unsound);
+        };
+
+        let now = OffsetDateTime::now_local().unwrap_or_else(|_| OffsetDateTime::now_utc());
+        let desc = time::macros::format_description!(
+            "Screenshot from [year]-[month]-[day] [hour]-[minute]-[second].png"
+        );
+        let name = now.format(desc).context("error formatting time")?;
+        path.push(name);
+
+        debug!("saving screenshot to {path:?}");
+
+        thread::spawn(move || {
+            if let Err(err) = image::save_buffer(
+                path,
+                &pixels,
+                size.w as u32,
+                size.h as u32,
+                image::ColorType::Rgba8,
+            ) {
+                warn!("error saving screenshot image: {err:?}");
+            }
+        });
+
+        Ok(())
+    }
 }
 
 render_elements! {
diff --git a/src/tty.rs b/src/tty.rs
index 528fb0bd..79b9e3c4 100644
--- a/src/tty.rs
+++ b/src/tty.rs
@@ -31,7 +31,7 @@ use smithay_drm_extras::drm_scanner::{DrmScanEvent, DrmScanner};
 use smithay_drm_extras::edid::EdidInfo;
 
 use crate::backend::Backend;
-use crate::input::CompositorMod;
+use crate::input::{BackendAction, CompositorMod};
 use crate::niri::OutputRenderElements;
 use crate::{LoopData, Niri};
 
@@ -173,10 +173,21 @@ impl Tty {
         event_loop
             .insert_source(input_backend, |mut event, _, data| {
                 let tty = data.tty.as_mut().unwrap();
-                let mut change_vt = |vt| tty.change_vt(vt);
-                data.niri.process_libinput_event(&mut event);
-                data.niri
-                    .process_input_event(&mut change_vt, CompositorMod::Super, event);
+                let niri = &mut data.niri;
+
+                niri.process_libinput_event(&mut event);
+                match niri.process_input_event(CompositorMod::Super, event) {
+                    BackendAction::None => (),
+                    BackendAction::ChangeVt(vt) => tty.change_vt(vt),
+                    BackendAction::Screenshot => {
+                        let active = niri.monitor_set.active_output().cloned();
+                        if let Some(active) = active {
+                            if let Err(err) = niri.screenshot(tty.renderer(), &active) {
+                                warn!("error taking screenshot: {err:?}");
+                            }
+                        }
+                    }
+                };
             })
             .unwrap();
 
diff --git a/src/winit.rs b/src/winit.rs
index c8b29203..69fad6ee 100644
--- a/src/winit.rs
+++ b/src/winit.rs
@@ -12,7 +12,7 @@ use smithay::reexports::winit::window::WindowBuilder;
 use smithay::utils::Transform;
 
 use crate::backend::Backend;
-use crate::input::CompositorMod;
+use crate::input::{BackendAction, CompositorMod};
 use crate::niri::OutputRenderElements;
 use crate::utils::get_monotonic_time;
 use crate::{LoopData, Niri};
@@ -129,6 +129,7 @@ impl Winit {
     }
 
     fn dispatch(&mut self, niri: &mut Niri) {
+        let renderer = self.backend.renderer();
         let res = self
             .winit_event_loop
             .dispatch_new_events(|event| match event {
@@ -145,7 +146,18 @@ impl Winit {
                     niri.output_resized(self.output.clone());
                 }
                 WinitEvent::Input(event) => {
-                    niri.process_input_event(&mut |_| (), CompositorMod::Alt, event)
+                    match niri.process_input_event(CompositorMod::Alt, event) {
+                        BackendAction::None => (),
+                        BackendAction::ChangeVt(_) => (),
+                        BackendAction::Screenshot => {
+                            let active = niri.monitor_set.active_output().cloned();
+                            if let Some(active) = active {
+                                if let Err(err) = niri.screenshot(renderer, &active) {
+                                    warn!("error taking screenshot: {err:?}");
+                                }
+                            }
+                        }
+                    }
                 }
                 WinitEvent::Focus(_) => (),
                 WinitEvent::Refresh => niri.queue_redraw(self.output.clone()),