use drm::node::DrmNode;
use std::{
collections::HashSet,
os::fd::{AsFd, BorrowedFd},
path::Path,
sync::atomic::{AtomicBool, Ordering},
};
use wayland_client::{
Connection, Dispatch, QueueHandle,
WEnum::{self, Value},
delegate_noop,
globals::GlobalListContents,
protocol::{
wl_buffer::WlBuffer,
wl_compositor::WlCompositor,
wl_output::{self, WlOutput},
wl_registry::{self, WlRegistry},
wl_shm::WlShm,
wl_shm_pool::WlShmPool,
wl_surface::WlSurface,
},
};
use wayland_protocols::{
ext::{
foreign_toplevel_list::v1::client::{
ext_foreign_toplevel_handle_v1::ExtForeignToplevelHandleV1,
ext_foreign_toplevel_list_v1::ExtForeignToplevelListV1,
},
image_capture_source::v1::client::{
ext_foreign_toplevel_image_capture_source_manager_v1::ExtForeignToplevelImageCaptureSourceManagerV1,
ext_image_capture_source_v1::ExtImageCaptureSourceV1,
ext_output_image_capture_source_manager_v1::ExtOutputImageCaptureSourceManagerV1,
},
image_copy_capture::v1::client::{
ext_image_copy_capture_frame_v1::{self, ExtImageCopyCaptureFrameV1, FailureReason},
ext_image_copy_capture_manager_v1::ExtImageCopyCaptureManagerV1,
ext_image_copy_capture_session_v1::{self, ExtImageCopyCaptureSessionV1},
},
},
wp::{
linux_dmabuf::zv1::client::{
zwp_linux_buffer_params_v1::{self, ZwpLinuxBufferParamsV1},
zwp_linux_dmabuf_v1::{self, ZwpLinuxDmabufV1},
},
viewporter::client::{wp_viewport::WpViewport, wp_viewporter::WpViewporter},
},
xdg::xdg_output::zv1::client::{
zxdg_output_manager_v1::ZxdgOutputManagerV1,
zxdg_output_v1::{self, ZxdgOutputV1},
},
};
use wayland_protocols_wlr::{
layer_shell::v1::client::{
zwlr_layer_shell_v1::ZwlrLayerShellV1,
zwlr_layer_surface_v1::{self, ZwlrLayerSurfaceV1},
},
screencopy::v1::client::{
zwlr_screencopy_frame_v1::{self, ZwlrScreencopyFrameV1},
zwlr_screencopy_manager_v1::ZwlrScreencopyManagerV1,
},
};
use crate::{
output::OutputInfo,
region::{LogicalRegion, Position, Size, TopLevel},
screencopy::{DMAFrameFormat, FrameFormat},
};
#[derive(Debug)]
pub struct OutputCaptureState {
pub outputs: Vec<OutputInfo>,
}
impl Dispatch<WlRegistry, ()> for OutputCaptureState {
#[tracing::instrument(skip(wl_registry, qh), ret, level = "trace")]
fn event(
state: &mut Self,
wl_registry: &WlRegistry,
event: wl_registry::Event,
_: &(),
_: &Connection,
qh: &QueueHandle<Self>,
) {
if let wl_registry::Event::Global {
name,
interface,
version,
} = event
&& interface == "wl_output"
{
if version >= 4 {
let output = wl_registry.bind::<wl_output::WlOutput, _, _>(name, 4, qh, ());
state.outputs.push(OutputInfo {
wl_output: output,
name: "".to_string(),
description: String::new(),
transform: wl_output::Transform::Normal,
physical_size: Size::default(),
logical_region: LogicalRegion::default(),
});
} else {
tracing::error!("Ignoring a wl_output with version < 4.");
}
}
}
}
impl Dispatch<WlOutput, ()> for OutputCaptureState {
#[tracing::instrument(skip(wl_output), ret, level = "trace")]
fn event(
state: &mut Self,
wl_output: &WlOutput,
event: wl_output::Event,
_: &(),
_: &Connection,
_: &QueueHandle<Self>,
) {
let output: &mut OutputInfo =
match state.outputs.iter_mut().find(|x| x.wl_output == *wl_output) {
Some(output) => output,
_ => {
tracing::error!(
"Received event for an output that is not registered: {event:#?}"
);
return;
}
};
match event {
wl_output::Event::Name { name } => {
output.name = name;
}
wl_output::Event::Description { description } => {
output.description = description;
}
wl_output::Event::Mode { width, height, .. } => {
output.physical_size = Size {
width: width as u32,
height: height as u32,
};
}
wl_output::Event::Geometry {
transform: WEnum::Value(transform),
..
} => {
output.transform = transform;
}
wl_output::Event::Scale { .. } => {}
wl_output::Event::Done => {}
_ => {}
}
}
}
delegate_noop!(OutputCaptureState: ignore ZxdgOutputManagerV1);
impl Dispatch<ZxdgOutputV1, usize> for OutputCaptureState {
#[tracing::instrument(ret, level = "trace")]
fn event(
state: &mut Self,
_: &ZxdgOutputV1,
event: zxdg_output_v1::Event,
index: &usize,
_: &Connection,
_: &QueueHandle<Self>,
) {
let output_info = match state.outputs.get_mut(*index) {
Some(output_info) => output_info,
_ => {
tracing::error!(
"Received event for output index {index} that is not registered: {event:#?}"
);
return;
}
};
match event {
zxdg_output_v1::Event::LogicalPosition { x, y } => {
output_info.logical_region.inner.position = Position { x, y };
}
zxdg_output_v1::Event::LogicalSize { width, height } => {
output_info.logical_region.inner.size = Size {
width: width as u32,
height: height as u32,
};
}
zxdg_output_v1::Event::Done => {}
zxdg_output_v1::Event::Name { .. } => {}
zxdg_output_v1::Event::Description { .. } => {}
_ => {}
};
}
}
#[derive(Debug, Copy, Clone, PartialEq)]
pub enum FrameState {
Failed,
FailedWithReason(WEnum<FailureReason>),
Finished,
}
pub struct CaptureFrameState {
pub formats: Vec<FrameFormat>,
pub dmabuf_formats: Vec<DMAFrameFormat>,
pub state: Option<FrameState>,
pub buffer_done: AtomicBool,
pub toplevels: Vec<TopLevel>,
pub(crate) session_done: bool,
pub(crate) gbm: Option<gbm::Device<Card>>,
find_gbm: bool,
session_size: Size,
}
impl CaptureFrameState {
pub(crate) fn new(find_gbm: bool) -> Self {
Self {
formats: Vec::new(),
dmabuf_formats: Vec::new(),
state: None,
buffer_done: AtomicBool::new(false),
toplevels: Vec::new(),
session_done: false,
gbm: None,
find_gbm,
session_size: Size {
width: 0,
height: 0,
},
}
}
}
impl Dispatch<ZwpLinuxDmabufV1, ()> for CaptureFrameState {
fn event(
_frame: &mut Self,
_proxy: &ZwpLinuxDmabufV1,
_event: zwp_linux_dmabuf_v1::Event,
_data: &(),
_conn: &Connection,
_qhandle: &wayland_client::QueueHandle<Self>,
) {
}
}
impl Dispatch<ZwpLinuxBufferParamsV1, ()> for CaptureFrameState {
fn event(
_state: &mut Self,
_proxy: &ZwpLinuxBufferParamsV1,
_event: zwp_linux_buffer_params_v1::Event,
_data: &(),
_conn: &Connection,
_qhandle: &QueueHandle<Self>,
) {
}
}
impl Dispatch<ExtImageCopyCaptureFrameV1, ()> for CaptureFrameState {
fn event(
state: &mut Self,
_proxy: &ExtImageCopyCaptureFrameV1,
event: <ExtImageCopyCaptureFrameV1 as wayland_client::Proxy>::Event,
_data: &(),
_conn: &Connection,
_qhandle: &QueueHandle<Self>,
) {
match event {
ext_image_copy_capture_frame_v1::Event::Ready => {
state.buffer_done.store(true, Ordering::Relaxed);
state.state = Some(FrameState::Finished);
}
ext_image_copy_capture_frame_v1::Event::Failed { reason } => {
state.buffer_done.store(true, Ordering::Relaxed);
state.state = Some(FrameState::FailedWithReason(reason));
}
ext_image_copy_capture_frame_v1::Event::Transform { .. } => {}
_ => {}
}
}
}
impl Dispatch<ExtImageCopyCaptureSessionV1, ()> for CaptureFrameState {
fn event(
state: &mut Self,
_proxy: &ExtImageCopyCaptureSessionV1,
event: <ExtImageCopyCaptureSessionV1 as wayland_client::Proxy>::Event,
_data: &(),
_conn: &Connection,
_qhandle: &QueueHandle<Self>,
) {
match event {
ext_image_copy_capture_session_v1::Event::BufferSize { width, height } => {
state.session_size.width = width;
state.session_size.height = height;
for format in &mut state.formats {
format.size = Size { width, height };
format.stride = 4 * width;
}
for DMAFrameFormat {
size:
Size {
width: dma_width,
height: dma_height,
},
..
} in &mut state.dmabuf_formats
{
*dma_width = width;
*dma_height = height;
}
}
ext_image_copy_capture_session_v1::Event::ShmFormat {
format: WEnum::Value(format),
} => {
state.formats.push(FrameFormat {
format,
size: state.session_size,
stride: 4 * state.session_size.width,
});
let set_format = &mut state.formats[0];
set_format.format = format;
}
ext_image_copy_capture_session_v1::Event::DmabufDevice { device } => {
if !state.find_gbm {
return;
}
let device = match device.try_into() {
Ok(bytes) => u64::from_le_bytes(bytes),
Err(_) => {
tracing::warn!(
"Received invalid device data from compositor (expected 8 bytes)"
);
return;
}
};
let Ok(node) = DrmNode::from_dev_id(device) else {
tracing::warn!("Failed to create DRM node from device ID: {}", device);
return;
};
let Some(pa) = node.dev_path() else {
tracing::warn!("DRM node has no device path");
return;
};
let Ok(card) = Card::open(&pa) else {
tracing::warn!("Failed to open DRM card at {:?}", pa);
return;
};
let Ok(gbm) = gbm::Device::new(card) else {
tracing::warn!("Failed to create GBM device");
return;
};
state.gbm = Some(gbm);
}
ext_image_copy_capture_session_v1::Event::DmabufFormat { format, .. } => {
state.dmabuf_formats.push(DMAFrameFormat {
format,
size: state.session_size,
});
}
ext_image_copy_capture_session_v1::Event::Done => {
state.session_done = true;
}
ext_image_copy_capture_session_v1::Event::Stopped => {
state.session_done = true;
state.state = Some(FrameState::Failed);
}
_ => {}
}
}
}
impl Dispatch<ExtForeignToplevelListV1, ()> for CaptureFrameState {
fn event(
state: &mut Self,
_proxy: &ExtForeignToplevelListV1,
event: <ExtForeignToplevelListV1 as wayland_client::Proxy>::Event,
_data: &(),
_conn: &Connection,
_qhandle: &QueueHandle<Self>,
) {
use wayland_protocols::ext::foreign_toplevel_list::v1::client::ext_foreign_toplevel_list_v1::Event;
if let Event::Toplevel { toplevel } = event {
state.toplevels.push(TopLevel::new(toplevel));
}
}
wayland_client::event_created_child!(CaptureFrameState, ExtForeignToplevelHandleV1, [
wayland_protocols::ext::foreign_toplevel_list::v1::client::ext_foreign_toplevel_list_v1::EVT_TOPLEVEL_OPCODE => (ExtForeignToplevelHandleV1, ())
]);
}
impl Dispatch<ExtForeignToplevelHandleV1, ()> for CaptureFrameState {
fn event(
state: &mut Self,
toplevel: &ExtForeignToplevelHandleV1,
event: <ExtForeignToplevelHandleV1 as wayland_client::Proxy>::Event,
_data: &(),
_conn: &Connection,
_qhandle: &QueueHandle<Self>,
) {
use wayland_protocols::ext::foreign_toplevel_list::v1::client::ext_foreign_toplevel_handle_v1::Event;
if let Some(toplevel_obj) = state.toplevels.iter_mut().find(|t| t.handle == *toplevel) {
match event {
Event::Title { title } => toplevel_obj.title = title,
Event::AppId { app_id } => toplevel_obj.app_id = app_id,
Event::Identifier { identifier } => toplevel_obj.identifier = identifier,
Event::Closed => toplevel_obj.active = false,
_ => {}
}
}
}
}
impl Dispatch<ZwlrScreencopyFrameV1, ()> for CaptureFrameState {
#[tracing::instrument(skip(frame), ret, level = "trace")]
fn event(
frame: &mut Self,
_: &ZwlrScreencopyFrameV1,
event: zwlr_screencopy_frame_v1::Event,
_: &(),
_: &Connection,
_: &QueueHandle<Self>,
) {
match event {
zwlr_screencopy_frame_v1::Event::Buffer {
format,
width,
height,
stride,
} => {
if let Value(f) = format {
tracing::debug!("Received Buffer event with format: {f:?}");
frame.formats.push(FrameFormat {
format: f,
size: Size { width, height },
stride,
})
} else {
tracing::debug!("Received Buffer event with unidentified format");
}
}
zwlr_screencopy_frame_v1::Event::Ready { .. } => {
frame.state.replace(FrameState::Finished);
}
zwlr_screencopy_frame_v1::Event::Failed => {
frame.state.replace(FrameState::Failed);
}
zwlr_screencopy_frame_v1::Event::Damage { .. } => {}
zwlr_screencopy_frame_v1::Event::LinuxDmabuf {
format,
width,
height,
} => {
tracing::debug!(
"Received wlr-screencopy linux_dmabuf event with format: {format} and size {width}x{height}"
);
frame.dmabuf_formats.push(DMAFrameFormat {
format,
size: Size { width, height },
});
}
zwlr_screencopy_frame_v1::Event::BufferDone => {
frame.buffer_done.store(true, Ordering::SeqCst);
}
_ => {}
};
}
}
delegate_noop!(CaptureFrameState: ignore WlShm);
delegate_noop!(CaptureFrameState: ignore WlShmPool);
delegate_noop!(CaptureFrameState: ignore WlBuffer);
delegate_noop!(CaptureFrameState: ignore ZwlrScreencopyManagerV1);
delegate_noop!(CaptureFrameState: ignore ExtImageCopyCaptureManagerV1);
delegate_noop!(CaptureFrameState: ignore ExtOutputImageCaptureSourceManagerV1);
delegate_noop!(CaptureFrameState: ignore ExtImageCaptureSourceV1);
delegate_noop!(CaptureFrameState: ignore ExtForeignToplevelImageCaptureSourceManagerV1);
pub struct WayshotState {}
delegate_noop!(WayshotState: ignore ZwpLinuxDmabufV1);
impl wayland_client::Dispatch<wl_registry::WlRegistry, GlobalListContents> for WayshotState {
fn event(
_: &mut WayshotState,
_: &wl_registry::WlRegistry,
_: wl_registry::Event,
_: &GlobalListContents,
_: &Connection,
_: &QueueHandle<WayshotState>,
) {
}
}
pub struct LayerShellState {
pub configured_outputs: HashSet<WlOutput>,
}
delegate_noop!(LayerShellState: ignore WlCompositor);
delegate_noop!(LayerShellState: ignore WlShm);
delegate_noop!(LayerShellState: ignore WlShmPool);
delegate_noop!(LayerShellState: ignore WlBuffer);
delegate_noop!(LayerShellState: ignore ZwlrLayerShellV1);
delegate_noop!(LayerShellState: ignore WlSurface);
delegate_noop!(LayerShellState: ignore WpViewport);
delegate_noop!(LayerShellState: ignore WpViewporter);
impl wayland_client::Dispatch<ZwlrLayerSurfaceV1, WlOutput> for LayerShellState {
fn event(
state: &mut Self,
proxy: &ZwlrLayerSurfaceV1,
event: <ZwlrLayerSurfaceV1 as wayland_client::Proxy>::Event,
data: &WlOutput,
_conn: &Connection,
_qhandle: &QueueHandle<Self>,
) {
match event {
zwlr_layer_surface_v1::Event::Configure {
serial,
width: _,
height: _,
} => {
tracing::debug!("Acking configure");
state.configured_outputs.insert(data.clone());
proxy.ack_configure(serial);
tracing::trace!("Acked configure");
}
zwlr_layer_surface_v1::Event::Closed => {
tracing::debug!("Closed")
}
_ => {}
}
}
}
pub(crate) struct Card(std::fs::File);
impl AsFd for Card {
fn as_fd(&self) -> BorrowedFd<'_> {
self.0.as_fd()
}
}
impl drm::Device for Card {}
impl Card {
pub fn open<P: AsRef<Path>>(path: P) -> std::io::Result<Self> {
let mut options = std::fs::OpenOptions::new();
options.read(true);
options.write(true);
Ok(Card(options.open(path)?))
}
}
#[derive(Debug)]
pub(crate) struct DMABUFState {
pub linux_dmabuf: ZwpLinuxDmabufV1,
pub gbmdev: gbm::Device<Card>,
}