1use crate::{
2 capture::capture_on,
3 protocols::{
4 export::server::{
5 hyprland_toplevel_export_frame_v1::{
6 self as frame, HyprlandToplevelExportFrameV1 as Frame,
7 },
8 hyprland_toplevel_export_manager_v1::{
9 self as export, HyprlandToplevelExportManagerV1 as Export,
10 },
11 },
12 mapping::server::{
13 hyprland_toplevel_mapping_manager_v1::{
14 self as mapping, HyprlandToplevelMappingManagerV1 as Mapping,
15 },
16 hyprland_toplevel_window_mapping_handle_v1::{
17 self as handle, HyprlandToplevelWindowMappingHandleV1 as Handle,
18 },
19 },
20 },
21};
22use std::{
23 fs::File,
24 io::{BufReader, Write},
25 os::unix::{fs::FileExt, net::UnixStream},
26 sync::{Arc, Mutex},
27 thread,
28};
29use wayland_protocols_wlr::foreign_toplevel::v1::server::{
30 zwlr_foreign_toplevel_handle_v1::{self as toplevel, ZwlrForeignToplevelHandleV1 as Toplevel},
31 zwlr_foreign_toplevel_manager_v1::{self as manager, ZwlrForeignToplevelManagerV1 as Manager},
32};
33use wayland_server::{
34 Client, DataInit, Dispatch, Display, DisplayHandle, GlobalDispatch, New, Resource,
35 backend::{ClientData, ClientId, DisconnectReason},
36 protocol::{
37 wl_buffer::{self, WlBuffer},
38 wl_shm::{self, WlShm},
39 wl_shm_pool::{self, WlShmPool},
40 },
41};
42
43const TARGET: u64 = 0x12345678abcdef01;
44// Same low 32 bits: a truncated-address capture would pick the wrong window.
45const OTHER: u64 = 0xfedcba98abcdef01;
46#[derive(Clone, Copy)]
47enum Mode {
48 Success,
49 Denied,
50 BadBuffer,
51}
52struct State {
53 mode: Mode,
54 captures: Arc<Mutex<Vec<u64>>>,
55}
56#[derive(Debug)]
57struct ClientState;
58impl ClientData for ClientState {
59 fn initialized(&self, _: ClientId) {}
60 fn disconnected(&self, _: ClientId, _: DisconnectReason) {}
61}
62struct Fixture {
63 stop: UnixStream,
64 thread: Option<thread::JoinHandle<()>>,
65 captures: Arc<Mutex<Vec<u64>>>,
66}
67impl Fixture {
68 fn start(mode: Mode) -> (Self, wayland_client::Connection) {
69 let mut display = Display::<State>::new().unwrap();
70 let handle = display.handle();
71 handle.create_global::<State, WlShm, _>(1, ());
72 handle.create_global::<State, Mapping, _>(1, ());
73 handle.create_global::<State, Export, _>(2, ());
74 handle.create_global::<State, Manager, _>(3, ());
75 let (client, server) = UnixStream::pair().unwrap();
76 display
77 .handle()
78 .insert_client(server, Arc::new(ClientState))
79 .unwrap();
80 let (stop, wake) = UnixStream::pair().unwrap();
81 let captures = Arc::new(Mutex::new(Vec::new()));
82 let mut state = State {
83 mode,
84 captures: captures.clone(),
85 };
86 let thread = thread::spawn(move || {
87 loop {
88 display.dispatch_clients(&mut state).unwrap();
89 display.flush_clients().unwrap();
90 let mut fds = [
91 rustix::event::PollFd::new(&display, rustix::event::PollFlags::IN),
92 rustix::event::PollFd::new(&wake, rustix::event::PollFlags::IN),
93 ];
94 rustix::event::poll(&mut fds, None).unwrap();
95 if !fds[1].revents().is_empty() {
96 break;
97 }
98 }
99 });
100 (
101 Self {
102 stop,
103 thread: Some(thread),
104 captures,
105 },
106 wayland_client::Connection::from_socket(client).unwrap(),
107 )
108 }
109}
110impl Drop for Fixture {
111 fn drop(&mut self) {
112 let _ = self.stop.write_all(&[1]);
113 if let Some(thread) = self.thread.take() {
114 thread.join().unwrap();
115 }
116 }
117}
118
119macro_rules! global {
120 ($ty:ty) => {
121 impl GlobalDispatch<$ty, ()> for State {
122 fn bind(
123 _: &mut Self,
124 _: &DisplayHandle,
125 _: &Client,
126 resource: New<$ty>,
127 _: &(),
128 init: &mut DataInit<, Self>,
129 ) {
130 init.init(resource, ());
131 }
132 }
133 };
134}
135global!(Export);
136global!(Mapping);
137impl GlobalDispatch<WlShm, ()> for State {
138 fn bind(
139 _: &mut Self,
140 _: &DisplayHandle,
141 _: &Client,
142 resource: New<WlShm>,
143 _: &(),
144 init: &mut DataInit<, Self>,
145 ) {
146 init.init(resource, ()).format(wl_shm::Format::Argb8888);
147 }
148}
149impl GlobalDispatch<Manager, ()> for State {
150 fn bind(
151 _: &mut Self,
152 dh: &DisplayHandle,
153 client: &Client,
154 resource: New<Manager>,
155 _: &(),
156 init: &mut DataInit<, Self>,
157 ) {
158 let manager = init.init(resource, ());
159 for address in [OTHER, TARGET] {
160 let window = client
161 .create_resource::<Toplevel, _, Self>(dh, 3, address)
162 .unwrap();
163 manager.toplevel(&window);
164 window.title("same title".into());
165 window.app_id("same-app".into());
166 window.done();
167 }
168 }
169}
170impl Dispatch<Manager, ()> for State {
171 fn request(
172 _: &mut Self,
173 _: &Client,
174 _: &Manager,
175 _: manager::Request,
176 _: &(),
177 _: &DisplayHandle,
178 _: &mut DataInit<, Self>,
179 ) {
180 }
181}
182impl Dispatch<Toplevel, u64> for State {
183 fn request(
184 _: &mut Self,
185 _: &Client,
186 _: &Toplevel,
187 _: toplevel::Request,
188 _: &u64,
189 _: &DisplayHandle,
190 _: &mut DataInit<, Self>,
191 ) {
192 }
193}
194impl Dispatch<Handle, ()> for State {
195 fn request(
196 _: &mut Self,
197 _: &Client,
198 _: &Handle,
199 _: handle::Request,
200 _: &(),
201 _: &DisplayHandle,
202 _: &mut DataInit<, Self>,
203 ) {
204 }
205}
206impl Dispatch<Mapping, ()> for State {
207 fn request(
208 _: &mut Self,
209 _: &Client,
210 _: &Mapping,
211 request: mapping::Request,
212 _: &(),
213 _: &DisplayHandle,
214 init: &mut DataInit<, Self>,
215 ) {
216 if let mapping::Request::GetWindowForToplevelWlr { handle, toplevel } = request {
217 let address = *toplevel.data::<u64>().unwrap();
218 init.init(handle, ())
219 .window_address((address >> 32) as u32, address as u32);
220 }
221 }
222}
223impl Dispatch<Export, ()> for State {
224 fn request(
225 state: &mut Self,
226 _: &Client,
227 _: &Export,
228 request: export::Request,
229 _: &(),
230 _: &DisplayHandle,
231 init: &mut DataInit<, Self>,
232 ) {
233 match request {
234 export::Request::CaptureToplevelWithWlrToplevelHandle {
235 frame,
236 overlay_cursor,
237 handle,
238 } => {
239 assert_eq!(overlay_cursor, 0);
240 state
241 .captures
242 .lock()
243 .unwrap()
244 .push(*handle.data::<u64>().unwrap());
245 let frame = init.init(frame, ());
246 match state.mode {
247 Mode::Denied => frame.failed(),
248 Mode::BadBuffer => {
249 frame.buffer(wl_shm::Format::Argb8888, 16384, 16384, 65536);
250 frame.buffer_done();
251 }
252 Mode::Success => {
253 frame.buffer(wl_shm::Format::Argb8888, 2, 1, 12);
254 frame.buffer_done();
255 }
256 }
257 }
258 export::Request::CaptureToplevel { .. } => {
259 panic!("Do not use truncated window addresses")
260 }
261 _ => {}
262 }
263 }
264}
265impl Dispatch<WlShm, ()> for State {
266 fn request(
267 _: &mut Self,
268 _: &Client,
269 _: &WlShm,
270 request: wl_shm::Request,
271 _: &(),
272 _: &DisplayHandle,
273 init: &mut DataInit<, Self>,
274 ) {
275 if let wl_shm::Request::CreatePool { id, fd, .. } = request {
276 init.init(id, Arc::new(File::from(fd)));
277 }
278 }
279}
280struct BufferData {
281 file: Arc<File>,
282 offset: u64,
283}
284impl Dispatch<WlShmPool, Arc<File>> for State {
285 fn request(
286 _: &mut Self,
287 _: &Client,
288 _: &WlShmPool,
289 request: wl_shm_pool::Request,
290 file: &Arc<File>,
291 _: &DisplayHandle,
292 init: &mut DataInit<, Self>,
293 ) {
294 if let wl_shm_pool::Request::CreateBuffer {
295 id,
296 offset,
297 width,
298 height,
299 stride,
300 ..
301 } = request
302 {
303 assert_eq!((width, height, stride), (2, 1, 12));
304 init.init(
305 id,
306 BufferData {
307 file: file.clone(),
308 offset: offset as u64,
309 },
310 );
311 }
312 }
313}
314impl Dispatch<WlBuffer, BufferData> for State {
315 fn request(
316 _: &mut Self,
317 _: &Client,
318 _: &WlBuffer,
319 _: wl_buffer::Request,
320 _: &BufferData,
321 _: &DisplayHandle,
322 _: &mut DataInit<, Self>,
323 ) {
324 }
325}
326impl Dispatch<Frame, ()> for State {
327 fn request(
328 _: &mut Self,
329 _: &Client,
330 resource: &Frame,
331 request: frame::Request,
332 _: &(),
333 _: &DisplayHandle,
334 _: &mut DataInit<, Self>,
335 ) {
336 if let frame::Request::Copy {
337 buffer,
338 ignore_damage,
339 } = request
340 {
341 assert_eq!(ignore_damage, 1, "Static windows must not wait for damage");
342 let data = buffer.data::<BufferData>().unwrap();
343 data.file
344 .write_all_at(&[30, 20, 10, 255, 60, 50, 40, 255, 0, 0, 0, 0], data.offset)
345 .unwrap();
346 resource.flags(frame::Flags::empty());
347 resource.ready(0, 1, 0);
348 }
349 }
350}
351
352#[test]
353fn exports_exact_window_over_real_fd_transport() {
354 let (fixture, connection) = Fixture::start(Mode::Success);
355 let directory = tempfile::tempdir().unwrap();
356 capture_on(connection, TARGET, directory.path()).unwrap();
357 assert_eq!(*fixture.captures.lock().unwrap(), [TARGET]);
358 let mut reader = png::Decoder::new(BufReader::new(
359 File::open(directory.path().join("capture.png")).unwrap(),
360 ))
361 .read_info()
362 .unwrap();
363 let mut pixels = vec![0; reader.output_buffer_size().unwrap()];
364 reader.next_frame(&mut pixels).unwrap();
365 assert_eq!(pixels, [10, 20, 30, 255, 40, 50, 60, 255]);
366 let png = std::fs::read(directory.path().join("capture.png")).unwrap();
367 let idat = png.windows(4).position(|tag| tag == b"IDAT").unwrap();
368 // zlib FLEVEL 1 selects levels 2-5; the slow default level 6 reports 2.
369 assert_eq!(png[idat + 5] >> 6, 1);
370}
371#[test]
372fn failed_export_produces_no_image_and_no_second_capture() {
373 let (fixture, connection) = Fixture::start(Mode::Denied);
374 let directory = tempfile::tempdir().unwrap();
375 assert!(
376 capture_on(connection, TARGET, directory.path())
377 .unwrap_err()
378 .to_string()
379 .contains("permission")
380 );
381 assert_eq!(*fixture.captures.lock().unwrap(), [TARGET]);
382 assert!(!directory.path().join("capture.png").exists());
383}
384#[test]
385fn rejects_oversized_buffers_before_allocating_shared_memory() {
386 let (_fixture, connection) = Fixture::start(Mode::BadBuffer);
387 let directory = tempfile::tempdir().unwrap();
388 assert!(
389 capture_on(connection, TARGET, directory.path())
390 .unwrap_err()
391 .to_string()
392 .contains("too large")
393 );
394 assert!(!directory.path().join("capture.png").exists());
395}