Capture screenshots in the compositor instead of grim

compositor/src/handlers/mod.rs

@@ -3,6 +3,8 @@ mod compositor;
 mod input_method;
 mod layer_shell;
 pub(crate) mod screencopy;
+pub(crate) mod screenshot;
+pub(crate) mod screenshot_worker;
 mod xdg_shell;
 
 use smithay::{

compositor/src/handlers/screencopy.rs

@@ -1,7 +1,12 @@
+use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::{Arc, Mutex};
 
+use super::screenshot::{CaptureBuffer, PartOutcome, PartSender};
 use smithay::{
-    backend::{allocator::Fourcc, renderer::ExportMem},
+    backend::{
+        allocator::Fourcc,
+        renderer::{ExportMem, TextureMapping},
+    },
     output::Output,
     reexports::{
         wayland_protocols_wlr::screencopy::v1::server::{
@@ -31,18 +36,120 @@ pub(crate) struct FrameData {
 
 #[derive(Debug)]
 pub(crate) struct PendingScreencopy {
-    pub(crate) frame: ZwlrScreencopyFrameV1,
-    buffer: WlBuffer,
+    sink: CaptureSink,
     pub(crate) output: Output,
     region: Rectangle<i32, Buffer>,
     overlay_cursor: bool,
     with_damage: bool,
 }
 
-pub(crate) fn init_global(display: &DisplayHandle) {
-    let enabled = std::env::var_os("FERESE_ENABLE_SCREENCOPY").is_some_and(|value| value == "1");
+#[derive(Debug)]
+pub(crate) enum CaptureSink {
+    Shm {
+        frame: ZwlrScreencopyFrameV1,
+        buffer: WlBuffer,
+    },
+    Owned {
+        request: u64,
+        part: usize,
+        complete: PartSender,
+        published: AtomicBool,
+    },
+}
 
-    if enabled {
+impl CaptureSink {
+    pub(crate) fn owned(request: u64, part: usize, complete: PartSender) -> Self {
+        Self::Owned {
+            request,
+            part,
+            complete,
+            published: AtomicBool::new(false),
+        }
+    }
+
+    // Every part reaches exactly one terminal state: pixels or an error. A
+    // cancelled part is still answered, otherwise its request would wait
+    // forever for a part that can no longer arrive.
+    pub(crate) fn fail(&self) {
+        match self {
+            CaptureSink::Shm { frame, .. } => frame.failed(),
+            CaptureSink::Owned {
+                request,
+                part,
+                complete,
+                published,
+            } => {
+                if !published.swap(true, Ordering::AcqRel) {
+                    let _ = complete.send(PartOutcome {
+                        request: *request,
+                        part: *part,
+                        result: Err("Screenshot part failed".to_string()),
+                    });
+                }
+            }
+        }
+    }
+
+    // One-shot: a part publishes at most once, so a second completion can
+    // never overwrite pixels that were already handed off.
+    pub(crate) fn publish(&self, buffer: CaptureBuffer) -> bool {
+        let CaptureSink::Owned {
+            request,
+            part,
+            complete,
+            published,
+        } = self
+        else {
+            return false;
+        };
+        if published.swap(true, Ordering::AcqRel) {
+            return false;
+        }
+        complete
+            .send(PartOutcome {
+                request: *request,
+                part: *part,
+                result: Ok(buffer),
+            })
+            .is_ok()
+    }
+}
+
+impl PendingScreencopy {
+    pub(crate) fn owned(
+        request: u64,
+        part: usize,
+        complete: PartSender,
+        output: Output,
+        region: Rectangle<i32, Buffer>,
+    ) -> Self {
+        Self {
+            sink: CaptureSink::owned(request, part, complete),
+            output,
+            region,
+            // The compositor-owned path deliberately omits the cursor.
+            overlay_cursor: false,
+            with_damage: false,
+        }
+    }
+
+    pub(crate) fn fail(&self) {
+        self.sink.fail();
+    }
+}
+
+pub(crate) fn capture_allowed() -> bool {
+    static ALLOWED: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
+    *ALLOWED
+        .get_or_init(|| capture_opted_in(std::env::var_os("FERESE_ENABLE_SCREENCOPY").as_deref()))
+}
+
+fn capture_opted_in(value: Option<&std::ffi::OsStr>) -> bool {
+    value.is_some_and(|value| value == "1")
+}
+
+pub(crate) fn init_global(display: &DisplayHandle) {
+    if capture_allowed() {
         display.create_global::<Ferese, ZwlrScreencopyManagerV1, ()>(3, ());
         tracing::info!("authorized screencopy is enabled for this session");
     }
@@ -161,8 +268,10 @@ impl Ferese {
         };
 
         self.pending_screencopies.push(PendingScreencopy {
-            frame: frame.clone(),
-            buffer,
+            sink: CaptureSink::Shm {
+                frame: frame.clone(),
+                buffer,
+            },
             output,
             region: data.region,
             overlay_cursor: data.overlay_cursor,
@@ -192,14 +301,16 @@ impl Ferese {
                 remaining.push(capture);
                 continue;
             }
-            if !capture.frame.is_alive() || !capture.buffer.is_alive() {
+            if let CaptureSink::Shm { frame, buffer } = &capture.sink
+                && (!frame.is_alive() || !buffer.is_alive())
+            {
                 continue;
             }
 
             framebuffer_binding_changed = true;
             if complete_capture(renderer, framebuffer, &capture, self.start_time.elapsed()).is_err()
             {
-                capture.frame.failed();
+                capture.fail();
             }
         }
 
@@ -283,17 +394,25 @@ fn capture_region(
     capture_region_for_geometry(mode.size, scale, transform, requested)
 }
 
-fn capture_region_for_geometry(
+pub(crate) fn output_logical_size(
+    mode_size: Size<i32, smithay::utils::Physical>,
+    scale: f64,
+    transform: smithay::utils::Transform,
+) -> Size<i32, Logical> {
+    transform
+        .transform_size(mode_size)
+        .to_f64()
+        .to_logical(scale)
+        .to_i32_round()
+}
+
+pub(crate) fn capture_region_for_geometry(
     mode_size: Size<i32, smithay::utils::Physical>,
     scale: f64,
     transform: smithay::utils::Transform,
     requested: Option<Rectangle<i32, Logical>>,
 ) -> Option<Rectangle<i32, Buffer>> {
-    let logical_size = transform
-        .transform_size(mode_size)
-        .to_f64()
-        .to_logical(scale)
-        .to_i32_round();
+    let logical_size = output_logical_size(mode_size, scale, transform);
     let output_region = Rectangle::from_size(logical_size);
     let logical_region = requested
         .unwrap_or(output_region)
@@ -342,50 +461,79 @@ where
         .map_err(|error| {
             tracing::warn!(?error, "failed to copy output framebuffer");
         })?;
+    if mapping.format() != Fourcc::Argb8888 {
+        return Err(());
+    }
     let source = renderer.map_texture(&mapping).map_err(|error| {
         tracing::warn!(?error, "failed to map output framebuffer copy");
     })?;
     let expected = capture.region.size;
-    let copied = with_buffer_contents_mut(&capture.buffer, |destination, length, data| {
-        let bytes = expected.w as usize * expected.h as usize * BYTES_PER_PIXEL;
-        if source.len() < bytes || data.offset < 0 || data.offset as usize + bytes > length {
-            return false;
-        }
+    let stride = (expected.w as usize)
+        .checked_mul(BYTES_PER_PIXEL)
+        .ok_or(())?;
+    let bytes = stride.checked_mul(expected.h as usize).ok_or(())?;
+    if source.len() < bytes {
+        return Err(());
+    }
+    // Smithay's output projection already accounts for OpenGL's Y axis, so these
+    // rows must not be flipped vertically again. They are in output-buffer order
+    // including the output transform, which the conversion stage normalizes.
+    let copied = match &capture.sink {
+        CaptureSink::Shm { buffer, .. } => {
+            with_buffer_contents_mut(buffer, |destination, length, data| {
+                let Ok(offset) = usize::try_from(data.offset) else {
+                    return false;
+                };
+                if offset > length || bytes > length - offset {
+                    return false;
+                }
 
-        // Smithay's output projection already accounts for OpenGL's Y axis.
-        // These rows are in output-buffer order, including its output transform.
-        unsafe {
-            std::ptr::copy_nonoverlapping(
-                source.as_ptr(),
-                destination.add(data.offset as usize),
-                bytes,
-            );
+                // SAFETY: the mapped slice is at least `bytes` long, and the
+                // checks above establish `bytes` writable bytes at `offset`
+                // within the shm buffer. The regions cannot overlap: one is
+                // renderer-owned and the other client-owned.
+                unsafe {
+                    std::ptr::copy_nonoverlapping(source.as_ptr(), destination.add(offset), bytes);
+                }
+                true
+            })
+            .unwrap_or(false)
         }
-        true
-    })
-    .map_err(|_| ())?;
+        CaptureSink::Owned { .. } => {
+            let mut pixels = Vec::new();
+            if pixels.try_reserve_exact(bytes).is_err() {
+                false
+            } else {
+                pixels.extend_from_slice(&source[..bytes]);
+                capture.sink.publish(CaptureBuffer {
+                    width: expected.w,
+                    height: expected.h,
+                    stride,
+                    pixels,
+                })
+            }
+        }
+    };
 
     if !copied {
         return Err(());
     }
 
-    if capture.with_damage {
-        capture
-            .frame
-            .damage(0, 0, expected.w as u32, expected.h as u32);
-    }
-    // TextureMapping::flipped describes importing the mapping as a texture;
-    // using it here would make clients flip the already-correct output rows.
-    capture
-        .frame
-        .flags(zwlr_screencopy_frame_v1::Flags::empty());
+    if let CaptureSink::Shm { frame, .. } = &capture.sink {
+        if capture.with_damage {
+            frame.damage(0, 0, expected.w as u32, expected.h as u32);
+        }
+        // TextureMapping::flipped describes importing the mapping as a texture;
+        // using it here would make clients flip the already-correct output rows.
+        frame.flags(zwlr_screencopy_frame_v1::Flags::empty());
 
-    let seconds = timestamp.as_secs();
-    capture.frame.ready(
-        (seconds >> 32) as u32,
-        seconds as u32,
-        timestamp.subsec_nanos(),
-    );
+        let seconds = timestamp.as_secs();
+        frame.ready(
+            (seconds >> 32) as u32,
+            seconds as u32,
+            timestamp.subsec_nanos(),
+        );
+    }
     Ok(())
 }
 
@@ -394,6 +542,19 @@ mod tests {
     use super::*;
     use smithay::utils::{Physical, Transform};
 
+    #[test]
+    fn capture_is_off_unless_the_opt_in_is_exactly_one() {
+        use std::ffi::OsStr;
+
+        assert!(!capture_opted_in(None));
+        assert!(!capture_opted_in(Some(OsStr::new("0"))));
+        assert!(!capture_opted_in(Some(OsStr::new(""))));
+        assert!(!capture_opted_in(Some(OsStr::new("true"))));
+        assert!(!capture_opted_in(Some(OsStr::new("01"))));
+        assert!(!capture_opted_in(Some(OsStr::new(" 1"))));
+        assert!(capture_opted_in(Some(OsStr::new("1"))));
+    }
+
     #[test]
     fn capture_region_clips_then_scales_to_buffer_coordinates() {
         let mode = Size::<i32, Physical>::from((1_920, 1_080));

compositor/src/handlers/screenshot_worker.rs

@@ -0,0 +1,292 @@
+use std::fs;
+use std::io;
+use std::os::unix::fs::{OpenOptionsExt, PermissionsExt};
+use std::path::PathBuf;
+use std::sync::mpsc::{Receiver, SyncSender, TrySendError, sync_channel};
+use std::thread;
+use std::time::{Duration, SystemTime};
+
+use smithay::reexports::calloop::channel;
+
+use super::screenshot::{Canvas, OutputFrame, compose};
+
+const QUEUE_CAPACITY: usize = 4;
+/// A client normally opens and unlinks its screenshot immediately. If it dies
+/// first, the file is left behind, so staged files are swept well after any
+/// realistic request could still be reading one.
+const STAGE_TTL: Duration = Duration::from_secs(60 * 60);
+pub(crate) const SWEEP_INTERVAL: Duration = Duration::from_secs(5 * 60);
+const DIR_MODE: u32 = 0o700;
+const FILE_MODE: u32 = 0o600;
+
+pub(crate) struct Job {
+    pub(crate) request: u64,
+    pub(crate) frames: Vec<OutputFrame>,
+}
+
+pub(crate) struct Encoded {
+    pub(crate) request: u64,
+    pub(crate) result: Result<PathBuf, String>,
+}
+
+pub(crate) struct Worker {
+    queue: SyncSender<Job>,
+}
+
+impl Worker {
+    pub(crate) fn spawn(results: channel::SyncSender<Encoded>) -> Option<Self> {
+        let (queue, receiver) = sync_channel::<Job>(QUEUE_CAPACITY);
+        thread::Builder::new()
+            .name("ferese-screenshot".to_owned())
+            .spawn(move || run(receiver, results))
+            .ok()
+            .map(|_| Self { queue })
+    }
+
+    // The compositor thread must never block, so a full queue is reported back
+    // to the caller instead of waiting.
+    pub(crate) fn submit(&self, job: Job) -> Result<(), Job> {
+        match self.queue.try_send(job) {
+            Ok(()) => Ok(()),
+            Err(TrySendError::Full(job)) | Err(TrySendError::Disconnected(job)) => Err(job),
+        }
+    }
+}
+
+fn run(receiver: Receiver<Job>, results: channel::SyncSender<Encoded>) {
+    while let Ok(job) = receiver.recv() {
+        let result = encode(&job);
+        // The loop drains this channel, so a full queue only delays the worker.
+        if results
+            .send(Encoded {
+                request: job.request,
+                result,
+            })
+            .is_err()
+        {
+            return;
+        }
+    }
+}
+
+// Removes staged screenshots whose client never collected them. Only files
+// older than the TTL are touched, so an in-flight request is never disturbed.
+pub(crate) fn sweep_stale_files() {
+    let Ok(directory) = private_directory() else {
+        return;
+    };
+    let Ok(entries) = fs::read_dir(&directory) else {
+        return;
+    };
+    for entry in entries.flatten() {
+        let path = entry.path();
+        if path.extension().is_none_or(|extension| extension != "png") {
+            continue;
+        }
+        let Ok(metadata) = entry.metadata() else {
+            continue;
+        };
+        if !metadata.is_file() {
+            continue;
+        }
+        // A file dated in the future is treated as fresh rather than swept.
+        let stale = metadata
+            .modified()
+            .ok()
+            .and_then(|modified| SystemTime::now().duration_since(modified).ok())
+            .is_some_and(|age| age > STAGE_TTL);
+        if stale && fs::remove_file(&path).is_ok() {
+            tracing::debug!(path = %path.display(), "removed an uncollected screenshot");
+        }
+    }
+}
+
+fn encode(job: &Job) -> Result<PathBuf, String> {
+    let canvas = compose(&job.frames)?;
+    let png = encode_png(&canvas)?;
+    write_private_file(&png)
+}
+
+fn encode_png(canvas: &Canvas) -> Result<Vec<u8>, String> {
+    let (width, height) = (canvas.width, canvas.height);
+    // A final sanity check on the composed size before handing it to the
+    // encoder, which would otherwise panic on a zero dimension.
+    if width == 0 || height == 0 {
+        return Err("Screenshot is empty".into());
+    }
+    let image = image::RgbaImage::from_raw(width, height, canvas.pixels.clone())
+        .ok_or("Screenshot buffer does not match its dimensions")?;
+    let mut out = Vec::new();
+    let encoder = image::codecs::png::PngEncoder::new_with_quality(
+        std::io::Cursor::new(&mut out),
+        image::codecs::png::CompressionType::Fast,
+        image::codecs::png::FilterType::Adaptive,
+    );
+    image::ImageEncoder::write_image(
+        encoder,
+        image.as_raw(),
+        width,
+        height,
+        image::ExtendedColorType::Rgba8,
+    )
+    .map_err(|error| format!("Could not encode screenshot: {error}"))?;
+    tracing::debug!(encoded = out.len(), "encoded screenshot");
+    Ok(out)
+}
+
+// The client is told to open this path and unlink it immediately, so the
+// directory is ours, private, and files are created with the final mode.
+fn write_private_file(png: &[u8]) -> Result<PathBuf, String> {
+    let directory = private_directory()?;
+    for attempt in 0..64u32 {
+        let path = directory.join(format!("shot-{}-{attempt}.png", std::process::id()));
+        let mut options = fs::OpenOptions::new();
+        options.write(true).create_new(true).mode(FILE_MODE);
+        match options.open(&path) {
+            Ok(mut file) => {
+                use io::Write;
+                file.write_all(png)
+                    .and_then(|()| file.sync_all())
+                    .map_err(|error| format!("Could not write screenshot: {error}"))?;
+                return Ok(path);
+            }
+            Err(error) if error.kind() == io::ErrorKind::AlreadyExists => continue,
+            Err(error) => return Err(format!("Could not write screenshot: {error}")),
+        }
+    }
+    Err("Could not find an unused screenshot filename".into())
+}
+
+fn private_directory() -> Result<PathBuf, String> {
+    let runtime = std::env::var_os("XDG_RUNTIME_DIR")
+        .ok_or("XDG_RUNTIME_DIR is not set, so screenshots cannot be staged")?;
+    let root = PathBuf::from(runtime);
+    let directory = root.join("ferese-screenshots");
+    match fs::symlink_metadata(&directory) {
+        Ok(metadata) if metadata.file_type().is_symlink() => {
+            return Err("Refusing to stage screenshots in a symlinked directory".into());
+        }
+        Ok(metadata) => {
+            use std::os::unix::fs::MetadataExt;
+            if !metadata.is_dir() || metadata.uid() != unsafe { libc::geteuid() } {
+                return Err("Refusing to stage screenshots in an untrusted directory".into());
+            }
+        }
+        Err(error) if error.kind() == io::ErrorKind::NotFound => {
+            fs::create_dir(&directory)
+                .and_then(|()| {
+                    fs::set_permissions(&directory, fs::Permissions::from_mode(DIR_MODE))
+                })
+                .map_err(|error| format!("Could not create screenshot directory: {error}"))?;
+        }
+        Err(error) => return Err(format!("Could not inspect screenshot directory: {error}")),
+    }
+    Ok(directory)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use smithay::utils::Transform;
+
+    fn frame(width: i32, height: i32) -> OutputFrame {
+        let stride = width as usize * 4;
+        let mut pixels = vec![0u8; stride * height as usize];
+        let (quads, _) = pixels.as_chunks_mut::<4>();
+        for pixel in quads {
+            pixel[3] = 255;
+        }
+        OutputFrame {
+            width,
+            height,
+            stride,
+            pixels,
+            transform: Transform::Normal,
+            location: (0, 0),
+            scale: 1.0,
+            logical_width: width,
+            logical_height: height,
+        }
+    }
+
+    #[test]
+    fn encode_produces_a_decodable_png() {
+        let job = Job {
+            request: 1,
+            frames: vec![frame(4, 3)],
+        };
+        let canvas = compose(&job.frames).unwrap();
+        let png = encode_png(&canvas).expect("png");
+        assert_eq!(&png[..8], b"\x89PNG\r\n\x1a\n", "png signature");
+
+        let decoded = image::load_from_memory_with_format(&png, image::ImageFormat::Png)
+            .expect("decodes back")
+            .to_rgba8();
+        assert_eq!(decoded.dimensions(), (4, 3));
+        // The opaque alpha we force is preserved through the round trip.
+        assert_eq!(decoded.get_pixel(0, 0)[3], 255);
+    }
+
+    #[test]
+    fn staged_files_are_owner_only() {
+        let png = b"\x89PNG\r\n\x1a\nnot really a png".to_vec();
+        let path = write_private_file(&png).expect("staged");
+        use std::os::unix::fs::PermissionsExt;
+        let mode = fs::metadata(&path).unwrap().permissions().mode() & 0o777;
+        assert_eq!(
+            mode, FILE_MODE,
+            "staged screenshots are not group/world readable"
+        );
+        fs::remove_file(&path).unwrap();
+    }
+
+    #[test]
+    fn the_sweep_keeps_fresh_files_and_ignores_other_types() {
+        let png = b"\x89PNG\r\n\x1a\n".to_vec();
+        let fresh = write_private_file(&png).expect("staged");
+
+        let directory = private_directory().unwrap();
+        let other = directory.join("keep-me.txt");
+        fs::write(&other, b"not a screenshot").unwrap();
+
+        sweep_stale_files();
+        assert!(
+            fresh.exists(),
+            "a freshly staged screenshot is still in use"
+        );
+        assert!(other.exists(), "unrelated files are left alone");
+
+        fs::remove_file(&fresh).unwrap();
+        fs::remove_file(&other).unwrap();
+    }
+
+    #[test]
+    fn the_sweep_removes_files_past_the_ttl() {
+        let png = b"\x89PNG\r\n\x1a\n".to_vec();
+        let path = write_private_file(&png).expect("staged");
+        // Backdate it past the TTL without waiting an hour.
+        let aged = SystemTime::now() - STAGE_TTL - Duration::from_secs(60);
+        fs::File::options()
+            .write(true)
+            .open(&path)
+            .unwrap()
+            .set_modified(aged)
+            .expect("backdate");
+
+        sweep_stale_files();
+        assert!(!path.exists(), "an uncollected screenshot is reclaimed");
+    }
+
+    #[test]
+    fn a_missing_runtime_dir_is_reported_rather_than_ignored() {
+        // private_directory must never fall back to /tmp or the cwd.
+        let result = private_directory();
+        match result {
+            Ok(path) => assert!(
+                path.ends_with("ferese-screenshots"),
+                "screenshots are staged in the runtime dir, not a fallback location"
+            ),
+            Err(error) => assert!(error.contains("XDG_RUNTIME_DIR") || error.contains("untrusted")),
+        }
+    }
+}