Add isolated window screenshot targets with native selection

compositor/src/backends/direct.rs

@@ -85,6 +85,30 @@ impl DirectBackendState {
             remaining,
         )
     }
+
+    pub(crate) fn capture_window_buffer(
+        &mut self,
+        window: &smithay::desktop::Window,
+        geometry: smithay::utils::Rectangle<i32, smithay::utils::Logical>,
+        output: &Output,
+    ) -> Result<crate::handlers::screenshot::CaptureBuffer, String> {
+        let device = self
+            .devices
+            .values_mut()
+            .find(|device| {
+                device
+                    .outputs
+                    .values()
+                    .any(|candidate| &candidate.output == output)
+            })
+            .ok_or("Window output is unavailable")?;
+        crate::winit::capture_window_buffer(
+            &mut device.renderer,
+            window,
+            geometry,
+            output.current_scale().fractional_scale(),
+        )
+    }
 }
 
 #[derive(Clone, Debug)]

compositor/src/handlers/screenshot.rs

@@ -10,6 +10,7 @@ use smithay::utils::{Buffer, Logical, Physical, Point, Rectangle, Size, Transfor
 const BYTES_PER_PIXEL: usize = 4;
 
 pub(crate) struct OutputFrame {
+    pub(crate) preserve_alpha: bool,
     pub(crate) width: i32,
     pub(crate) height: i32,
     pub(crate) stride: usize,
@@ -114,7 +115,7 @@ fn scale_and_convert(
                 out[to] = pixel[2];
                 out[to + 1] = pixel[1];
                 out[to + 2] = pixel[0];
-                out[to + 3] = 255;
+                out[to + 3] = if frame.preserve_alpha { pixel[3] } else { 255 };
             }
         }
         return Some(out);
@@ -147,7 +148,7 @@ fn scale_and_convert(
             out[to] = pixel[2];
             out[to + 1] = pixel[1];
             out[to + 2] = pixel[0];
-            out[to + 3] = 255;
+            out[to + 3] = if frame.preserve_alpha { pixel[3] } else { 255 };
         }
     }
     Some(out)
