1//! One-shot KWin adapter. The installed executable, not Electron or the AppImage
2//! mount, is the identity authorized by KDE's desktop-entry permission check.
3use serde::{Deserialize, Serialize};
4use std::{
5 collections::HashMap,
6 fs::OpenOptions,
7 io::{Read, Write},
8 os::{fd::AsFd, unix::net::UnixStream},
9 path::Path,
10 sync::mpsc,
11 time::Duration,
12};
13use zbus::{
14 blocking::{Connection, Proxy, connection::Builder},
15 zvariant::{Fd, OwnedValue, Value},
16};
17
18type Result<T> = std::result::Result<T, Box<dyn std::error::Error>>;
19const SCREENSHOT: &str = "org.kde.KWin.ScreenShot2";
20const MAX_BYTES: usize = 128 * 1024 * 1024;
21const DEADLINE: Duration = Duration::from_secs(5);
22
23mod feedback;
24
25#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
26#[serde(rename_all = "camelCase")]
27struct Window {
28 id: String,
29 title: String,
30 app_name: String,
31 app_identifier: String,
32 process_id: u32,
33 bounds: Bounds,
34 client_bounds: Bounds,
35}
36
37impl Window {
38 fn matches_after_capture(&self, current: &Self) -> bool {
39 self.id == current.id
40 && self.process_id == current.process_id
41 && self.app_identifier == current.app_identifier
42 && self.app_name == current.app_name
43 && self.bounds == current.bounds
44 && self.client_bounds == current.client_bounds
45 && capture_title(&self.title) == capture_title(&current.title)
46 }
47}
48
49/// Match the same leading CLI spinner accepted by the Wayland AT-SPI window matcher.
50fn capture_title(title: &str) -> &str {
51 let title = title.trim();
52 let mut chars = title.chars();
53 if matches!(
54 chars.next(),
55 Some('⠋' | '⠙' | '⠹' | '⠸' | '⠼' | '⠴' | '⠦' | '⠧' | '⠇' | '⠏')
56 ) && chars.clone().next().is_some_and(char::is_whitespace)
57 {
58 let label = chars.as_str().trim_start();
59 if !label.is_empty() {
60 return label;
61 }
62 }
63 title
64}
65
66#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
67struct Bounds {
68 x: i32,
69 y: i32,
70 width: u32,
71 height: u32,
72}
73
74struct ScriptReply {
75 owner: String,
76 result: mpsc::SyncSender<String>,
77}
78#[zbus::interface(name = "com.t3tools.KdeCapture")]
79impl ScriptReply {
80 fn reply(
81 &self,
82 value: &str,
83 #[zbus(header)] header: zbus::message::Header<'_>,
84 ) -> zbus::fdo::Result<()> {
85 if header.sender().map(|name| name.as_str()) != Some(self.owner.as_str()) {
86 return Err(zbus::fdo::Error::AccessDenied("Not KWin".into()));
87 }
88 if value.len() > 64 * 1024 {
89 return Err(zbus::fdo::Error::InvalidArgs(
90 "Oversized window metadata".into(),
91 ));
92 }
93 let _ = self.result.try_send(value.to_owned());
94 Ok(())
95 }
96}
97
98/// Scripts exist only for this operation. Unload even when evaluation/callback fails.
99fn script(connection: &Connection, directory: &Path, body: &str) -> Result<String> {
100 let dbus = zbus::blocking::fdo::DBusProxy::new(connection)?;
101 let owner = dbus.get_name_owner("org.kde.KWin".try_into()?)?.to_string();
102 let (send, receive) = mpsc::sync_channel(1);
103 connection.object_server().at(
104 "/com/t3tools/KdeCapture",
105 ScriptReply {
106 owner,
107 result: send,
108 },
109 )?;
110 let path = directory.join("window.js");
111 let name = format!("t3-capture-{}", std::process::id());
112 let destination = serde_json::to_string(
113 connection
114 .unique_name()
115 .ok_or("Missing bus identity")?
116 .as_str(),
117 )?;
118 let source = format!(
119 "function reply(value) {{ callDBus({destination}, '/com/t3tools/KdeCapture', 'com.t3tools.KdeCapture', 'Reply', JSON.stringify(value)); }}\ntry {{ {body} }} catch (error) {{ reply({{error: String(error)}}); }}"
120 );
121 std::fs::write(&path, source)?;
122 let scripting = Proxy::new(
123 connection,
124 "org.kde.KWin",
125 "/Scripting",
126 "org.kde.kwin.Scripting",
127 )?;
128 let result = (|| {
129 let id: i32 = scripting.call(
130 "loadScript",
131 &(path.to_str().ok_or("Invalid script path")?, &name),
132 )?;
133 if id < 0 {
134 return Err("KWin could not load the capture helper script".into());
135 }
136 let result = (|| {
137 let object = format!("/Scripting/Script{id}");
138 let loaded = Proxy::new(
139 connection,
140 "org.kde.KWin",
141 object.as_str(),
142 "org.kde.kwin.Script",
143 )?;
144 let _: () = loaded.call("run", &())?;
145 Ok(receive.recv_timeout(DEADLINE)?)
146 })();
147 let _: std::result::Result<bool, _> = scripting.call("unloadScript", &(&name,));
148 result
149 })();
150 connection
151 .object_server()
152 .remove::<ScriptReply, _>("/com/t3tools/KdeCapture")?;
153 let _ = std::fs::remove_file(path);
154 result
155}
156
157fn window(connection: &Connection, directory: &Path, id: Option<&str>) -> Result<Option<Window>> {
158 let selector = match id {
159 Some(id) => format!(
160 "workspace.windowList().find(w => String(w.internalId) === {})",
161 serde_json::to_string(id)?
162 ),
163 None => "workspace.activeWindow".into(),
164 };
165 let value = script(
166 connection,
167 directory,
168 &format!(
169 "function bounds(rect) {{ return {{ x: Math.round(rect.x), y: Math.round(rect.y), width: Math.round(rect.width), height: Math.round(rect.height) }}; }}
170 const w = {selector}; reply(!w ? null : {{ id: String(w.internalId), title: w.caption, appName: w.resourceClass || w.desktopFileName || 'Application', appIdentifier: w.desktopFileName || w.resourceClass || '', processId: w.pid, bounds: bounds(w.frameGeometry), clientBounds: bounds(w.clientGeometry) }});"
171 ),
172 )?;
173 Ok(serde_json::from_str(&value)?)
174}
175
176fn screenshot_proxy(connection: &Connection) -> zbus::Result<Proxy<'_>> {
177 Proxy::new(
178 connection,
179 SCREENSHOT,
180 "/org/kde/KWin/ScreenShot2",
181 SCREENSHOT,
182 )
183}
184
185/// InvalidWindow follows the authorization check and cannot produce a screenshot.
186fn check(connection: &Connection) -> Result<()> {
187 let proxy = screenshot_proxy(connection)?;
188 if proxy.get_property::<u32>("Version")? < 2 {
189 return Err("KDE's screenshot API is too old".into());
190 }
191 let sink = OpenOptions::new().write(true).open("/dev/null")?;
192 let reply: zbus::Result<HashMap<String, OwnedValue>> = proxy.call(
193 "CaptureWindow",
194 &(
195 "t3-permission-check-not-a-window",
196 HashMap::<&str, Value<'_>>::new(),
197 Fd::from(sink.as_fd()),
198 ),
199 );
200 match reply {
201 Err(zbus::Error::MethodError(name, _, _))
202 if name.as_str() == "org.kde.KWin.ScreenShot2.Error.InvalidWindow" =>
203 {
204 Ok(())
205 }
206 Err(error) => Err(error.into()),
207 Ok(_) => Err("KWin returned an unexpected permission-check result".into()),
208 }
209}
210
211fn uint(metadata: &HashMap<String, OwnedValue>, key: &str) -> Result<usize> {
212 Ok(u32::try_from(
213 metadata
214 .get(key)
215 .ok_or_else(|| format!("Missing screenshot {key}"))?,
216 )? as usize)
217}
218
219/// KWin supplies Qt ARGB32/RGB32 or RGBA8888 pixels. Refuse unknown formats.
220fn rgba(raw: &[u8], width: usize, height: usize, stride: usize, format: usize) -> Result<Vec<u8>> {
221 let row = width
222 .checked_mul(4)
223 .ok_or("Invalid screenshot dimensions")?;
224 let bytes = stride
225 .checked_mul(height)
226 .ok_or("Invalid screenshot dimensions")?;
227 if width == 0 || height == 0 || stride < row || bytes > MAX_BYTES || raw.len() != bytes {
228 return Err("Invalid or oversized screenshot dimensions".into());
229 }
230 if ![4, 5, 6, 17, 18].contains(&format) {
231 return Err("Unsupported KDE screenshot pixel format".into());
232 }
233 let mut output = Vec::with_capacity(row * height);
234 for y in 0..height {
235 for pixel in raw[y * stride..y * stride + row].chunks_exact(4) {
236 let (r, g, b, a) = if format >= 17 {
237 (pixel[0], pixel[1], pixel[2], pixel[3])
238 } else {
239 let argb = u32::from_ne_bytes(pixel.try_into()?);
240 (
241 (argb >> 16) as u8,
242 (argb >> 8) as u8,
243 argb as u8,
244 if format == 4 { 255 } else { (argb >> 24) as u8 },
245 )
246 };
247 let straight = |v: u8| {
248 if a == 0 {
249 0
250 } else {
251 ((v as u32 * 255 + a as u32 / 2) / a as u32).min(255) as u8
252 }
253 };
254 if format == 6 || format == 18 {
255 output.extend_from_slice(&[straight(r), straight(g), straight(b), a]);
256 } else {
257 output.extend_from_slice(&[r, g, b, a]);
258 }
259 }
260 }
261 Ok(output)
262}
263
264fn capture(connection: &Connection, directory: &Path) -> Result<Option<Window>> {
265 // Lack of scripting only loses identity/text; never guess which app was captured.
266 let before = window(connection, directory, None).ok().flatten();
267 let proxy = screenshot_proxy(connection)?;
268 let (mut reader, writer) = UnixStream::pair()?;
269 reader.set_read_timeout(Some(DEADLINE))?;
270 let options = HashMap::from([
271 ("include-decoration", Value::from(true)),
272 ("include-shadow", Value::from(false)),
273 ("include-cursor", Value::from(false)),
274 ("native-resolution", Value::from(true)),
275 ]);
276 let metadata: HashMap<String, OwnedValue> = if let Some(window) = &before {
277 proxy.call(
278 "CaptureWindow",
279 &(&window.id, options, Fd::from(writer.as_fd())),
280 )?
281 } else {
282 proxy.call("CaptureActiveWindow", &(options, Fd::from(writer.as_fd())))?
283 };
284 drop(writer);
285 if <&str>::try_from(metadata.get("type").ok_or("Missing screenshot type")?)? != "raw" {
286 return Err("Unsupported KDE screenshot type".into());
287 }
288 let width = uint(&metadata, "width")?;
289 let height = uint(&metadata, "height")?;
290 let stride = metadata
291 .get("stride")
292 .map(|_| uint(&metadata, "stride"))
293 .transpose()?
294 .unwrap_or(width.checked_mul(4).ok_or("Invalid width")?);
295 let bytes = stride
296 .checked_mul(height)
297 .filter(|bytes| *bytes <= MAX_BYTES)
298 .ok_or("Screenshot too large")?;
299 let mut raw = vec![0; bytes];
300 reader.read_exact(&mut raw)?;
301 let pixels = rgba(&raw, width, height, stride, uint(&metadata, "format")?)?;
302 let image = OpenOptions::new()
303 .write(true)
304 .create_new(true)
305 .open(directory.join("capture.png"))?;
306 let mut encoder = png::Encoder::new(image, width as u32, height as u32);
307 encoder.set_color(png::ColorType::Rgba);
308 encoder.set_depth(png::BitDepth::Eight);
309 let mut image = encoder.write_header()?;
310 image.write_image_data(&pixels)?;
311 image.finish()?;
312 let identity = before.filter(|before| {
313 window(connection, directory, Some(&before.id))
314 .ok()
315 .flatten()
316 .is_some_and(|current| before.matches_after_capture(&current))
317 });
318 Ok(identity)
319}
320
321fn run() -> Result<()> {
322 let args: Vec<String> = std::env::args().collect();
323 let connection = Builder::session()?.method_timeout(DEADLINE).build()?;
324 match args.get(1).map(String::as_str) {
325 Some("check") => {
326 check(&connection)?;
327 println!(
328 "{}",
329 serde_json::json!({"ready": true, "feedbackAvailable": feedback::supported(&connection)})
330 );
331 }
332 Some("capture") if args.len() == 3 => {
333 let window = capture(&connection, Path::new(&args[2]))?;
334 println!("{}", serde_json::json!({"window": window}));
335 }
336 Some("feedback") if args.len() == 4 => {
337 feedback::run(&connection, Path::new(&args[2]), &args[3])?;
338 }
339 Some("activate") if args.len() == 5 => {
340 let pid: u32 = args[3].parse()?;
341 let title = serde_json::to_string(&args[4])?;
342 let value = script(
343 &connection,
344 Path::new(&args[2]),
345 &format!(
346 "const targetPid = {pid}; const targetTitle = {title};\n{}",
347 include_str!("activate.js")
348 ),
349 )?;
350 if serde_json::from_str::<serde_json::Value>(&value)?["activated"] != true {
351 return Err("KDE could not identify The Knife window to activate".into());
352 }
353 println!("{{\"activated\":true}}");
354 }
355 _ => {
356 return Err(
357 "Expected check, capture <directory>, feedback <directory> <options>, or activate <directory> <pid> <title>"
358 .into(),
359 );
360 }
361 }
362 Ok(())
363}
364
365#[cfg(test)]
366mod transport_tests;
367
368fn main() {
369 if let Err(error) = run() {
370 let _ = writeln!(std::io::stderr(), "{error}");
371 std::process::exit(1);
372 }
373}
374
375#[cfg(test)]
376mod tests {
377 use super::*;
378
379 fn window() -> Window {
380 Window {
381 id: "{window-1}".into(),
382 title: "⠋ Document".into(),
383 app_name: "Editor".into(),
384 app_identifier: "editor".into(),
385 process_id: 123,
386 bounds: Bounds {
387 x: 0,
388 y: 0,
389 width: 800,
390 height: 600,
391 },
392 client_bounds: Bounds {
393 x: 0,
394 y: 29,
395 width: 800,
396 height: 571,
397 },
398 }
399 }
400
401 #[test]
402 fn accepts_only_a_leading_cli_spinner_with_a_nonempty_label() {
403 let before = window();
404 for title in [
405 "⠋ Document",
406 "⠙ Document",
407 "⠹ Document",
408 "⠸ Document",
409 "⠼ Document",
410 "⠴ Document",
411 "⠦ Document",
412 "⠧ Document",
413 "⠇ Document",
414 "⠏ Document",
415 "Document",
416 ] {
417 let mut after = before.clone();
418 after.title = title.into();
419 assert!(before.matches_after_capture(&after), "{title}");
420 assert!(after.matches_after_capture(&before), "{title}");
421 }
422 for title in ["⠙ Other", "⠙Document", "⠁ Document", "Document ⠙", "⠙", ""] {
423 let mut after = before.clone();
424 after.title = title.into();
425 assert!(!before.matches_after_capture(&after), "{title}");
426 }
427 assert_ne!(capture_title("⠋"), capture_title("⠙"));
428 }
429
430 #[test]
431 fn spinner_normalization_does_not_weaken_window_identity_checks() {
432 let before = window();
433 for field in [
434 "id",
435 "process_id",
436 "app_identifier",
437 "app_name",
438 "bounds",
439 "client_bounds",
440 ] {
441 let mut after = before.clone();
442 after.title = "⠙ Document".into();
443 match field {
444 "id" => after.id = "{window-2}".into(),
445 "process_id" => after.process_id = 456,
446 "app_identifier" => after.app_identifier = "other".into(),
447 "app_name" => after.app_name = "Other".into(),
448 "bounds" => after.bounds.width = 400,
449 "client_bounds" => after.client_bounds.height = 500,
450 _ => unreachable!(),
451 }
452 assert!(!before.matches_after_capture(&after), "{field}");
453 }
454 }
455
456 #[test]
457 fn decodes_native_argb_and_row_padding() {
458 let raw = [0xff112233_u32.to_ne_bytes(), [0; 4]].concat();
459 assert_eq!(rgba(&raw, 1, 1, 8, 5).unwrap(), [0x11, 0x22, 0x33, 255]);
460 }
461 #[test]
462 fn unpremultiplies_transparency() {
463 assert_eq!(
464 rgba(&[64, 32, 16, 128], 1, 1, 4, 18).unwrap(),
465 [128, 64, 32, 128]
466 );
467 assert_eq!(rgba(&[0, 0, 0, 0], 1, 1, 4, 18).unwrap(), [0, 0, 0, 0]);
468 }
469 #[test]
470 fn rejects_invalid_dimensions_truncation_and_formats() {
471 for (width, height, stride, format) in [
472 (0, 1, 4, 5),
473 (2, 1, 4, 5),
474 (1, 2, 4, 5),
475 (1, 1, 4, 99),
476 (usize::MAX, 1, 4, 5),
477 ] {
478 assert!(rgba(&[0; 4], width, height, stride, format).is_err());
479 }
480 }
481}