remote trackpad and keyboard server
git clone https://git.lucas.co/cce-remote.git
src/screencopy.rs (11.5K)
1 //! Persistent wlr-screencopy capture — the robust replacement for forking
2 //! `grim` per frame.
3 //!
4 //! One long-lived Wayland connection per stream: each frame is a
5 //! `capture_output_region` of the focused window's rect, throttled by
6 //! `copy_with_damage` — the compositor withholds the frame until the region
7 //! actually changes, so an idle desktop costs nothing and an active one
8 //! streams at compositor pace instead of fork/exec pace. Shm buffers are
9 //! reused across frames; pixels are box-downscaled and JPEG-encoded
10 //! in-process (`jpeg-encoder`).
11 //!
12 //! Single-output assumption: the region is passed in output-local logical
13 //! coordinates, which equals layout coordinates when the (only) output sits
14 //! at 0,0 — true for this DE's eDP-1 setup, same assumption grim ran under.
15
16 use std::os::fd::AsFd;
17 use std::time::{Duration, Instant};
18
19 use wayland_client::globals::{registry_queue_init, GlobalListContents};
20 use wayland_client::protocol::{wl_buffer, wl_output, wl_registry, wl_shm, wl_shm_pool};
21 use wayland_client::{delegate_noop, Connection, Dispatch, EventQueue, QueueHandle, WEnum};
22 use wayland_protocols_wlr::screencopy::v1::client::{
23 zwlr_screencopy_frame_v1::{self, ZwlrScreencopyFrameV1},
24 zwlr_screencopy_manager_v1::ZwlrScreencopyManagerV1,
25 };
26
27
28 #[derive(Default)]
29 struct CapState {
30 // per-frame handshake state, reset before each request
31 buffer_meta: Option<(u32, u32, u32, wl_shm::Format)>, // w, h, stride, format
32 buffer_done: bool,
33 ready: bool,
34 failed: bool,
35 }
36
37 impl Dispatch<wl_registry::WlRegistry, GlobalListContents> for CapState {
38 fn event(
39 _: &mut Self,
40 _: &wl_registry::WlRegistry,
41 _: wl_registry::Event,
42 _: &GlobalListContents,
43 _: &Connection,
44 _: &QueueHandle<Self>,
45 ) {
46 }
47 }
48
49 delegate_noop!(CapState: ignore wl_shm::WlShm);
50 delegate_noop!(CapState: ignore wl_output::WlOutput);
51 delegate_noop!(CapState: ignore wl_shm_pool::WlShmPool);
52 delegate_noop!(CapState: ignore wl_buffer::WlBuffer);
53 delegate_noop!(CapState: ignore ZwlrScreencopyManagerV1);
54
55 impl Dispatch<ZwlrScreencopyFrameV1, ()> for CapState {
56 fn event(
57 state: &mut Self,
58 _: &ZwlrScreencopyFrameV1,
59 event: zwlr_screencopy_frame_v1::Event,
60 _: &(),
61 _: &Connection,
62 _: &QueueHandle<Self>,
63 ) {
64 match event {
65 zwlr_screencopy_frame_v1::Event::Buffer { format, width, height, stride } => {
66 if let WEnum::Value(f) = format {
67 // prefer 32-bit formats we know how to swizzle
68 if matches!(
69 f,
70 wl_shm::Format::Xrgb8888
71 | wl_shm::Format::Argb8888
72 | wl_shm::Format::Xbgr8888
73 | wl_shm::Format::Abgr8888
74 ) {
75 state.buffer_meta = Some((width, height, stride, f));
76 }
77 }
78 }
79 zwlr_screencopy_frame_v1::Event::BufferDone => state.buffer_done = true,
80 zwlr_screencopy_frame_v1::Event::Ready { .. } => state.ready = true,
81 zwlr_screencopy_frame_v1::Event::Failed => state.failed = true,
82 _ => {}
83 }
84 }
85 }
86
87 struct ShmSlot {
88 _file: std::fs::File,
89 map: memmap2::MmapMut,
90 pool: wl_shm_pool::WlShmPool,
91 buffer: wl_buffer::WlBuffer,
92 meta: (u32, u32, u32, wl_shm::Format),
93 }
94
95 pub struct CaptureSession {
96 conn: Connection,
97 queue: EventQueue<CapState>,
98 qh: QueueHandle<CapState>,
99 state: CapState,
100 manager: ZwlrScreencopyManagerV1,
101 output: wl_output::WlOutput,
102 shm: wl_shm::WlShm,
103 slot: Option<ShmSlot>,
104 }
105
106 impl CaptureSession {
107 pub fn new() -> Result<Self, String> {
108 let conn = Connection::connect_to_env().map_err(|e| e.to_string())?;
109 let (globals, queue) =
110 registry_queue_init::<CapState>(&conn).map_err(|e| e.to_string())?;
111 let qh = queue.handle();
112 let manager: ZwlrScreencopyManagerV1 = globals
113 .bind(&qh, 3..=3, ())
114 .map_err(|e| format!("screencopy v3 unavailable: {e}"))?;
115 let shm: wl_shm::WlShm = globals.bind(&qh, 1..=1, ()).map_err(|e| e.to_string())?;
116 let output: wl_output::WlOutput =
117 globals.bind(&qh, 1..=4, ()).map_err(|e| e.to_string())?;
118 Ok(Self {
119 conn,
120 queue,
121 qh,
122 state: CapState::default(),
123 manager,
124 output,
125 shm,
126 slot: None,
127 })
128 }
129
130 /// Pump the event queue until `pred(state)` holds or the deadline passes.
131 /// Returns false on timeout (Wayland connection still healthy).
132 fn wait_until(
133 &mut self,
134 deadline: Instant,
135 pred: impl Fn(&CapState) -> bool,
136 ) -> Result<bool, String> {
137 loop {
138 self.queue
139 .dispatch_pending(&mut self.state)
140 .map_err(|e| e.to_string())?;
141 if pred(&self.state) {
142 return Ok(true);
143 }
144 if Instant::now() >= deadline {
145 return Ok(false);
146 }
147 self.conn.flush().map_err(|e| e.to_string())?;
148 if let Some(guard) = self.queue.prepare_read() {
149 let fd = guard.connection_fd();
150 let mut fds = [rustix::event::PollFd::new(
151 &fd,
152 rustix::event::PollFlags::IN,
153 )];
154 let remaining = deadline.saturating_duration_since(Instant::now());
155 let ms = remaining.as_millis().min(200) as i32;
156 let _ = rustix::event::poll(&mut fds, ms.max(1));
157 let readable = fds[0].revents().contains(rustix::event::PollFlags::IN);
158 drop(fds);
159 if readable {
160 let _ = guard.read();
161 } // else: drop the guard without reading and re-check
162 }
163 }
164 }
165
166 fn ensure_slot(&mut self, meta: (u32, u32, u32, wl_shm::Format)) -> Result<(), String> {
167 if let Some(s) = &self.slot {
168 if s.meta == meta {
169 return Ok(());
170 }
171 }
172 if let Some(old) = self.slot.take() {
173 old.buffer.destroy();
174 old.pool.destroy();
175 }
176 let (w, h, stride, format) = meta;
177 let size = (stride * h) as usize;
178 let dir = std::path::Path::new("/dev/shm");
179 let dir = if dir.is_dir() { dir } else { std::path::Path::new("/tmp") };
180 let path = dir.join(format!("cce-remote-shm-{}", std::process::id()));
181 let file = std::fs::OpenOptions::new()
182 .read(true)
183 .write(true)
184 .create(true)
185 .truncate(true)
186 .open(&path)
187 .map_err(|e| e.to_string())?;
188 // unlink immediately — the fd keeps it alive, nothing lingers on disk
189 let _ = std::fs::remove_file(&path);
190 file.set_len(size as u64).map_err(|e| e.to_string())?;
191 let map = unsafe { memmap2::MmapMut::map_mut(&file) }.map_err(|e| e.to_string())?;
192 let pool = self.shm.create_pool(file.as_fd(), size as i32, &self.qh, ());
193 let buffer = pool.create_buffer(
194 0,
195 w as i32,
196 h as i32,
197 stride as i32,
198 format,
199 &self.qh,
200 (),
201 );
202 self.slot = Some(ShmSlot { _file: file, map, pool, buffer, meta });
203 Ok(())
204 }
205
206 /// Capture one frame of `rect` (output-local logical px). With
207 /// `use_damage`, blocks until the region changes or `timeout` — a timeout
208 /// returns Ok(None) so the caller can force a keepalive frame. The pixels
209 /// are returned RAW (copied out of the shm slot, which is reused);
210 /// encoding happens at send time so dropped frames cost nothing.
211 pub fn next_frame(
212 &mut self,
213 rect: (i32, i32, i32, i32),
214 use_damage: bool,
215 timeout: Duration,
216 ) -> Result<Option<crate::stream::Payload>, String> {
217 self.state = CapState::default();
218 let frame = self.manager.capture_output_region(
219 1, // overlay the cursor — the remote wants to see it
220 &self.output,
221 rect.0,
222 rect.1,
223 rect.2,
224 rect.3,
225 &self.qh,
226 (),
227 );
228 // Phase 1: buffer negotiation (always fast).
229 let ok = self.wait_until(Instant::now() + Duration::from_secs(5), |s| {
230 s.buffer_done || s.failed
231 })?;
232 let Some(meta) = self.state.buffer_meta else {
233 frame.destroy();
234 return Err("no usable shm format offered".into());
235 };
236 if !ok || self.state.failed {
237 frame.destroy();
238 return if self.state.failed { Err("capture failed".into()) } else { Ok(None) };
239 }
240 self.ensure_slot(meta)?;
241 let buffer = &self.slot.as_ref().unwrap().buffer;
242 if use_damage {
243 frame.copy_with_damage(buffer);
244 } else {
245 frame.copy(buffer);
246 }
247 // Phase 2: damage-gated (this is the idle throttle).
248 let ok = self.wait_until(Instant::now() + timeout, |s| s.ready || s.failed)?;
249 let failed = self.state.failed;
250 if !ok || failed {
251 frame.destroy();
252 return if failed { Err("copy failed".into()) } else { Ok(None) };
253 }
254 frame.destroy();
255 let slot = self.slot.as_ref().unwrap();
256 let (w, h, stride, format) = slot.meta;
257 // byte offsets of R,G,B within each little-endian 32-bit pixel
258 let rgb = match format {
259 wl_shm::Format::Xrgb8888 | wl_shm::Format::Argb8888 => (2usize, 1usize, 0usize),
260 _ => (0usize, 1usize, 2usize), // Xbgr8888 / Abgr8888
261 };
262 let size = (stride * h) as usize;
263 Ok(Some(crate::stream::Payload::Raw {
264 data: slot.map[..size].to_vec(),
265 w,
266 h,
267 stride,
268 rgb,
269 }))
270 }
271 }
272
273 /// Shared by both raw frame sources (screencopy shm and the compositor's
274 /// window-stream RGBA): box-downscale 32-bit pixels to ≤ `max_edge` and JPEG
275 /// them at `quality`. `rgb_at` gives the byte offsets of R,G,B within each
276 /// 4-byte pixel. Called per SENT frame, with the (edge, quality) the
277 /// adaptation ladder picked for the link.
278 pub fn downscale_encode(
279 data: &[u8],
280 w: u32,
281 h: u32,
282 stride: u32,
283 (ri, gi, bi): (usize, usize, usize),
284 max_edge: u32,
285 quality: u8,
286 ) -> Result<Vec<u8>, String> {
287 if w == 0 || h == 0 || (stride * h) as usize > data.len() {
288 return Err("bad frame dimensions".into());
289 }
290 let f = ((w.max(h) + max_edge - 1) / max_edge).max(1);
291 let (ow, oh) = (w / f, h / f);
292 let mut rgb = Vec::with_capacity((ow * oh * 3) as usize);
293 let fsq = (f * f) as u32;
294 for oy in 0..oh {
295 for ox in 0..ow {
296 let (mut r, mut g, mut b) = (0u32, 0u32, 0u32);
297 for sy in 0..f {
298 let row = ((oy * f + sy) * stride) as usize;
299 for sx in 0..f {
300 let px = row + ((ox * f + sx) * 4) as usize;
301 r += data[px + ri] as u32;
302 g += data[px + gi] as u32;
303 b += data[px + bi] as u32;
304 }
305 }
306 rgb.push((r / fsq) as u8);
307 rgb.push((g / fsq) as u8);
308 rgb.push((b / fsq) as u8);
309 }
310 }
311 let mut out = Vec::new();
312 let encoder = jpeg_encoder::Encoder::new(&mut out, quality);
313 encoder
314 .encode(&rgb, ow as u16, oh as u16, jpeg_encoder::ColorType::Rgb)
315 .map_err(|e| e.to_string())?;
316 Ok(out)
317 }