1use crate::diagnostics::hydrate_session_bus_env;
2use crate::x11_display::{is_native_x11_session, with_x11_display, X11_CAPTURE_TIMEOUT};
3use anyhow::{anyhow, bail, Context, Result};
4use base64::{engine::general_purpose::STANDARD, Engine};
5use futures_util::StreamExt;
6use image::codecs::jpeg::JpegEncoder;
7use image::imageops::FilterType;
8use schemars::JsonSchema;
9use serde::{Deserialize, Serialize};
10use std::{
11 collections::HashMap,
12 fs::{self, OpenOptions},
13 io::{Cursor, Read},
14 os::unix::fs::OpenOptionsExt,
15 path::{Path, PathBuf},
16 process::Stdio,
17 time::{Duration, SystemTime, UNIX_EPOCH},
18};
19use tokio::process::Command;
20use zbus::{
21 message::{Message, Type as MessageType},
22 names::{BusName, OwnedUniqueName, WellKnownName},
23 zvariant::{OwnedObjectPath, OwnedValue, Value},
24 MatchRule, MessageStream, Proxy,
25};
26
27const PORTAL_DESKTOP_SERVICE: &str = "org.freedesktop.portal.Desktop";
28const PORTAL_DESKTOP_PATH: &str = "/org/freedesktop/portal/desktop";
29const PORTAL_SCREENSHOT_INTERFACE: &str = "org.freedesktop.portal.Screenshot";
30const PORTAL_REQUEST_INTERFACE: &str = "org.freedesktop.portal.Request";
31const PORTAL_REQUEST_PATH_NAMESPACE: &str = "/org/freedesktop/portal/desktop/request";
32const MAX_SCREENSHOT_SOURCE_BYTES: u64 = 64 * 1024 * 1024;
33const MAX_SCREENSHOT_SOURCE_PIXELS: u64 = 128 * 1024 * 1024;
34
35pub const DEFAULT_SCREENSHOT_MAX_DIMENSION: u32 = 1920;
36pub const DEFAULT_SCREENSHOT_MAX_BYTES: usize = 2 * 1024 * 1024;
37pub const ABSOLUTE_SCREENSHOT_MAX_DIMENSION: u32 = 4096;
38pub const ABSOLUTE_SCREENSHOT_MAX_BYTES: usize = 4 * 1024 * 1024;
39pub const DEFAULT_SCREENSHOT_JPEG_QUALITY: u8 = 80;
40pub const MIN_SCREENSHOT_JPEG_QUALITY: u8 = 1;
41pub const MAX_SCREENSHOT_JPEG_QUALITY: u8 = 95;
42const MIN_SCREENSHOT_MAX_BYTES: usize = 1024;
43
44#[derive(Debug, Clone)]
45pub struct RawScreenshotCapture {
46 pub mime_type: String,
47 pub bytes: Vec<u8>,
48 pub source: String,
49 pub width: u32,
50 pub height: u32,
51}
52
53#[derive(Debug, Clone, Serialize, JsonSchema)]
54pub struct ScreenshotCapture {
55 pub mime_type: String,
56 pub data_url: String,
57 pub source: String,
58 pub width: u32,
60 pub height: u32,
62 pub coordinate_width: u32,
64 pub coordinate_height: u32,
66 pub scale: f32,
68 pub resized: bool,
69 pub bytes: usize,
70 pub original_bytes: usize,
71 pub max_bytes: usize,
72 pub format: ScreenshotOutputFormat,
73 pub quality: Option<u8>,
74}
75
76#[derive(Debug, Clone, Copy, Default)]
77pub struct ScreenshotPayloadOptions {
78 pub max_width: Option<u32>,
79 pub max_height: Option<u32>,
80 pub max_bytes: Option<usize>,
81 pub scale: Option<f32>,
82 pub format: Option<ScreenshotOutputFormat>,
83 pub quality: Option<u8>,
84}
85
86#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
87#[serde(rename_all = "lowercase")]
88pub enum ScreenshotOutputFormat {
89 Png,
90 Jpeg,
91}
92
93impl ScreenshotOutputFormat {
94 fn mime_type(self) -> &'static str {
95 match self {
96 Self::Png => "image/png",
97 Self::Jpeg => "image/jpeg",
98 }
99 }
100}
101
102#[derive(Debug, Clone, Copy)]
103struct ResolvedScreenshotPayloadOptions {
104 max_width: u32,
105 max_height: u32,
106 max_bytes: usize,
107 scale: f32,
108 format: ScreenshotOutputFormat,
109 quality: u8,
110}
111
112#[derive(Debug, Clone, PartialEq, Eq)]
113enum ScreenshotCleanup {
114 DeletePath(PathBuf),
115 Preserve,
116}
117
118impl ScreenshotPayloadOptions {
119 fn resolve(self) -> ResolvedScreenshotPayloadOptions {
120 let max_width = self
121 .max_width
122 .unwrap_or(DEFAULT_SCREENSHOT_MAX_DIMENSION)
123 .clamp(1, ABSOLUTE_SCREENSHOT_MAX_DIMENSION);
124 let max_height = self
125 .max_height
126 .unwrap_or(DEFAULT_SCREENSHOT_MAX_DIMENSION)
127 .clamp(1, ABSOLUTE_SCREENSHOT_MAX_DIMENSION);
128 let max_bytes = self
129 .max_bytes
130 .unwrap_or(DEFAULT_SCREENSHOT_MAX_BYTES)
131 .clamp(MIN_SCREENSHOT_MAX_BYTES, ABSOLUTE_SCREENSHOT_MAX_BYTES);
132 let scale = self
133 .scale
134 .filter(|value| value.is_finite() && *value > 0.0)
135 .unwrap_or(1.0)
136 .min(1.0);
137 let format = self.format.unwrap_or(ScreenshotOutputFormat::Png);
138 let quality = self
139 .quality
140 .unwrap_or(DEFAULT_SCREENSHOT_JPEG_QUALITY)
141 .clamp(MIN_SCREENSHOT_JPEG_QUALITY, MAX_SCREENSHOT_JPEG_QUALITY);
142
143 ResolvedScreenshotPayloadOptions {
144 max_width,
145 max_height,
146 max_bytes,
147 scale,
148 format,
149 quality,
150 }
151 }
152}
153
154const SCREENSHOT_BACKEND_ENV: &str = "COMPUTER_USE_LINUX_SCREENSHOT_BACKEND";
157
158#[derive(Debug, Clone, Copy, PartialEq, Eq)]
159enum ScreenshotBackend {
160 GnomeShell,
161 Portal,
162 X11,
163 GnomeScreenshot,
164}
165
166impl ScreenshotBackend {
167 fn parse(value: &str) -> Option<Self> {
168 match value.trim().to_ascii_lowercase().as_str() {
169 "gnome-shell" | "gnome_shell" | "shell" => Some(Self::GnomeShell),
170 "portal" | "xdg-portal" | "xdg_portal" => Some(Self::Portal),
171 "x11" | "x11-native" | "x11_native" | "xgetimage" => Some(Self::X11),
172 "gnome-screenshot" | "gnome_screenshot" => Some(Self::GnomeScreenshot),
173 _ => None,
174 }
175 }
176
177 async fn capture(self) -> Result<RawScreenshotCapture> {
178 match self {
179 Self::GnomeShell => capture_with_gnome_shell().await,
180 Self::Portal => capture_with_portal().await,
181 Self::X11 => capture_with_x11().await,
182 Self::GnomeScreenshot => capture_with_gnome_screenshot().await,
183 }
184 }
185}
186
187pub async fn capture_screenshot_raw() -> Result<RawScreenshotCapture> {
188 hydrate_session_bus_env();
189
190 if let Some(forced) = forced_backend()? {
194 return forced.capture().await;
195 }
196
197 let gnome_error = match capture_with_gnome_shell().await {
204 Ok(capture) => return Ok(capture),
205 Err(error) => error,
206 };
207 let portal_error = match capture_with_portal().await {
208 Ok(capture) => return Ok(capture),
209 Err(error) => error,
210 };
211 let x11_error = match capture_with_x11().await {
215 Ok(capture) => return Ok(capture),
216 Err(error) => error,
217 };
218 let cli_error = match capture_with_gnome_screenshot().await {
219 Ok(capture) => return Ok(capture),
220 Err(error) => error,
221 };
222
223 Err(anyhow!(
224 "GNOME Shell screenshot failed: {gnome_error:#}; \
225 XDG portal screenshot failed: {portal_error:#}; \
226 native X11 screenshot failed: {x11_error:#}; \
227 gnome-screenshot fallback failed: {cli_error:#}"
228 ))
229}
230
231async fn capture_with_x11() -> Result<RawScreenshotCapture> {
234 if !is_native_x11_session() {
235 bail!("not a native X11 session (needs DISPLAY on an X11, not Wayland, session)");
236 }
237 let image = with_x11_display(X11_CAPTURE_TIMEOUT, |display| display.capture_root()).await??;
238 let (width, height) = (image.width, image.height);
239 let bytes =
240 tokio::task::spawn_blocking(move || encode_rgb_png(image.width, image.height, image.rgb))
241 .await
242 .context("X11 screenshot encoder task failed")??;
243 Ok(RawScreenshotCapture {
244 mime_type: "image/png".to_string(),
245 bytes,
246 source: "x11".to_string(),
247 width,
248 height,
249 })
250}
251
252fn encode_rgb_png(width: u32, height: u32, rgb: Vec<u8>) -> Result<Vec<u8>> {
253 let buffer = image::RgbImage::from_raw(width, height, rgb)
254 .context("X11 root image did not match its dimensions")?;
255 let mut out = Vec::new();
256 image::DynamicImage::ImageRgb8(buffer)
257 .write_to(&mut Cursor::new(&mut out), image::ImageFormat::Png)
258 .context("failed to encode X11 screenshot PNG")?;
259 Ok(out)
260}
261
262fn forced_backend() -> Result<Option<ScreenshotBackend>> {
263 match std::env::var(SCREENSHOT_BACKEND_ENV) {
264 Ok(value) if !value.trim().is_empty() => {
265 ScreenshotBackend::parse(&value).map(Some).ok_or_else(|| {
266 anyhow!(
267 "{SCREENSHOT_BACKEND_ENV}={value:?} is not a recognized backend \
268 (expected gnome-shell, portal, x11, or gnome-screenshot)"
269 )
270 })
271 }
272 _ => Ok(None),
273 }
274}
275
276pub async fn capture_screenshot() -> Result<ScreenshotCapture> {
277 let raw = capture_screenshot_raw().await?;
278 prepare_screenshot_payload(raw, ScreenshotPayloadOptions::default())
279}
280
281pub fn prepare_screenshot_payload(
282 raw: RawScreenshotCapture,
283 options: ScreenshotPayloadOptions,
284) -> Result<ScreenshotCapture> {
285 if raw.bytes.is_empty() {
286 bail!("screenshot file was empty");
287 }
288 let (coordinate_width, coordinate_height) = png_dimensions(&raw.bytes)?;
289 let original_bytes = raw.bytes.len();
290 let options = options.resolve();
291 let (target_width, target_height) =
292 target_dimensions(coordinate_width, coordinate_height, options);
293
294 let (bytes, width, height) = if options.format == ScreenshotOutputFormat::Png
295 && target_width == coordinate_width
296 && target_height == coordinate_height
297 && original_bytes <= options.max_bytes
298 {
299 (raw.bytes, coordinate_width, coordinate_height)
300 } else {
301 encode_screenshot_to_fit_bytes(
302 &raw.bytes,
303 coordinate_width,
304 coordinate_height,
305 target_width,
306 target_height,
307 options,
308 )?
309 };
310
311 let encoded = STANDARD.encode(&bytes);
312 let scale = if coordinate_width == 0 {
313 1.0
314 } else {
315 width as f32 / coordinate_width as f32
316 };
317
318 Ok(ScreenshotCapture {
319 mime_type: options.format.mime_type().to_string(),
320 data_url: format!("data:{};base64,{encoded}", options.format.mime_type()),
321 source: raw.source,
322 width,
323 height,
324 coordinate_width,
325 coordinate_height,
326 scale,
327 resized: width != coordinate_width || height != coordinate_height,
328 bytes: bytes.len(),
329 original_bytes,
330 max_bytes: options.max_bytes,
331 format: options.format,
332 quality: (options.format == ScreenshotOutputFormat::Jpeg).then_some(options.quality),
333 })
334}
335
336async fn capture_with_gnome_shell() -> Result<RawScreenshotCapture> {
337 let connection = zbus::Connection::session()
338 .await
339 .context("failed to connect to session bus")?;
340 let proxy = Proxy::new(
341 &connection,
342 "org.gnome.Shell.Screenshot",
343 "/org/gnome/Shell/Screenshot",
344 "org.gnome.Shell.Screenshot",
345 )
346 .await
347 .context("failed to create GNOME Shell screenshot proxy")?;
348 let path = temp_png_path("gnome-shell");
349 let filename = path
350 .to_str()
351 .context("temporary screenshot path is not valid UTF-8")?;
352 let result = proxy.call("Screenshot", &(false, false, filename)).await;
353 let (success, filename_used): (bool, String) = match result {
354 Ok(result) => result,
355 Err(error) => {
356 cleanup_gnome_requested_path(&path);
357 return Err(error).context("GNOME Shell Screenshot call failed");
358 }
359 };
360
361 if !success {
362 cleanup_gnome_requested_path(&path);
363 bail!("GNOME Shell reported screenshot failure");
364 }
365
366 read_png_as_capture(
367 PathBuf::from(filename_used),
368 "gnome-shell",
369 ScreenshotCleanup::DeletePath(path),
370 )
371 .await
372}
373
374async fn capture_with_portal() -> Result<RawScreenshotCapture> {
375 let connection = zbus::Connection::session()
376 .await
377 .context("failed to connect to session bus")?;
378 let portal_owner = portal_desktop_owner(&connection).await?;
379 let token = request_token();
380 let mut response_stream = portal_response_stream(&connection, &portal_owner).await?;
383
384 let portal_proxy = Proxy::new(
385 &connection,
386 portal_owner.clone(),
387 PORTAL_DESKTOP_PATH,
388 PORTAL_SCREENSHOT_INTERFACE,
389 )
390 .await
391 .context("failed to create XDG portal screenshot proxy")?;
392 let mut options: HashMap<&str, Value<'_>> = HashMap::new();
393 options.insert("handle_token", Value::from(token.as_str()));
394 options.insert("interactive", Value::from(false));
395 let handle: OwnedObjectPath = portal_proxy
396 .call("Screenshot", &("", options))
397 .await
398 .context("XDG portal Screenshot call failed")?;
399
400 let (response_code, results) = tokio::time::timeout(
401 Duration::from_secs(20),
402 wait_for_portal_response(&mut response_stream, handle.as_str(), portal_owner.as_str()),
403 )
404 .await
405 .context("timed out waiting for XDG portal screenshot response")??;
406
407 match response_code {
408 0 => {}
409 1 => bail!("XDG portal screenshot was cancelled by the user (response code 1)"),
410 2 => bail!(
411 "XDG portal screenshot ended without success (response code 2); \
412 check the desktop portal logs for permission or backend errors"
413 ),
414 _ => bail!("XDG portal screenshot returned unknown response code {response_code}"),
415 }
416
417 let uri_value = results
418 .get("uri")
419 .context("XDG portal screenshot response did not include a uri")?;
420 let uri: String = uri_value
421 .try_clone()
422 .context("failed to clone XDG portal screenshot uri")?
423 .try_into()
424 .context("XDG portal screenshot uri was not a string")?;
425 let path = file_uri_to_path(&uri)?;
426
427 read_png_as_capture(path, "xdg-desktop-portal", ScreenshotCleanup::Preserve).await
428}
429
430const GNOME_SCREENSHOT_TIMEOUT: Duration = Duration::from_secs(20);
433
434async fn capture_with_gnome_screenshot() -> Result<RawScreenshotCapture> {
435 let path = temp_png_path("gnome-screenshot");
436 let filename = path
437 .to_str()
438 .context("temporary screenshot path is not valid UTF-8")?;
439
440 let mut command = Command::new("gnome-screenshot");
444 command
445 .args(["-f", filename])
446 .stdin(Stdio::null())
447 .stdout(Stdio::null())
448 .stderr(Stdio::null());
449 let mut child = match crate::command_runner::spawn_retrying_busy(&mut command).await {
450 Ok(child) => child,
451 Err(error) => {
452 cleanup_gnome_requested_path(&path);
453 return Err(error).context("failed to spawn gnome-screenshot");
454 }
455 };
456
457 let status = match tokio::time::timeout(GNOME_SCREENSHOT_TIMEOUT, child.wait()).await {
460 Ok(Ok(status)) => status,
461 Ok(Err(error)) => {
462 cleanup_gnome_requested_path(&path);
463 return Err(error).context("failed to wait for gnome-screenshot");
464 }
465 Err(_) => {
466 let _ = child.kill().await;
467 cleanup_gnome_requested_path(&path);
468 bail!("gnome-screenshot timed out");
469 }
470 };
471
472 if !status.success() {
473 cleanup_gnome_requested_path(&path);
474 bail!("gnome-screenshot exited with {status}");
475 }
476
477 read_png_as_capture(
478 path.clone(),
479 "gnome-screenshot",
480 ScreenshotCleanup::DeletePath(path),
481 )
482 .await
483}
484
485async fn portal_desktop_owner(connection: &zbus::Connection) -> Result<OwnedUniqueName> {
486 let dbus = zbus::fdo::DBusProxy::new(connection)
487 .await
488 .context("failed to create session-bus identity proxy")?;
489 let _: u32 = dbus
490 .start_service_by_name(WellKnownName::try_from(PORTAL_DESKTOP_SERVICE)?, 0)
491 .await
492 .context("failed to activate the XDG desktop portal")?;
493 dbus.get_name_owner(BusName::try_from(PORTAL_DESKTOP_SERVICE)?)
494 .await
495 .context("failed to resolve the XDG desktop portal owner")
496}
497
498async fn portal_response_stream(
499 connection: &zbus::Connection,
500 portal_owner: &OwnedUniqueName,
501) -> Result<MessageStream> {
502 let response_rule = MatchRule::builder()
503 .msg_type(MessageType::Signal)
504 .sender(portal_owner.clone())?
505 .interface(PORTAL_REQUEST_INTERFACE)?
506 .member("Response")?
507 .path_namespace(PORTAL_REQUEST_PATH_NAMESPACE)?
508 .build();
509
510 MessageStream::for_match_rule(response_rule, connection, None)
511 .await
512 .context("failed to subscribe to XDG portal screenshot responses")
513}
514
515async fn wait_for_portal_response(
516 response_stream: &mut MessageStream,
517 request_path: &str,
518 portal_owner: &str,
519) -> Result<(u32, HashMap<String, OwnedValue>)> {
520 loop {
521 let response = response_stream
522 .next()
523 .await
524 .context("XDG portal screenshot response stream ended")?
525 .context("XDG portal screenshot response stream failed")?;
526
527 if !portal_response_matches(&response, request_path, portal_owner) {
528 continue;
529 }
530
531 return response
532 .body()
533 .deserialize()
534 .context("failed to decode XDG portal screenshot response");
535 }
536}
537
538fn portal_response_matches(response: &Message, request_path: &str, portal_owner: &str) -> bool {
539 let header = response.header();
540 header
541 .sender()
542 .is_some_and(|sender| sender.as_str() == portal_owner)
543 && header
544 .path()
545 .is_some_and(|path| path.as_str() == request_path)
546}
547
548async fn read_png_as_capture(
549 path: PathBuf,
550 source: &str,
551 cleanup: ScreenshotCleanup,
552) -> Result<RawScreenshotCapture> {
553 let result = read_png_as_capture_inner(&path, source);
554 if let ScreenshotCleanup::DeletePath(path) = cleanup {
555 let _ = fs::remove_file(path);
556 }
557 result
558}
559
560fn read_png_as_capture_inner(path: &Path, source: &str) -> Result<RawScreenshotCapture> {
561 let file = OpenOptions::new()
562 .read(true)
563 .custom_flags(libc::O_NOFOLLOW | libc::O_NONBLOCK)
564 .open(path)
565 .with_context(|| format!("failed to open screenshot file {}", path.display()))?;
566 let metadata = file
567 .metadata()
568 .with_context(|| format!("failed to inspect screenshot file {}", path.display()))?;
569 if !metadata.file_type().is_file() {
570 bail!("screenshot path is not a regular file: {}", path.display());
571 }
572 if metadata.len() > MAX_SCREENSHOT_SOURCE_BYTES {
573 bail!(
574 "screenshot file is {} bytes, over the {}-byte source limit",
575 metadata.len(),
576 MAX_SCREENSHOT_SOURCE_BYTES
577 );
578 }
579 let mut bytes = Vec::with_capacity(metadata.len() as usize);
580 file.take(MAX_SCREENSHOT_SOURCE_BYTES + 1)
581 .read_to_end(&mut bytes)
582 .with_context(|| format!("failed to read screenshot file {}", path.display()))?;
583 if bytes.len() as u64 > MAX_SCREENSHOT_SOURCE_BYTES {
584 bail!(
585 "screenshot file grew beyond the {}-byte source limit",
586 MAX_SCREENSHOT_SOURCE_BYTES
587 );
588 }
589 if bytes.is_empty() {
590 bail!("screenshot file was empty: {}", path.display());
591 }
592 let (width, height) = png_dimensions(&bytes)?;
593 Ok(RawScreenshotCapture {
594 mime_type: "image/png".to_string(),
595 bytes,
596 source: source.to_string(),
597 width,
598 height,
599 })
600}
601
602fn target_dimensions(
603 width: u32,
604 height: u32,
605 options: ResolvedScreenshotPayloadOptions,
606) -> (u32, u32) {
607 let width_scale = options.max_width as f64 / width as f64;
608 let height_scale = options.max_height as f64 / height as f64;
609 let scale = f64::from(options.scale)
610 .min(width_scale)
611 .min(height_scale)
612 .min(1.0);
613
614 let target_width = ((width as f64 * scale).round() as u32).clamp(1, width);
615 let target_height = ((height as f64 * scale).round() as u32).clamp(1, height);
616 (target_width, target_height)
617}
618
619fn encode_screenshot_to_fit_bytes(
620 raw: &[u8],
621 original_width: u32,
622 original_height: u32,
623 mut target_width: u32,
624 mut target_height: u32,
625 options: ResolvedScreenshotPayloadOptions,
626) -> Result<(Vec<u8>, u32, u32)> {
627 let img = image::load_from_memory_with_format(raw, image::ImageFormat::Png)
628 .context("failed to decode screenshot PNG for encoding")?;
629
630 loop {
631 let bytes = if options.format == ScreenshotOutputFormat::Png
632 && target_width == original_width
633 && target_height == original_height
634 {
635 raw.to_vec()
636 } else {
637 let output = if target_width == original_width && target_height == original_height {
638 img.clone()
639 } else {
640 img.resize_exact(target_width, target_height, FilterType::Lanczos3)
641 };
642 encode_image(&output, options)?
643 };
644
645 if bytes.len() <= options.max_bytes {
646 return Ok((bytes, target_width, target_height));
647 }
648
649 if target_width == 1 && target_height == 1 {
650 bail!(
651 "screenshot payload is {} bytes at 1x1, over max_bytes {}",
652 bytes.len(),
653 options.max_bytes
654 );
655 }
656
657 (target_width, target_height) = next_dimensions_for_byte_cap(
658 target_width,
659 target_height,
660 bytes.len(),
661 options.max_bytes,
662 );
663 }
664}
665
666fn encode_image(
667 img: &image::DynamicImage,
668 options: ResolvedScreenshotPayloadOptions,
669) -> Result<Vec<u8>> {
670 let mut out = Vec::new();
671 match options.format {
672 ScreenshotOutputFormat::Png => {
673 img.write_to(&mut Cursor::new(&mut out), image::ImageFormat::Png)
674 .context("failed to encode screenshot PNG")?;
675 }
676 ScreenshotOutputFormat::Jpeg => {
677 let rgb = img.to_rgb8();
678 JpegEncoder::new_with_quality(&mut out, options.quality)
679 .encode_image(&rgb)
680 .context("failed to encode screenshot JPEG")?;
681 }
682 }
683 Ok(out)
684}
685
686fn next_dimensions_for_byte_cap(
687 width: u32,
688 height: u32,
689 encoded_bytes: usize,
690 max_bytes: usize,
691) -> (u32, u32) {
692 let shrink = ((max_bytes as f64 / encoded_bytes as f64).sqrt() * 0.9).clamp(0.1, 0.95);
693 let mut next_width = ((width as f64 * shrink).floor() as u32).max(1);
694 let mut next_height = ((height as f64 * shrink).floor() as u32).max(1);
695
696 if next_width >= width && width > 1 {
697 next_width = width - 1;
698 }
699 if next_height >= height && height > 1 {
700 next_height = height - 1;
701 }
702
703 (next_width, next_height)
704}
705
706fn cleanup_gnome_requested_path(path: &Path) {
707 let _ = fs::remove_file(path);
708}
709
710fn png_dimensions(bytes: &[u8]) -> Result<(u32, u32)> {
711 const PNG_SIGNATURE: &[u8; 8] = b"\x89PNG\r\n\x1a\n";
712 if bytes.len() < 24 || &bytes[..8] != PNG_SIGNATURE || &bytes[12..16] != b"IHDR" {
713 bail!("screenshot file was not a valid PNG");
714 }
715 let width = u32::from_be_bytes(bytes[16..20].try_into().unwrap());
716 let height = u32::from_be_bytes(bytes[20..24].try_into().unwrap());
717 if width == 0 || height == 0 {
718 bail!("screenshot PNG had invalid dimensions {width}x{height}");
719 }
720 let pixels = u64::from(width) * u64::from(height);
721 if pixels > MAX_SCREENSHOT_SOURCE_PIXELS {
722 bail!(
723 "screenshot PNG has {pixels} pixels, over the {MAX_SCREENSHOT_SOURCE_PIXELS}-pixel source limit"
724 );
725 }
726 Ok((width, height))
727}
728
729fn file_uri_to_path(uri: &str) -> Result<PathBuf> {
730 let Some(rest) = uri.strip_prefix("file://") else {
731 bail!("unsupported screenshot uri: {uri}");
732 };
733 let path = PathBuf::from(percent_decode(rest));
734 if !path.is_absolute() {
735 bail!("screenshot uri did not contain an absolute file path: {uri}");
736 }
737 Ok(path)
738}
739
740fn percent_decode(value: &str) -> String {
741 let bytes = value.as_bytes();
742 let mut decoded = Vec::with_capacity(bytes.len());
743 let mut index = 0;
744
745 while index < bytes.len() {
746 if bytes[index] == b'%' && index + 2 < bytes.len() {
747 if let Ok(hex) = std::str::from_utf8(&bytes[index + 1..index + 3]) {
748 if let Ok(byte) = u8::from_str_radix(hex, 16) {
749 decoded.push(byte);
750 index += 3;
751 continue;
752 }
753 }
754 }
755
756 decoded.push(bytes[index]);
757 index += 1;
758 }
759
760 String::from_utf8_lossy(&decoded).into_owned()
761}
762
763fn temp_png_path(source: &str) -> PathBuf {
764 std::env::temp_dir().join(format!(
765 "computer-use-linux-{source}-{}.png",
766 unique_suffix()
767 ))
768}
769
770fn request_token() -> String {
771 format!("computer_use_linux_{}", unique_suffix().replace('-', "_"))
772}
773
774fn unique_suffix() -> String {
775 let nanos = SystemTime::now()
776 .duration_since(UNIX_EPOCH)
777 .map(|duration| duration.as_nanos())
778 .unwrap_or_default();
779 format!("{}-{nanos}", std::process::id())
780}
781
782#[cfg(test)]
783mod tests {
784 use super::*;
785
786 fn test_path(name: &str) -> PathBuf {
787 std::env::temp_dir().join(format!(
788 "computer-use-linux-screenshot-test-{name}-{}",
789 unique_suffix()
790 ))
791 }
792
793 fn valid_png(width: u32, height: u32) -> Vec<u8> {
794 let mut png = Vec::new();
795 png.extend_from_slice(b"\x89PNG\r\n\x1a\n");
796 png.extend_from_slice(&13_u32.to_be_bytes());
797 png.extend_from_slice(b"IHDR");
798 png.extend_from_slice(&width.to_be_bytes());
799 png.extend_from_slice(&height.to_be_bytes());
800 png.extend_from_slice(&[8, 6, 0, 0, 0]);
801 png
802 }
803
804 fn solid_png(width: u32, height: u32) -> Vec<u8> {
805 let img = image::RgbaImage::from_pixel(width, height, image::Rgba([24, 96, 160, 255]));
806 encode_test_png(img)
807 }
808
809 fn noisy_png(width: u32, height: u32) -> Vec<u8> {
810 let mut img = image::RgbaImage::new(width, height);
811 for (x, y, pixel) in img.enumerate_pixels_mut() {
812 let r = ((x * 31 + y * 17) % 256) as u8;
813 let g = ((x * 13 + y * 47) % 256) as u8;
814 let b = ((x * 97 + y * 7) % 256) as u8;
815 *pixel = image::Rgba([r, g, b, 255]);
816 }
817 encode_test_png(img)
818 }
819
820 fn encode_test_png(img: image::RgbaImage) -> Vec<u8> {
821 let mut out = Vec::new();
822 image::DynamicImage::ImageRgba8(img)
823 .write_to(&mut Cursor::new(&mut out), image::ImageFormat::Png)
824 .unwrap();
825 out
826 }
827
828 fn raw_capture(bytes: Vec<u8>) -> RawScreenshotCapture {
829 let (width, height) = png_dimensions(&bytes).unwrap();
830 RawScreenshotCapture {
831 mime_type: "image/png".to_string(),
832 bytes,
833 source: "test".to_string(),
834 width,
835 height,
836 }
837 }
838
839 #[test]
840 fn decodes_file_uri_percent_escapes() {
841 assert_eq!(
842 file_uri_to_path("file:///tmp/Codex%20Screenshot.png").unwrap(),
843 PathBuf::from("/tmp/Codex Screenshot.png")
844 );
845 }
846
847 #[test]
848 fn file_uri_requires_an_absolute_path() {
849 assert!(file_uri_to_path("file://relative/path.png").is_err());
850 }
851
852 #[test]
853 fn portal_response_requires_the_pinned_sender_and_exact_path() {
854 let request_path = "/org/freedesktop/portal/desktop/request/1_7/token";
855 let response = Message::signal(request_path, PORTAL_REQUEST_INTERFACE, "Response")
856 .unwrap()
857 .sender(":1.42")
858 .unwrap()
859 .build(&())
860 .unwrap();
861
862 assert!(portal_response_matches(&response, request_path, ":1.42"));
863 assert!(!portal_response_matches(&response, request_path, ":1.99"));
864 assert!(!portal_response_matches(
865 &response,
866 "/org/freedesktop/portal/desktop/request/1_7/other",
867 ":1.42"
868 ));
869 }
870
871 #[test]
872 fn parses_known_backend_names() {
873 assert_eq!(
874 ScreenshotBackend::parse("gnome-shell"),
875 Some(ScreenshotBackend::GnomeShell)
876 );
877 assert_eq!(
878 ScreenshotBackend::parse(" Portal "),
879 Some(ScreenshotBackend::Portal)
880 );
881 assert_eq!(
882 ScreenshotBackend::parse("GNOME_SCREENSHOT"),
883 Some(ScreenshotBackend::GnomeScreenshot)
884 );
885 assert_eq!(
886 ScreenshotBackend::parse("x11"),
887 Some(ScreenshotBackend::X11)
888 );
889 assert_eq!(
890 ScreenshotBackend::parse(" X11-Native "),
891 Some(ScreenshotBackend::X11)
892 );
893 assert_eq!(ScreenshotBackend::parse("nonsense"), None);
894 }
895
896 #[test]
897 fn forced_backend_reads_env_override() {
898 std::env::set_var(SCREENSHOT_BACKEND_ENV, "gnome-screenshot");
900 assert_eq!(
901 forced_backend().unwrap(),
902 Some(ScreenshotBackend::GnomeScreenshot)
903 );
904
905 std::env::set_var(SCREENSHOT_BACKEND_ENV, " ");
906 assert_eq!(forced_backend().unwrap(), None);
907
908 std::env::set_var(SCREENSHOT_BACKEND_ENV, "bogus");
909 let error = forced_backend().unwrap_err();
910 assert!(error.to_string().contains("not a recognized backend"));
911
912 std::env::remove_var(SCREENSHOT_BACKEND_ENV);
913 assert_eq!(forced_backend().unwrap(), None);
914 }
915
916 #[test]
917 fn request_token_is_portal_safe() {
918 let token = request_token();
919 assert!(token.starts_with("computer_use_linux_"));
920 assert!(token.chars().all(|c| c.is_ascii_alphanumeric() || c == '_'));
921 }
922
923 #[test]
924 fn reads_png_dimensions_from_ihdr() {
925 let png = valid_png(3840, 1080);
926
927 assert_eq!(png_dimensions(&png).unwrap(), (3840, 1080));
928 }
929
930 #[test]
931 fn default_payload_downscales_long_edge() {
932 let capture =
933 prepare_screenshot_payload(raw_capture(solid_png(4000, 1000)), Default::default())
934 .unwrap();
935
936 assert_eq!((capture.width, capture.height), (1920, 480));
937 assert_eq!(
938 (capture.coordinate_width, capture.coordinate_height),
939 (4000, 1000)
940 );
941 assert!(capture.resized);
942 assert!(capture.bytes <= DEFAULT_SCREENSHOT_MAX_BYTES);
943 assert!(capture.data_url.starts_with("data:image/png;base64,"));
944 }
945
946 #[test]
947 fn larger_bounded_request_can_keep_more_detail() {
948 let capture = prepare_screenshot_payload(
949 raw_capture(solid_png(3000, 1000)),
950 ScreenshotPayloadOptions {
951 max_width: Some(3000),
952 max_height: Some(3000),
953 max_bytes: Some(DEFAULT_SCREENSHOT_MAX_BYTES),
954 ..Default::default()
955 },
956 )
957 .unwrap();
958
959 assert_eq!((capture.width, capture.height), (3000, 1000));
960 assert_eq!(
961 (capture.coordinate_width, capture.coordinate_height),
962 (3000, 1000)
963 );
964 assert!(!capture.resized);
965 }
966
967 #[test]
968 fn byte_cap_downscales_until_payload_fits() {
969 let capture = prepare_screenshot_payload(
970 raw_capture(noisy_png(512, 512)),
971 ScreenshotPayloadOptions {
972 max_width: Some(512),
973 max_height: Some(512),
974 max_bytes: Some(20_000),
975 ..Default::default()
976 },
977 )
978 .unwrap();
979
980 assert!(capture.bytes <= 20_000);
981 assert!(capture.width < 512);
982 assert_eq!(
983 (capture.coordinate_width, capture.coordinate_height),
984 (512, 512)
985 );
986 assert!(capture.resized);
987 }
988
989 #[test]
990 fn jpeg_format_compresses_when_requested() {
991 let capture = prepare_screenshot_payload(
992 raw_capture(noisy_png(512, 512)),
993 ScreenshotPayloadOptions {
994 max_width: Some(512),
995 max_height: Some(512),
996 max_bytes: Some(DEFAULT_SCREENSHOT_MAX_BYTES),
997 format: Some(ScreenshotOutputFormat::Jpeg),
998 quality: Some(60),
999 ..Default::default()
1000 },
1001 )
1002 .unwrap();
1003
1004 assert_eq!(capture.mime_type, "image/jpeg");
1005 assert_eq!(capture.format, ScreenshotOutputFormat::Jpeg);
1006 assert_eq!(capture.quality, Some(60));
1007 assert_eq!((capture.width, capture.height), (512, 512));
1008 assert_eq!(
1009 (capture.coordinate_width, capture.coordinate_height),
1010 (512, 512)
1011 );
1012 assert!(capture.bytes < capture.original_bytes);
1013 assert!(capture.data_url.starts_with("data:image/jpeg;base64,"));
1014 }
1015
1016 #[tokio::test]
1017 async fn portal_capture_preserves_valid_returned_path() {
1018 let path = test_path("portal-valid");
1019 fs::write(&path, valid_png(1, 1)).unwrap();
1020
1021 let capture = read_png_as_capture(
1022 path.clone(),
1023 "xdg-desktop-portal",
1024 ScreenshotCleanup::Preserve,
1025 )
1026 .await
1027 .unwrap();
1028
1029 assert_eq!(capture.source, "xdg-desktop-portal");
1030 assert!(path.exists());
1031 let _ = fs::remove_file(path);
1032 }
1033
1034 #[tokio::test]
1035 async fn portal_capture_preserves_invalid_returned_path() {
1036 let path = test_path("portal-invalid");
1037 fs::write(&path, b"").unwrap();
1038
1039 let error = read_png_as_capture(
1040 path.clone(),
1041 "xdg-desktop-portal",
1042 ScreenshotCleanup::Preserve,
1043 )
1044 .await
1045 .unwrap_err();
1046
1047 assert!(error.to_string().contains("screenshot file was empty"));
1048 assert!(path.exists());
1049 let _ = fs::remove_file(path);
1050 }
1051
1052 #[tokio::test]
1053 async fn screenshot_reader_rejects_symlinks() {
1054 let target = test_path("symlink-target");
1055 let link = test_path("symlink-link");
1056 fs::write(&target, valid_png(1, 1)).unwrap();
1057 std::os::unix::fs::symlink(&target, &link).unwrap();
1058
1059 let error = read_png_as_capture(link.clone(), "test", ScreenshotCleanup::Preserve)
1060 .await
1061 .unwrap_err();
1062
1063 assert!(error.to_string().contains("failed to open screenshot file"));
1064 let _ = fs::remove_file(link);
1065 let _ = fs::remove_file(target);
1066 }
1067
1068 #[tokio::test]
1069 async fn screenshot_reader_rejects_oversized_source_files() {
1070 let path = test_path("oversized-source");
1071 let file = fs::File::create(&path).unwrap();
1072 file.set_len(MAX_SCREENSHOT_SOURCE_BYTES + 1).unwrap();
1073
1074 let error = read_png_as_capture(path.clone(), "test", ScreenshotCleanup::Preserve)
1075 .await
1076 .unwrap_err();
1077
1078 assert!(error.to_string().contains("over the"));
1079 let _ = fs::remove_file(path);
1080 }
1081
1082 #[test]
1083 fn png_dimensions_reject_decompression_bomb_dimensions() {
1084 let png = valid_png(u32::MAX, u32::MAX);
1085
1086 assert!(png_dimensions(&png)
1087 .unwrap_err()
1088 .to_string()
1089 .contains("pixel source limit"));
1090 }
1091
1092 #[tokio::test]
1093 async fn gnome_capture_deletes_backend_temp_path_on_success() {
1094 let path = test_path("gnome-valid");
1095 fs::write(&path, valid_png(1, 1)).unwrap();
1096
1097 let capture = read_png_as_capture(
1098 path.clone(),
1099 "gnome-shell",
1100 ScreenshotCleanup::DeletePath(path.clone()),
1101 )
1102 .await
1103 .unwrap();
1104
1105 assert_eq!(capture.source, "gnome-shell");
1106 assert!(!path.exists());
1107 }
1108
1109 #[tokio::test]
1110 async fn gnome_capture_deletes_backend_temp_path_on_parse_failure() {
1111 let path = test_path("gnome-invalid");
1112 fs::write(&path, b"").unwrap();
1113
1114 let error = read_png_as_capture(
1115 path.clone(),
1116 "gnome-shell",
1117 ScreenshotCleanup::DeletePath(path.clone()),
1118 )
1119 .await
1120 .unwrap_err();
1121
1122 assert!(error.to_string().contains("screenshot file was empty"));
1123 assert!(!path.exists());
1124 }
1125
1126 #[test]
1127 fn gnome_failure_cleanup_removes_requested_temp_path() {
1128 let path = test_path("gnome-pre-read-failure");
1129 fs::write(&path, b"partial").unwrap();
1130
1131 cleanup_gnome_requested_path(&path);
1132
1133 assert!(!path.exists());
1134 }
1135
1136 #[tokio::test]
1137 async fn gnome_deletes_requested_temp_path_and_preserves_unexpected_returned_path() {
1138 let requested = test_path("gnome-requested");
1139 let returned = test_path("gnome-returned");
1140 fs::write(&requested, b"partial").unwrap();
1141 fs::write(&returned, valid_png(1, 1)).unwrap();
1142
1143 let capture = read_png_as_capture(
1144 returned.clone(),
1145 "gnome-shell",
1146 ScreenshotCleanup::DeletePath(requested.clone()),
1147 )
1148 .await
1149 .unwrap();
1150
1151 assert_eq!(capture.source, "gnome-shell");
1152 assert!(!requested.exists());
1153 assert!(returned.exists());
1154 let _ = fs::remove_file(returned);
1155 }
1156}