Implement ext-image-capture-source and ext-image-copy-capture protocols (staging)

src/handlers/mod.rs

@@ -13,19 +13,22 @@ use std::time::Duration;
 use smithay::backend::allocator::dmabuf::Dmabuf;
 use smithay::backend::drm::DrmNode;
 use smithay::backend::input::{InputEvent, TabletToolDescriptor};
+use smithay::backend::renderer::damage::OutputDamageTracker;
+use smithay::backend::renderer::gles::GlesRenderer;
+use smithay::backend::renderer::{BufferType, buffer_type};
 use smithay::desktop::{PopupKind, PopupManager};
 use smithay::input::dnd::{self, DnDGrab, DndGrabHandler, DndTarget};
 use smithay::input::pointer::{CursorIcon, CursorImageStatus, Focus, PointerHandle};
 use smithay::input::{keyboard, Seat, SeatHandler, SeatState};
-use smithay::output::Output;
+use smithay::output::{Output, WeakOutput};
 use smithay::reexports::rustix::fs::{fcntl_setfl, OFlags};
 use smithay::reexports::wayland_protocols_wlr::screencopy::v1::server::zwlr_screencopy_manager_v1::ZwlrScreencopyManagerV1;
 use smithay::reexports::wayland_server::protocol::wl_output::WlOutput;
 use smithay::reexports::wayland_server::protocol::wl_surface::WlSurface;
 use smithay::reexports::wayland_server::Resource;
-use smithay::utils::{Logical, Point, Rectangle, Serial};
+use smithay::utils::{Logical, Point, Rectangle, Scale, Serial, Transform};
 use smithay::wayland::compositor::{get_parent, with_states};
-use smithay::wayland::dmabuf::{DmabufGlobal, DmabufHandler, DmabufState, ImportNotifier};
+use smithay::wayland::dmabuf::{self as wl_dmabuf, DmabufGlobal, DmabufHandler, DmabufState, ImportNotifier};
 use smithay::wayland::drm_lease::{
     DrmLease, DrmLeaseBuilder, DrmLeaseHandler, DrmLeaseRequest, DrmLeaseState, LeaseRejected,
 };
@@ -71,15 +74,28 @@ use smithay::{
     delegate_single_pixel_buffer, delegate_tablet_manager, delegate_text_input_manager,
     delegate_viewporter, delegate_virtual_keyboard_manager, delegate_xdg_activation,
 };
+use smithay::wayland::image_capture_source::{
+    ImageCaptureSource, ImageCaptureSourceHandler, OutputCaptureSourceHandler,
+    OutputCaptureSourceState,
+};
+use smithay::wayland::image_copy_capture::{
+    BufferConstraints, ImageCopyCaptureHandler, ImageCopyCaptureState, Session, SessionRef, Frame,
+};
 
 pub use crate::handlers::xdg_shell::KdeDecorationsModeState;
 use crate::layout::workspace::WorkspaceId;
 use crate::layout::ActivateWindow;
-use crate::niri::{DndIcon, NewClient, State};
+use crate::niri::{DndIcon, NewClient, Niri, OutputRenderElements, State};
+use crate::window::mapped::WindowCastRenderElements;
 use crate::protocols::ext_workspace::{self, ExtWorkspaceHandler, ExtWorkspaceManagerState};
+use crate::render_helpers::{render_to_dmabuf, render_to_shm, RenderCtx, RenderTarget};
+use crate::utils::get_monotonic_time;
 use crate::protocols::foreign_toplevel::{
     self, ForeignToplevelHandler, ForeignToplevelManagerState,
 };
+use crate::protocols::toplevel_image_capture_source::{
+    ToplevelImageCaptureHandler, ToplevelImageCaptureManagerState,
+};
 use crate::protocols::gamma_control::{GammaControlHandler, GammaControlManagerState};
 use crate::protocols::mutter_x11_interop::MutterX11InteropHandler;
 use crate::protocols::output_management::{OutputManagementHandler, OutputManagementManagerState};
@@ -93,7 +109,7 @@ use crate::utils::{output_size, send_scale_transform};
 use crate::{
     delegate_ext_workspace, delegate_foreign_toplevel, delegate_gamma_control,
     delegate_mutter_x11_interop, delegate_output_management, delegate_screencopy,
-    delegate_virtual_pointer,
+    delegate_toplevel_image_capture_source, delegate_virtual_pointer,
 };
 
 pub const XDG_ACTIVATION_TOKEN_TIMEOUT: Duration = Duration::from_secs(10);
@@ -597,6 +613,256 @@ impl ForeignToplevelHandler for State {
 }
 delegate_foreign_toplevel!(State);
 
