use std::fs::File;
use std::os::fd::{AsFd, OwnedFd};
use color_eyre::eyre::{Result, WrapErr, eyre};
use memmap2::MmapMut;
use rustix::fs::{MemfdFlags, memfd_create};
use wayland_client::protocol::{
wl_buffer::WlBuffer,
wl_output::{self, WlOutput},
wl_registry::WlRegistry,
wl_shm::{Format, WlShm},
wl_shm_pool::WlShmPool,
};
use wayland_client::{Connection, Dispatch, Proxy, QueueHandle, WEnum};
use wayland_protocols_wlr::screencopy::v1::client::{
zwlr_screencopy_frame_v1::{self, ZwlrScreencopyFrameV1},
zwlr_screencopy_manager_v1::ZwlrScreencopyManagerV1,
};
use crate::render::FrozenFrame;
pub(crate) struct OutputFrame {
pub name: Option<String>,
pub frame: FrozenFrame,
}
struct Output {
proxy: WlOutput,
name: Option<String>,
}
#[derive(Default)]
struct FrameState {
format: Option<Format>,
width: u32,
height: u32,
stride: u32,
buffer_done: bool,
ready: bool,
failed: bool,
mmap: Option<MmapMut>,
}
#[derive(Default)]
struct State {
manager: Option<ZwlrScreencopyManagerV1>,
shm: Option<WlShm>,
outputs: Vec<Output>,
frame: Option<FrameState>,
}
pub(crate) fn capture_outputs() -> Result<Vec<OutputFrame>> {
let connection = Connection::connect_to_env().wrap_err("connect for frozen screencopy")?;
let mut queue = connection.new_event_queue();
let mut state = State::default();
connection.display().get_registry(&queue.handle(), ());
queue
.roundtrip(&mut state)
.wrap_err("discover screencopy globals")?;
queue
.roundtrip(&mut state)
.wrap_err("read screencopy output state")?;
if state.manager.is_none() {
return Err(eyre!(
"compositor does not support zwlr_screencopy_manager_v1"
));
}
if state.shm.is_none() {
return Err(eyre!(
"compositor does not support shared-memory screencopy"
));
}
let mut frames = Vec::with_capacity(state.outputs.len());
for index in 0..state.outputs.len() {
let name = state.outputs[index].name.clone();
let frame = capture_output(&mut queue, &mut state, index)?;
frames.push(OutputFrame { name, frame });
}
Ok(frames)
}
fn capture_output(
queue: &mut wayland_client::EventQueue<State>,
state: &mut State,
index: usize,
) -> Result<FrozenFrame> {
let output = state.outputs[index].proxy.clone();
let manager = state.manager.as_ref().expect("manager checked").clone();
let qh = queue.handle();
state.frame = Some(FrameState::default());
let frame = manager.capture_output(0, &output, &qh, ());
while !state
.frame
.as_ref()
.is_some_and(|frame| frame.buffer_done || frame.failed)
{
queue
.blocking_dispatch(state)
.wrap_err("receive screencopy format")?;
}
let captured = state.frame.as_ref().expect("capture state exists");
if captured.failed {
return Err(eyre!("screencopy failed before buffer creation"));
}
let format = captured
.format
.ok_or_else(|| eyre!("screencopy supplied no pixel format"))?;
let size = captured
.stride
.checked_mul(captured.height)
.and_then(|size| usize::try_from(size).ok())
.ok_or_else(|| eyre!("screencopy buffer is too large"))?;
let fd: OwnedFd =
memfd_create("sip-freeze", MemfdFlags::CLOEXEC).wrap_err("create screencopy buffer")?;
let file = File::from(fd);
file.set_len(u64::try_from(size)?)
.wrap_err("size screencopy buffer")?;
let mmap = unsafe { MmapMut::map_mut(&file) }.wrap_err("map screencopy buffer")?;
let shm = state.shm.as_ref().expect("shm checked").clone();
let pool = shm.create_pool(file.as_fd(), i32::try_from(size)?, &qh, ());
let buffer = pool.create_buffer(
0,
i32::try_from(captured.width)?,
i32::try_from(captured.height)?,
i32::try_from(captured.stride)?,
format,
&qh,
(),
);
pool.destroy();
state.frame.as_mut().expect("capture state exists").mmap = Some(mmap);
frame.copy(&buffer);
while !state
.frame
.as_ref()
.is_some_and(|frame| frame.ready || frame.failed)
{
queue
.blocking_dispatch(state)
.wrap_err("receive frozen screencopy")?;
}
buffer.destroy();
frame.destroy();
let captured = state.frame.take().expect("capture state exists");
if captured.failed {
return Err(eyre!("screencopy failed"));
}
rgba(captured, format)
}
fn rgba(frame: FrameState, format: Format) -> Result<FrozenFrame> {
let raw = frame
.mmap
.ok_or_else(|| eyre!("screencopy supplied no pixels"))?;
let capacity = frame
.width
.checked_mul(frame.height)
.and_then(|pixels| pixels.checked_mul(4))
.and_then(|size| usize::try_from(size).ok())
.ok_or_else(|| eyre!("frozen image is too large"))?;
let mut rgba = Vec::with_capacity(capacity);
for row in 0..usize::try_from(frame.height)? {
let start = row * usize::try_from(frame.stride)?;
let end = start + usize::try_from(frame.width)? * 4;
for pixel in raw
.get(start..end)
.ok_or_else(|| eyre!("short screencopy buffer"))?
.chunks_exact(4)
{
match format {
Format::Argb8888 => rgba.extend_from_slice(&[pixel[2], pixel[1], pixel[0], 255]),
Format::Xrgb8888 => rgba.extend_from_slice(&[pixel[2], pixel[1], pixel[0], 255]),
_ => return Err(eyre!("unsupported screencopy pixel format {format:?}")),
}
}
}
Ok(FrozenFrame {
width: frame.width,
height: frame.height,
rgba,
})
}
impl Dispatch<WlRegistry, ()> for State {
fn event(
state: &mut Self,
registry: &WlRegistry,
event: <WlRegistry as Proxy>::Event,
_: &(),
_: &Connection,
qh: &QueueHandle<Self>,
) {
let wayland_client::protocol::wl_registry::Event::Global {
name,
interface,
version,
} = event
else {
return;
};
match interface.as_str() {
"zwlr_screencopy_manager_v1" => {
state.manager = Some(registry.bind(name, version.min(3), qh, ()))
}
"wl_shm" => state.shm = Some(registry.bind(name, 1, qh, ())),
"wl_output" => {
let index = state.outputs.len();
state.outputs.push(Output {
proxy: registry.bind(name, version.min(4), qh, index),
name: None,
});
}
_ => {}
}
}
}
impl Dispatch<WlOutput, usize> for State {
fn event(
state: &mut Self,
_: &WlOutput,
event: wl_output::Event,
index: &usize,
_: &Connection,
_: &QueueHandle<Self>,
) {
if let wl_output::Event::Name { name } = event {
state.outputs[*index].name = Some(name);
}
}
}
macro_rules! ignore {
($interface:ty) => {
impl Dispatch<$interface, ()> for State {
fn event(
_: &mut Self,
_: &$interface,
_: <$interface as Proxy>::Event,
_: &(),
_: &Connection,
_: &QueueHandle<Self>,
) {
}
}
};
}
ignore!(WlShm);
ignore!(WlShmPool);
ignore!(WlBuffer);
ignore!(ZwlrScreencopyManagerV1);
impl Dispatch<ZwlrScreencopyFrameV1, ()> for State {
fn event(
state: &mut Self,
_: &ZwlrScreencopyFrameV1,
event: zwlr_screencopy_frame_v1::Event,
_: &(),
_: &Connection,
_: &QueueHandle<Self>,
) {
let Some(frame) = state.frame.as_mut() else {
return;
};
match event {
zwlr_screencopy_frame_v1::Event::Buffer {
format: WEnum::Value(format),
width,
height,
stride,
} => {
frame.format = Some(format);
frame.width = width;
frame.height = height;
frame.stride = stride;
}
zwlr_screencopy_frame_v1::Event::BufferDone => frame.buffer_done = true,
zwlr_screencopy_frame_v1::Event::Ready { .. } => frame.ready = true,
zwlr_screencopy_frame_v1::Event::Failed => frame.failed = true,
_ => {}
}
}
}