@@ -362,6 +363,7 @@ pub(crate) type PartSender = smithay::reexports::calloop::channel::Sender<PartOu
 // Geometry snapshotted when the request was admitted, so the encoding worker
 // never re-queries an output whose transform or scale has since changed.
 pub(crate) struct PartSpec {
+    pub(crate) preserve_alpha: bool,
     pub(crate) transform: Transform,
     pub(crate) scale: f64,
     pub(crate) location: (i32, i32),
@@ -510,6 +512,7 @@ impl Coordinator {
                 } else {
                     pixels.extend_from_slice(&buffer.pixels);
                     Some(OutputFrame {
+                        preserve_alpha: spec.preserve_alpha,
                         width: buffer.width,
                         height: buffer.height,
                         stride: buffer.stride,
@@ -664,6 +667,7 @@ pub(crate) struct PlannedPart {
 impl PlannedPart {
     pub(crate) fn spec(&self, output: &OutputLayout) -> PartSpec {
         PartSpec {
+            preserve_alpha: false,
             transform: output.transform,
             scale: output.scale,
             location: (self.logical.loc.x, self.logical.loc.y),
@@ -984,6 +988,7 @@ mod completion {
 
     fn spec(width: i32, height: i32) -> PartSpec {
         PartSpec {
+            preserve_alpha: false,
             transform: Transform::Normal,
             scale: 1.0,
             location: (0, 0),
@@ -1276,6 +1281,7 @@ mod completion {
         assert!(overflow.is_err(), "outstanding requests are bounded");
 
         let huge = PartSpec {
+            preserve_alpha: false,
             buffer_width: 1 << 20,
             buffer_height: 1 << 20,
             ..spec(2, 2)
@@ -1351,6 +1357,7 @@ mod tests {
             pixels.extend_from_slice(&[b'a' + index as u8, 0, 0, 255]);
         }
         OutputFrame {
+            preserve_alpha: false,
             width,
             height,
             stride: width as usize * BYTES_PER_PIXEL,
@@ -1377,6 +1384,7 @@ mod tests {
             pixels.extend_from_slice(&[(index % 251) as u8, 0, 0, 255]);
         }
         OutputFrame {
+            preserve_alpha: false,
             width,
             height,
             stride: width as usize * BYTES_PER_PIXEL,

compositor/src/handlers/screenshot_worker.rs

@@ -200,6 +200,7 @@ mod tests {
             pixel[3] = 255;
         }
         OutputFrame {
+            preserve_alpha: false,
             width,
             height,
             stride,

compositor/src/ipc.rs

@@ -125,7 +125,9 @@ pub(crate) fn init(
                         return;
                     }
                 };
-                if call.request.command == "screenshot" {
+                if call.request.command == "screenshot-window" {
+                    state.start_window_screenshot(call.request, call.response);
+                } else if call.request.command == "screenshot" {
                     // Deferred: this path answers the caller itself, exactly
                     // once, whenever the request finishes or is terminated.
                     state.start_screenshot(call.request, call.response);
@@ -309,6 +311,8 @@ fn serve_connection(
         };
         let (response, receiver) = sync_channel(1);
         let exit_requested = request.command == "exit";
+        let screenshot_requested =
+            matches!(request.command.as_str(), "screenshot" | "screenshot-window");
         let call = IpcEvent::Call(IpcCall {
             owner,
             request,
@@ -323,6 +327,9 @@ fn serve_connection(
         };
         let exit_accepted = exit_requested && response.error.is_none();
         if let Err(error) = write_frame(&mut stream, &response) {
+            if screenshot_requested {
+                discard_undelivered_screenshot(&response);
+            }
             tracing::debug!(%error, "IPC client disconnected before receiving its response");
             return;
         }
@@ -335,6 +342,19 @@ fn serve_connection(
     }
 }
 
+fn discard_undelivered_screenshot(response: &Response) {
+    if response.error.is_none()
+        && let Some(path) = response
+            .result
+            .as_ref()
+            .and_then(|result| result["path"].as_str())
+        && let Err(error) = fs::remove_file(path)
+        && error.kind() != io::ErrorKind::NotFound
+    {
+        tracing::debug!(%error, "failed to remove an undelivered screenshot");
+    }
+}
+
 impl Ferese {
     fn handle_ipc_request(&mut self, owner: u64, request: Request) -> Response {
         if self.session_lock.active && request.command != "portal-shortcuts-poll" {
@@ -520,6 +540,93 @@ impl Ferese {
         crate::backends::direct::render_all(self);
     }
 
+    fn start_window_screenshot(&mut self, request: Request, response: SyncSender<Response>) {
+        let reject = |message: String| {
+            let _ = response.try_send(Response::error(
+                request.id,
+                "window_capture_failed",
+                message,
+            ));
+        };
+        if self.session_lock.active || !screencopy::capture_allowed() {
+            reject("Screen capture is unavailable".into());
+            return;
+        }
+        if let Err(error) = validate_request(&request) {
+            reject(error.message);
+            return;
+        }
+        let Some(id) = request
+            .args
+            .get("window")
+            .and_then(Value::as_u64)
+            .map(ferese_layout::WindowId)
+        else {
+            reject("Expected a window ID".into());
+            return;
+        };
+        let Some(window) = self
+            .window_ids
+            .iter()
+            .find(|(_, candidate)| **candidate == id)
+            .map(|(window, _)| window.clone())
+        else {
+            reject("Window no longer exists".into());
+            return;
+        };
+        let Some(output) = self
+            .space
+            .outputs()
+            .find(|output| self.window_belongs_to_output(id, output))
+            .cloned()
+        else {
+            reject("Window output is unavailable".into());
+            return;
+        };
+        let geometry = window.geometry();
+        let scale = output.current_scale().fractional_scale();
+        let size = geometry.size.to_physical_precise_round(scale);
+        let spec = crate::handlers::screenshot::PartSpec {
+            preserve_alpha: true,
+            transform: Transform::Normal,
+            scale,
+            location: (0, 0),
+            logical_width: geometry.size.w,
+            logical_height: geometry.size.h,
+            buffer_width: size.w,
+            buffer_height: size.h,
+        };
+        let fallback = response.clone();
+        let capture = match self.screenshot.admit(request.id, response, vec![spec]) {
+            Ok(capture) => capture,
+            Err(error) => {
+                let _ =
+                    fallback.try_send(Response::error(request.id, "screenshot_rejected", error));
+                return;
+            }
+        };
+        let result = if let Some(backend) = &self.nested_backend {
+            match backend.try_borrow_mut() {
+                Ok(mut backend) => crate::winit::capture_window_buffer(
+                    backend.renderer(),
+                    &window,
+                    geometry,
+                    scale,
+                ),
+                Err(_) => Err("Window renderer is busy".into()),
+            }
+        } else if let Some(backend) = &mut self.direct_backend {
+            backend.capture_window_buffer(&window, geometry, &output)
+        } else {
+            Err("Window renderer is unavailable".into())
+        };
+        self.on_screenshot_part(PartOutcome {
+            request: capture,
+            part: 0,
+            result,
+        });
+    }
+
     pub(crate) fn on_screenshot_part(&mut self, outcome: PartOutcome) {
         let action = self
             .screenshot
@@ -1030,6 +1137,20 @@ mod tests {
         worker.join().unwrap();
     }
 
+    #[test]
+    fn failed_screenshot_delivery_removes_only_its_staging_file() {
+        let root = unique_test_directory("undelivered-screenshot");
+        fs::create_dir(&root).unwrap();
+        let path = root.join("capture.png");
+        fs::write(&path, b"capture").unwrap();
+        let response = Response::success(1, serde_json::json!({"path": path}));
+
+        discard_undelivered_screenshot(&response);
+        assert!(!path.exists());
+        discard_undelivered_screenshot(&response);
+        fs::remove_dir(root).unwrap();
+    }
+
     #[test]
     fn runtime_socket_directory_is_private() {
         let root = unique_test_directory("private-parent");

compositor/src/winit.rs

@@ -14,7 +14,8 @@ use smithay::{
     backend::{
         allocator::Fourcc,
         renderer::{
-            Bind, Color32F, ErasedContextId, Frame, ImportDma, Offscreen, Renderer, Texture,
+            Bind, Color32F, ErasedContextId, ExportMem, Frame, ImportDma, Offscreen, Renderer,
+            Texture,
             damage::OutputDamageTracker,
             element::{
                 Element, Id, Kind as RenderElementKind, RenderElement, UnderlyingStorage,
@@ -169,6 +170,56 @@ pub(crate) fn capture_resize_snapshot(
     }))
 }
 
+pub(crate) fn capture_window_buffer(
+    renderer: &mut GlesRenderer,
+    window: &smithay::desktop::Window,
+    geometry: Rectangle<i32, Logical>,
+    scale: f64,
+) -> Result<crate::handlers::screenshot::CaptureBuffer, String> {
+    let snapshot = capture_resize_snapshot(renderer, window, geometry, scale, 128 * 1024 * 1024)
+        .map_err(|error| error.to_string())?
+        .ok_or("Window has no capturable content")?;
+    let size = snapshot.texture.size();
+    let mapping = renderer
+        .copy_texture(
+            &snapshot.texture,
+            Rectangle::from_size(size),
+            Fourcc::Abgr8888,
+        )
+        .map_err(|error| error.to_string())?;
+    if mapping.format() != Some(Fourcc::Abgr8888) {
+        return Err("Unsupported window readback format".into());
+    }
+    let bytes = (size.w as usize)
+        .checked_mul(size.h as usize)
+        .and_then(|pixels| pixels.checked_mul(4))
+        .ok_or("Window is too large")?;
+    let source = renderer
+        .map_texture(&mapping)
+        .map_err(|error| error.to_string())?;
+    if source.len() < bytes {
+        return Err("Incomplete window readback".into());
+    }
+    let mut pixels = source[..bytes].to_vec();
+    for pixel in pixels.chunks_exact_mut(4) {
+        pixel.swap(0, 2);
+        let alpha = pixel[3] as u16;
+        for channel in &mut pixel[..3] {
+            *channel = if alpha == 0 {
+                0
+            } else {
+                ((*channel as u16 * 255 + alpha / 2) / alpha).min(255) as u8
+            };
+        }
+    }
+    Ok(crate::handlers::screenshot::CaptureBuffer {
+        width: size.w,
+        height: size.h,
+        stride: size.w as usize * 4,
+        pixels,
+    })
+}
+
 render_elements! {
     WindowContentRenderElement<=GlesRenderer>;
     Rounded=RoundedSurfaceRenderElement,