+impl ImageCaptureSourceHandler for State {
+    fn source_destroyed(&mut self, _source: ImageCaptureSource) {
+        // niri doesn't need to track sources
+    }
+}
+
+smithay::delegate_image_capture_source!(State);
+
+impl OutputCaptureSourceHandler for State {
+    fn output_capture_source_state(&mut self) -> &mut OutputCaptureSourceState {
+        &mut self.niri.output_capture_source_state
+    }
+
+    fn output_source_created(&mut self, source: ImageCaptureSource, output: &smithay::output::Output) {
+        source.user_data().insert_if_missing(|| output.downgrade());
+    }
+}
+
+smithay::delegate_output_capture_source!(State);
+
+impl ImageCopyCaptureHandler for State {
+    fn image_copy_capture_state(&mut self) -> &mut ImageCopyCaptureState {
+        &mut self.niri.image_copy_capture_state
+    }
+
+    fn capture_constraints(&mut self, source: &ImageCaptureSource) -> Option<BufferConstraints> {
+        // Check for output sources
+        if let Some(weak_output) = source.user_data().get::<WeakOutput>() {
+            let output = weak_output.upgrade()?;
+            let mode = output.current_mode()?;
+            let transform = output.current_transform();
+            let size = transform.transform_size(mode.size);
+
+            return Some(BufferConstraints {
+                size: size.to_logical(1).to_buffer(1, smithay::utils::Transform::Normal),
+                shm: vec![
+                    smithay::reexports::wayland_server::protocol::wl_shm::Format::Argb8888,
+                    smithay::reexports::wayland_server::protocol::wl_shm::Format::Xrgb8888,
+                ],
+                dma: None,
+            });
+        }
+
+        // Check for toplevel sources
+        if let Some(wl_surface) = source.user_data().get::<WlSurface>() {
+            let root = self.niri.find_root_shell_surface(wl_surface);
+            let (mapped, _) = self.niri.layout.find_window_and_output(&root)?;
+            let window = &mapped.window;
+            let scale = self.niri
+                .layout
+                .find_window_and_output(&root)
+                .and_then(|(_, output)| output)
+                .map(|o| Scale::from(o.current_scale().fractional_scale()))
+                .unwrap_or(Scale::from(1.0));
+            let bbox = window.bbox_with_popups().to_physical_precise_up(scale);
+
+            return Some(BufferConstraints {
+                size: bbox.size.to_logical(1).to_buffer(1, smithay::utils::Transform::Normal),
+                shm: vec![
+                    smithay::reexports::wayland_server::protocol::wl_shm::Format::Argb8888,
+                    smithay::reexports::wayland_server::protocol::wl_shm::Format::Xrgb8888,
+                ],
+                dma: None,
+            });
+        }
+
+        None
+    }
+
+    fn new_session(&mut self, _session: Session) {
+        // Sessions clean up when dropped
+    }
+
+    fn frame(&mut self, session: &SessionRef, frame: Frame) {
+        let source = session.source();
+        let buffer = frame.buffer();
+
+        let rendered = self.backend.with_primary_renderer(|renderer| -> bool {
+            if let Some(weak_output) = source.user_data().get::<WeakOutput>() {
+                let output = match weak_output.upgrade() {
+                    Some(o) => o,
+                    None => return false,
+                };
+                let niri = &mut self.niri;
+                render_output_capture(niri, renderer, &output, &buffer)
+            } else if let Some(wl_surface) = source.user_data().get::<WlSurface>() {
+                let niri = &mut self.niri;
+                render_toplevel_capture(niri, renderer, wl_surface, &buffer)
+            } else {
+                false
+            }
+        }).unwrap_or(false);
+
+        if rendered {
+            frame.success(Transform::Normal, None, get_monotonic_time());
+        } else {
+            frame.fail(smithay::wayland::image_copy_capture::CaptureFailureReason::Unknown);
+        }
+    }
+}
+
+smithay::delegate_image_copy_capture!(State);
+
+/// Render output content into a buffer for the image copy capture protocol.
+fn render_output_capture(
+    niri: &mut Niri,
+    renderer: &mut GlesRenderer,
+    output: &Output,
+    buffer: &smithay::reexports::wayland_server::protocol::wl_buffer::WlBuffer,
+) -> bool {
+    // Update render elements for this output
+    niri.update_render_elements(Some(output));
+
+    // Build render elements
+    let mut elements: Vec<OutputRenderElements<GlesRenderer>> = Vec::new();
+    let ctx = RenderCtx {
+        renderer,
+        target: RenderTarget::ScreenCapture,
+        xray: None,
+    };
+    niri.render(ctx, output, false, &mut |elem| {
+        elements.push(elem);
+    });
+
+    // Create damage tracker and compute states (age=0 = full damage)
+    let mut damage_tracker = OutputDamageTracker::from_output(output);
+    let Ok((_damages, states)) = damage_tracker.damage_output(0, &elements) else {
+        warn!("failed to compute output damage for image capture");
+        return false;
+    };
+
+    // Render based on buffer type
+    let result = match buffer_type(buffer) {
+        Some(BufferType::Shm) => {
+            render_to_shm(renderer, &mut damage_tracker, buffer, &elements, states)
+                .map(|_| ())
+        }
+        Some(BufferType::Dma) => {
+            match wl_dmabuf::get_dmabuf(buffer) {
+                Ok(dmabuf) => render_to_dmabuf(renderer, &mut damage_tracker, dmabuf.clone(), &elements, states)
+                    .map(|_| ()),
+                Err(err) => {
+                    warn!("failed to get dmabuf from buffer: {err:?}");
+                    Err(anyhow::anyhow!("dmabuf error"))
+                }
+            }
+        }
+        _ => {
+            warn!("unsupported buffer type for image capture");
+            return false;
+        }
+    };
+
+    match result {
+        Ok(()) => true,
+        Err(err) => {
+            warn!("error rendering output for image capture: {err:?}");
+            false
+        }
+    }
+}
+
+/// Render toplevel (window) content into a buffer for the image copy capture protocol.
+fn render_toplevel_capture(
+    niri: &mut Niri,
+    renderer: &mut GlesRenderer,
+    wl_surface: &WlSurface,
+    buffer: &smithay::reexports::wayland_server::protocol::wl_buffer::WlBuffer,
+) -> bool {
+    // Find the root surface and mapped window
+    let root = niri.find_root_shell_surface(wl_surface);
+    let Some((mapped, output_opt)) = niri.layout.find_window_and_output(&root) else {
+        warn!("toplevel image capture: window not found");
+        return false;
+    };
+
+    let window = &mapped.window;
+    let output = match output_opt {
+        Some(o) => o,
+        None => {
+            warn!("toplevel image capture: output not found");
+            return false;
+        }
+    };
+
+    let scale = Scale::from(output.current_scale().fractional_scale());
+    let bbox = window.bbox_with_popups().to_physical_precise_up(scale);
+    let size = bbox.size;
+
+    // Build window render elements
+    let mut elements: Vec<WindowCastRenderElements<GlesRenderer>> = Vec::new();
+    mapped.render_for_screen_cast(renderer, scale, &mut |elem| {
+        elements.push(elem);
+    });
+
+    // Create a damage tracker for the window size
+    let mut damage_tracker = OutputDamageTracker::new(size, scale, Transform::Normal);
+    let Ok((_damages, states)) = damage_tracker.damage_output(0, &elements) else {
+        warn!("failed to compute window damage for image capture");
+        return false;
+    };
+
+    // Render based on buffer type
+    let result = match buffer_type(buffer) {
+        Some(BufferType::Shm) => {
+            render_to_shm(renderer, &mut damage_tracker, buffer, &elements, states)
+                .map(|_| ())
+        }
+        Some(BufferType::Dma) => {
+            match wl_dmabuf::get_dmabuf(buffer) {
+                Ok(dmabuf) => render_to_dmabuf(renderer, &mut damage_tracker, dmabuf.clone(), &elements, states)
+                    .map(|_| ()),
+                Err(err) => {
+                    warn!("failed to get dmabuf from buffer: {err:?}");
+                    Err(anyhow::anyhow!("dmabuf error"))
+                }
+            }
+        }
+        _ => {
+            warn!("unsupported buffer type for image capture");
+            return false;
+        }
+    };
+
+    match result {
+        Ok(()) => true,
+        Err(err) => {
+            warn!("error rendering toplevel for image capture: {err:?}");
+            false
+        }
+    }
+}
+
+impl ToplevelImageCaptureHandler for State {
+    fn toplevel_image_capture_manager_state(&mut self) -> &mut ToplevelImageCaptureManagerState {
+        &mut self.niri.toplevel_image_capture_state
+    }
+
+    fn lookup_toplevel_surface(
+        &mut self,
+        handle: &smithay::reexports::wayland_protocols::ext::foreign_toplevel_list::v1::server::ext_foreign_toplevel_handle_v1::ExtForeignToplevelHandleV1,
+    ) -> Option<smithay::reexports::wayland_server::protocol::wl_surface::WlSurface> {
+        self.niri
+            .foreign_toplevel_state
+            .find_surface_for_handle(handle)
+            .cloned()
+    }
+}
+delegate_toplevel_image_capture_source!(State);
+
 impl ExtWorkspaceHandler for State {
     fn ext_workspace_manager_state(&mut self) -> &mut ExtWorkspaceManagerState {
         &mut self.niri.ext_workspace_state

src/niri.rs

@@ -81,6 +81,9 @@ use smithay::wayland::dmabuf::DmabufState;
 use smithay::wayland::fractional_scale::FractionalScaleManagerState;
 use smithay::wayland::idle_inhibit::IdleInhibitManagerState;
 use smithay::wayland::idle_notify::IdleNotifierState;
+use smithay::wayland::image_capture_source::ImageCaptureSourceState;
+use smithay::wayland::image_capture_source::OutputCaptureSourceState;
+use smithay::wayland::image_copy_capture::ImageCopyCaptureState;
 use smithay::wayland::input_method::InputMethodManagerState;
 use smithay::wayland::keyboard_shortcuts_inhibit::{
     KeyboardShortcutsInhibitState, KeyboardShortcutsInhibitor,
@@ -150,6 +153,7 @@ use crate::protocols::gamma_control::GammaControlManagerState;
 use crate::protocols::mutter_x11_interop::MutterX11InteropManagerState;
 use crate::protocols::output_management::OutputManagementManagerState;
 use crate::protocols::screencopy::{Screencopy, ScreencopyBuffer, ScreencopyManagerState};
+use crate::protocols::toplevel_image_capture_source::ToplevelImageCaptureManagerState;
 use crate::protocols::virtual_pointer::VirtualPointerManagerState;
 use crate::render_helpers::blur::BlurOptions;
 use crate::render_helpers::debug::push_opaque_regions;
@@ -278,6 +282,10 @@ pub struct Niri {
     pub layer_shell_state: WlrLayerShellState,
     pub session_lock_state: SessionLockManagerState,
     pub foreign_toplevel_state: ForeignToplevelManagerState,
+    pub image_capture_source_state: ImageCaptureSourceState,
+    pub output_capture_source_state: OutputCaptureSourceState,
+    pub image_copy_capture_state: ImageCopyCaptureState,
+    pub toplevel_image_capture_state: ToplevelImageCaptureManagerState,
     pub ext_workspace_state: ExtWorkspaceManagerState,
     pub screencopy_state: ScreencopyManagerState,
     pub output_management_state: OutputManagementManagerState,
@@ -2334,6 +2342,13 @@ impl Niri {
             VirtualPointerManagerState::new::<State, _>(&display_handle, client_is_unrestricted);
         let foreign_toplevel_state =
             ForeignToplevelManagerState::new::<State, _>(&display_handle, client_is_unrestricted);
+        let image_capture_source_state = ImageCaptureSourceState::new();
+        let output_capture_source_state =
+            OutputCaptureSourceState::new::<State>(&display_handle);
+        let image_copy_capture_state =
+            ImageCopyCaptureState::new::<State>(&display_handle);
+        let toplevel_image_capture_state =
+            ToplevelImageCaptureManagerState::new::<State, _>(&display_handle, client_is_unrestricted);
         let ext_workspace_state =
             ExtWorkspaceManagerState::new::<State, _>(&display_handle, client_is_unrestricted);
         let mut output_management_state =
@@ -2522,6 +2537,10 @@ impl Niri {
             layer_shell_state,
             session_lock_state,
             foreign_toplevel_state,
+            image_capture_source_state,
+            output_capture_source_state,
+            image_copy_capture_state,
+            toplevel_image_capture_state,
             ext_workspace_state,
             output_management_state,
             screencopy_state,

src/protocols/foreign_toplevel.rs

@@ -63,6 +63,17 @@ pub struct ForeignToplevelGlobalData {
 }
 
 impl ForeignToplevelManagerState {
+    /// Find the surface (WlSurface) associated with the given ext_foreign_toplevel_handle_v1.
+    pub fn find_surface_for_handle(
+        &self,
+        handle: &smithay::reexports::wayland_protocols::ext::foreign_toplevel_list::v1::server::ext_foreign_toplevel_handle_v1::ExtForeignToplevelHandleV1,
+    ) -> Option<&WlSurface> {
+        self.toplevels
+            .iter()
+            .find(|(_, data)| data.ext_list_instances.contains(handle))
+            .map(|(surface, _)| surface)
+    }
+
     pub fn new<D, F>(display: &DisplayHandle, filter: F) -> Self
     where
         D: GlobalDispatch<ZwlrForeignToplevelManagerV1, ForeignToplevelGlobalData>,