1use crate::{
2 Result,
3 protocols::{
4 export::{
5 hyprland_toplevel_export_frame_v1::{
6 self as frame, HyprlandToplevelExportFrameV1 as Frame,
7 },
8 hyprland_toplevel_export_manager_v1::HyprlandToplevelExportManagerV1 as Export,
9 },
10 mapping::{
11 hyprland_toplevel_mapping_manager_v1::HyprlandToplevelMappingManagerV1 as Mapping,
12 hyprland_toplevel_window_mapping_handle_v1::{
13 self as mapping, HyprlandToplevelWindowMappingHandleV1 as Handle,
14 },
15 },
16 },
17};
18use smithay_client_toolkit::{
19 delegate_registry, delegate_shm,
20 registry::{ProvidesRegistryState, RegistryState},
21 shm::{
22 Shm, ShmHandler,
23 slot::{Buffer, SlotPool},
24 },
25};
26use std::{
27 fs::OpenOptions,
28 io::BufWriter,
29 os::unix::fs::OpenOptionsExt,
30 path::Path,
31 time::{Duration, Instant},
32};
33use wayland_client::{
34 Connection, Dispatch, EventQueue, QueueHandle, WEnum, delegate_noop,
35 globals::{GlobalList, registry_queue_init},
36 protocol::wl_shm,
37};
38use wayland_protocols_wlr::foreign_toplevel::v1::client::{
39 zwlr_foreign_toplevel_handle_v1::{self as toplevel, ZwlrForeignToplevelHandleV1 as Toplevel},
40 zwlr_foreign_toplevel_manager_v1::{self as manager, ZwlrForeignToplevelManagerV1 as Manager},
41};
42
43struct Capture {
44 registry: RegistryState,
45 shm: Shm,
46 mapping: Mapping,
47 export: Export,
48 target: u64,
49 started: bool,
50 done: bool,
51 error: Option<String>,
52 pixels: Option<Pixels>,
53 inverted: bool,
54}
55struct Pixels {
56 pool: SlotPool,
57 buffer: Buffer,
58 width: u32,
59 height: u32,
60 stride: u32,
61 format: wl_shm::Format,
62}
63
64fn connect(
65 connection: Connection,
66) -> Result<(Connection, GlobalList, EventQueue<Capture>, Capture)> {
67 let (globals, queue) = registry_queue_init::<Capture>(&connection)?;
68 let qh = queue.handle();
69 let mapping = globals.bind::<Mapping,_,_>(&qh, 1..=1, ())
70 .map_err(|_| "Update Hyprland: its window-mapping protocol is required for exact active-window capture.")?;
71 let export = globals
72 .bind::<Export, _, _>(&qh, 2..=2, ())
73 .map_err(|_| "This Hyprland version does not support window export v2.")?;
74 let state = Capture {
75 registry: RegistryState::new(&globals),
76 shm: Shm::bind(&globals, &qh)?,
77 mapping,
78 export,
79 target: 0,
80 started: false,
81 done: false,
82 error: None,
83 pixels: None,
84 inverted: false,
85 };
86 Ok((connection, globals, queue, state))
87}
88pub fn check() -> Result<bool> {
89 let (_, globals, queue, _) = connect(Connection::connect_to_env()?)?;
90 let _manager: Manager = globals.bind(&queue.handle(), 1..=3, ())?;
91 Ok(globals.contents().with_list(|list| {
92 [
93 "zwlr_layer_shell_v1",
94 "wp_viewporter",
95 "wl_subcompositor",
96 "wp_presentation",
97 ]
98 .iter()
99 .all(|name| list.iter().any(|g| &g.interface == name))
100 }))
101}
102
103pub fn capture(address: u64, directory: &Path) -> Result<()> {
104 capture_on(Connection::connect_to_env()?, address, directory)
105}
106
107pub(super) fn capture_on(connection: Connection, address: u64, directory: &Path) -> Result<()> {
108 let (connection, globals, mut queue, mut state) = connect(connection)?;
109 state.target = address;
110 let _manager: Manager = globals.bind(&queue.handle(), 1..=3, ())?;
111 let deadline = Instant::now() + Duration::from_secs(15);
112 while !state.done && state.error.is_none() {
113 queue.dispatch_pending(&mut state)?;
114 if state.done || state.error.is_some() {
115 break;
116 }
117 connection.flush()?;
118 if let Some(read) = queue.prepare_read() {
119 let timeout = rustix::event::Timespec::try_from(deadline.checked_duration_since(Instant::now()).ok_or("Hyprland capture timed out. Approve its screen-sharing prompt, then try again.")?)?;
120 let mut fds = [rustix::event::PollFd::new(
121 &connection,
122 rustix::event::PollFlags::IN,
123 )];
124 if rustix::event::poll(&mut fds, Some(&timeout))? == 0 {
125 return Err("Hyprland capture timed out. Check screen-sharing permission.".into());
126 }
127 read.read()?;
128 }
129 }
130 if let Some(error) = state.error {
131 return Err(error.into());
132 }
133 let pixels = state.pixels.as_mut().ok_or("Hyprland returned no image.")?;
134 let bytes = pixels.pool.raw_data_mut(&pixels.buffer.slot());
135 let rgba = to_rgba(
136 bytes,
137 pixels.width,
138 pixels.height,
139 pixels.stride,
140 pixels.format,
141 state.inverted,
142 )?;
143 let file = OpenOptions::new()
144 .write(true)
145 .create_new(true)
146 .mode(0o600)
147 .open(directory.join("capture.png"))?;
148 let mut encoder = png::Encoder::new(BufWriter::new(file), pixels.width, pixels.height);
149 encoder.set_color(png::ColorType::Rgba);
150 encoder.set_depth(png::BitDepth::Eight);
151 encoder.write_header()?.write_image_data(&rgba)?;
152 Ok(())
153}
154
155fn image_size(width: u32, height: u32, stride: u32) -> Result<usize> {
156 if width == 0 || height == 0 || width > 16384 || height > 16384 || stride < width * 4 {
157 return Err("Invalid capture dimensions.".into());
158 }
159 let size = u64::from(height) * u64::from(stride);
160 if size > 128 * 1024 * 1024 {
161 return Err("Capture is too large.".into());
162 }
163 Ok(size as usize)
164}
165
166fn to_rgba(
167 bytes: &[u8],
168 width: u32,
169 height: u32,
170 stride: u32,
171 format: wl_shm::Format,
172 inverted: bool,
173) -> Result<Vec<u8>> {
174 let size = image_size(width, height, stride)?;
175 if bytes.len() < size {
176 return Err("Truncated capture buffer.".into());
177 }
178 let (red, blue, alpha) = match format {
179 wl_shm::Format::Argb8888 => (2, 0, true),
180 wl_shm::Format::Xrgb8888 => (2, 0, false),
181 wl_shm::Format::Abgr8888 => (0, 2, true),
182 wl_shm::Format::Xbgr8888 => (0, 2, false),
183 _ => return Err("Hyprland returned an unsupported capture pixel format.".into()),
184 };
185 let mut rgba = vec![0; (width * height * 4) as usize];
186 for y in 0..height {
187 let row = if inverted { height - 1 - y } else { y };
188 for x in 0..width {
189 let source = (row * stride + x * 4) as usize;
190 let dest = ((y * width + x) * 4) as usize;
191 let a = if alpha { bytes[source + 3] } else { 255 };
192 for (i, channel) in [red, 1, blue].into_iter().enumerate() {
193 rgba[dest + i] = if a == 0 {
194 0
195 } else {
196 (u32::from(bytes[source + channel]) * 255 / u32::from(a)).min(255) as u8
197 };
198 }
199 rgba[dest + 3] = a;
200 }
201 }
202 Ok(rgba)
203}
204
205impl Dispatch<Manager, ()> for Capture {
206 fn event(
207 state: &mut Self,
208 _: &Manager,
209 event: manager::Event,
210 _: &(),
211 _: &Connection,
212 qh: &QueueHandle<Self>,
213 ) {
214 if let manager::Event::Toplevel { toplevel } = event {
215 state
216 .mapping
217 .get_window_for_toplevel_wlr(&toplevel, qh, toplevel.clone());
218 }
219 }
220 wayland_client::event_created_child!(Capture, Manager, [0 => (Toplevel, ())]);
221}
222impl Dispatch<Toplevel, ()> for Capture {
223 fn event(
224 _: &mut Self,
225 handle: &Toplevel,
226 event: toplevel::Event,
227 _: &(),
228 _: &Connection,
229 _: &QueueHandle<Self>,
230 ) {
231 if let toplevel::Event::Closed = event {
232 handle.destroy();
233 }
234 }
235}
236impl Dispatch<Handle, Toplevel> for Capture {
237 fn event(
238 state: &mut Self,
239 handle: &Handle,
240 event: mapping::Event,
241 toplevel: &Toplevel,
242 _: &Connection,
243 qh: &QueueHandle<Self>,
244 ) {
245 if let mapping::Event::WindowAddress {
246 address_hi,
247 address,
248 } = event
249 {
250 if (u64::from(address_hi) << 32 | u64::from(address)) == state.target && !state.started
251 {
252 state.started = true;
253 state
254 .export
255 .capture_toplevel_with_wlr_toplevel_handle(0, toplevel, qh, ());
256 }
257 }
258 handle.destroy();
259 }
260}
261impl Dispatch<Frame, ()> for Capture {
262 fn event(
263 state: &mut Self,
264 frame: &Frame,
265 event: frame::Event,
266 _: &(),
267 _: &Connection,
268 _: &QueueHandle<Self>,
269 ) {
270 if state.error.is_some() || state.done {
271 return;
272 }
273 let result = (|| -> Result<()> {
274 match event {
275 frame::Event::Buffer { format: WEnum::Value(format), width, height, stride } => {
276 let mut pool = SlotPool::new(image_size(width, height, stride)?, &state.shm)?;
277 let (buffer, _) = pool.create_buffer(width as i32, height as i32, stride as i32, format)?;
278 state.pixels = Some(Pixels { pool, buffer, width, height, stride, format });
279 }
280 frame::Event::BufferDone => {
281 let pixels = state.pixels.as_ref().ok_or("Hyprland did not offer a shared-memory image.")?;
282 frame.copy(pixels.buffer.wl_buffer(), 1);
283 }
284 frame::Event::Flags { flags: WEnum::Value(flags) } => state.inverted = flags.contains(frame::Flags::YInvert),
285 frame::Event::Ready { .. } => { state.done = true; frame.destroy(); }
286 frame::Event::Failed => return Err("Hyprland could not capture this window. Check screen-sharing permission and try again.".into()),
287 _ => {}
288 }
289 Ok(())
290 })();
291 if let Err(error) = result {
292 state.error = Some(error.to_string());
293 }
294 }
295}
296delegate_noop!(Capture: ignore Export);
297delegate_noop!(Capture: ignore Mapping);
298delegate_registry!(Capture);
299delegate_shm!(Capture);
300impl ProvidesRegistryState for Capture {
301 fn registry(&mut self) -> &mut RegistryState {
302 &mut self.registry
303 }
304 smithay_client_toolkit::registry_handlers!();
305}
306impl ShmHandler for Capture {
307 fn shm_state(&mut self) -> &mut Shm {
308 &mut self.shm
309 }
310}
311
312#[cfg(test)]
313mod tests {
314 use super::*;
315 #[test]
316 fn bounds_buffer_allocations() {
317 assert!(image_size(100, 100, 399).is_err());
318 assert!(image_size(16384, 16384, 65536).is_err());
319 assert!(image_size(0, 10, 4).is_err());
320 }
321 #[test]
322 fn converts_channels_padding_flip_and_alpha() {
323 let bytes = [10, 20, 30, 255, 0, 0, 0, 0, 25, 50, 100, 128, 0, 0, 0, 0];
324 assert_eq!(
325 to_rgba(&bytes, 1, 2, 8, wl_shm::Format::Argb8888, true).unwrap(),
326 [199, 99, 49, 128, 30, 20, 10, 255]
327 );
328 assert_eq!(
329 to_rgba(&bytes, 1, 1, 8, wl_shm::Format::Xbgr8888, false).unwrap(),
330 [10, 20, 30, 255]
331 );
332 assert!(to_rgba(&bytes[..2], 1, 1, 4, wl_shm::Format::Argb8888, false).is_err());
333 }
334}