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, HashSet},
12 fs::{self, File, OpenOptions},
13 io::{Cursor, Read},
14 os::unix::fs::{MetadataExt, 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)]
113struct PortalScreenshotRequest {
114 requested_at: SystemTime,
115 pictures: Option<PortalPicturesSnapshot>,
116}
117
118#[derive(Debug)]
119struct PortalPicturesSnapshot {
120 path: PathBuf,
121 directory: File,
122 names: HashSet<std::ffi::OsString>,
123}
124
125impl PortalPicturesSnapshot {
126 fn new(path: PathBuf) -> std::io::Result<Self> {
127 let path = fs::canonicalize(path)?;
128 let directory = File::open(&path)?;
129 let names = fs::read_dir(&path)?
130 .map(|entry| entry.map(|entry| entry.file_name()))
131 .collect::<std::io::Result<HashSet<_>>>()?;
132 Ok(Self {
133 path,
134 directory,
135 names,
136 })
137 }
138
139 fn contains_new_path(&self, path: &Path) -> bool {
140 if path.parent() != Some(self.path.as_path()) {
141 return false;
142 }
143 let same_directory = self
144 .directory
145 .metadata()
146 .ok()
147 .zip(fs::metadata(&self.path).ok())
148 .is_some_and(|(before, current)| {
149 (before.dev(), before.ino()) == (current.dev(), current.ino())
150 });
151 same_directory
152 && path
153 .file_name()
154 .is_some_and(|name| !self.names.contains(name))
155 }
156}
157
158#[derive(Debug)]
159enum ScreenshotCleanup {
160 DeletePath(PathBuf),
161 #[cfg(test)]
162 Preserve,
163 PortalPath {
164 path: PathBuf,
165 request: PortalScreenshotRequest,
166 },
167}
168
169struct OwnedScreenshotCleanup {
170 strategy: Option<ScreenshotCleanup>,
171 portal_file: Option<File>,
172}
173
174impl OwnedScreenshotCleanup {
175 fn new(strategy: ScreenshotCleanup) -> Self {
176 let portal_file = match &strategy {
177 ScreenshotCleanup::PortalPath { path, request } => OpenOptions::new()
178 .read(true)
179 .custom_flags(libc::O_NOFOLLOW | libc::O_NONBLOCK)
180 .open(path)
181 .ok()
182 .filter(|file| {
183 file.metadata().is_ok_and(|metadata| {
184 if !metadata.is_file() || metadata.uid() != unsafe { libc::geteuid() } {
185 return false;
186 }
187 if let Some(pictures) = &request.pictures {
188 if path.parent() == Some(pictures.path.as_path()) {
189 return pictures.contains_new_path(path);
190 }
191 }
192 metadata
195 .created()
196 .is_ok_and(|created| created >= request.requested_at)
197 })
198 }),
199 _ => None,
200 };
201 Self {
202 strategy: Some(strategy),
203 portal_file,
204 }
205 }
206}
207
208impl OwnedScreenshotCleanup {
209 fn cleanup_now(&mut self) {
210 if let Some(strategy) = self.strategy.take() {
211 cleanup_screenshot_path(strategy, self.portal_file.take());
212 }
213 }
214}
215
216impl Drop for OwnedScreenshotCleanup {
217 fn drop(&mut self) {
218 let Some(strategy) = self.strategy.take() else {
219 return;
220 };
221 let portal_file = self.portal_file.take();
222 if matches!(strategy, ScreenshotCleanup::PortalPath { .. }) {
223 if let Ok(runtime) = tokio::runtime::Handle::try_current() {
224 runtime.spawn_blocking(move || cleanup_screenshot_path(strategy, portal_file));
227 return;
228 }
229 }
230 cleanup_screenshot_path(strategy, portal_file);
231 }
232}
233
234fn cleanup_screenshot_path(strategy: ScreenshotCleanup, portal_file: Option<File>) {
235 let path = match &strategy {
236 ScreenshotCleanup::DeletePath(path) => Some(path),
237 ScreenshotCleanup::PortalPath { path, .. } => portal_file.as_ref().and_then(|file| {
238 file.metadata()
239 .ok()
240 .zip(fs::symlink_metadata(path).ok())
241 .filter(|(original, current)| {
242 current.file_type().is_file()
243 && (original.dev(), original.ino()) == (current.dev(), current.ino())
244 })
245 .map(|_| path)
246 }),
247 #[cfg(test)]
248 ScreenshotCleanup::Preserve => None,
249 };
250 if let Some(path) = path {
251 if let Err(error) = fs::remove_file(path) {
252 if error.kind() != std::io::ErrorKind::NotFound {
253 eprintln!("Failed to remove screenshot capture file: {error}");
254 }
255 }
256 }
257}
258
259impl ScreenshotPayloadOptions {
260 fn resolve(self) -> ResolvedScreenshotPayloadOptions {
261 let max_width = self
262 .max_width
263 .unwrap_or(DEFAULT_SCREENSHOT_MAX_DIMENSION)
264 .clamp(1, ABSOLUTE_SCREENSHOT_MAX_DIMENSION);
265 let max_height = self
266 .max_height
267 .unwrap_or(DEFAULT_SCREENSHOT_MAX_DIMENSION)
268 .clamp(1, ABSOLUTE_SCREENSHOT_MAX_DIMENSION);
269 let max_bytes = self
270 .max_bytes
271 .unwrap_or(DEFAULT_SCREENSHOT_MAX_BYTES)
272 .clamp(MIN_SCREENSHOT_MAX_BYTES, ABSOLUTE_SCREENSHOT_MAX_BYTES);
273 let scale = self
274 .scale
275 .filter(|value| value.is_finite() && *value > 0.0)
276 .unwrap_or(1.0)
277 .min(1.0);
278 let format = self.format.unwrap_or(ScreenshotOutputFormat::Png);
279 let quality = self
280 .quality
281 .unwrap_or(DEFAULT_SCREENSHOT_JPEG_QUALITY)
282 .clamp(MIN_SCREENSHOT_JPEG_QUALITY, MAX_SCREENSHOT_JPEG_QUALITY);
283
284 ResolvedScreenshotPayloadOptions {
285 max_width,
286 max_height,
287 max_bytes,
288 scale,
289 format,
290 quality,
291 }
292 }
293}
294
295const SCREENSHOT_BACKEND_ENV: &str = "COMPUTER_USE_LINUX_SCREENSHOT_BACKEND";
298
299#[derive(Debug, Clone, Copy, PartialEq, Eq)]
300enum ScreenshotBackend {
301 GnomeShell,
302 Portal,
303 X11,
304 GnomeScreenshot,
305}
306
307impl ScreenshotBackend {
308 fn parse(value: &str) -> Option<Self> {
309 match value.trim().to_ascii_lowercase().as_str() {
310 "gnome-shell" | "gnome_shell" | "shell" => Some(Self::GnomeShell),
311 "portal" | "xdg-portal" | "xdg_portal" => Some(Self::Portal),
312 "x11" | "x11-native" | "x11_native" | "xgetimage" => Some(Self::X11),
313 "gnome-screenshot" | "gnome_screenshot" => Some(Self::GnomeScreenshot),
314 _ => None,
315 }
316 }
317
318 async fn capture(self) -> Result<RawScreenshotCapture> {
319 match self {
320 Self::GnomeShell => capture_with_gnome_shell().await,
321 Self::Portal => capture_with_portal().await,
322 Self::X11 => capture_with_x11().await,
323 Self::GnomeScreenshot => capture_with_gnome_screenshot().await,
324 }
325 }
326}
327
328pub async fn capture_screenshot_raw() -> Result<RawScreenshotCapture> {
329 hydrate_session_bus_env();
330
331 if let Some(forced) = forced_backend()? {
335 return forced.capture().await;
336 }
337
338 let gnome_error = match capture_with_gnome_shell().await {
345 Ok(capture) => return Ok(capture),
346 Err(error) => error,
347 };
348 let portal_error = match capture_with_portal().await {
349 Ok(capture) => return Ok(capture),
350 Err(error) => error,
351 };
352 let x11_error = match capture_with_x11().await {
356 Ok(capture) => return Ok(capture),
357 Err(error) => error,
358 };
359 let cli_error = match capture_with_gnome_screenshot().await {
360 Ok(capture) => return Ok(capture),
361 Err(error) => error,
362 };
363
364 Err(anyhow!(
365 "GNOME Shell screenshot failed: {gnome_error:#}; \
366 XDG portal screenshot failed: {portal_error:#}; \
367 native X11 screenshot failed: {x11_error:#}; \
368 gnome-screenshot fallback failed: {cli_error:#}"
369 ))
370}
371
372async fn capture_with_x11() -> Result<RawScreenshotCapture> {
375 if !is_native_x11_session() {
376 bail!("not a native X11 session (needs DISPLAY on an X11, not Wayland, session)");
377 }
378 let image = with_x11_display(X11_CAPTURE_TIMEOUT, |display| display.capture_root()).await??;
379 let (width, height) = (image.width, image.height);
380 let bytes =
381 tokio::task::spawn_blocking(move || encode_rgb_png(image.width, image.height, image.rgb))
382 .await
383 .context("X11 screenshot encoder task failed")??;
384 Ok(RawScreenshotCapture {
385 mime_type: "image/png".to_string(),
386 bytes,
387 source: "x11".to_string(),
388 width,
389 height,
390 })
391}
392
393fn encode_rgb_png(width: u32, height: u32, rgb: Vec<u8>) -> Result<Vec<u8>> {
394 let buffer = image::RgbImage::from_raw(width, height, rgb)
395 .context("X11 root image did not match its dimensions")?;
396 let mut out = Vec::new();
397 image::DynamicImage::ImageRgb8(buffer)
398 .write_to(&mut Cursor::new(&mut out), image::ImageFormat::Png)
399 .context("failed to encode X11 screenshot PNG")?;
400 Ok(out)
401}
402
403fn forced_backend() -> Result<Option<ScreenshotBackend>> {
404 match std::env::var(SCREENSHOT_BACKEND_ENV) {
405 Ok(value) if !value.trim().is_empty() => {
406 ScreenshotBackend::parse(&value).map(Some).ok_or_else(|| {
407 anyhow!(
408 "{SCREENSHOT_BACKEND_ENV}={value:?} is not a recognized backend \
409 (expected gnome-shell, portal, x11, or gnome-screenshot)"
410 )
411 })
412 }
413 _ => Ok(None),
414 }
415}
416
417pub async fn capture_screenshot() -> Result<ScreenshotCapture> {
418 let _pipeline = screenshot_pipeline_permit().await?;
419 let raw = capture_screenshot_raw().await?;
420 run_image_task(move || prepare_screenshot_payload(raw, ScreenshotPayloadOptions::default()))
421 .await
422}
423
424pub(crate) async fn screenshot_pipeline_permit() -> Result<tokio::sync::SemaphorePermit<'static>> {
425 static CAPTURE_PIPELINE: tokio::sync::Semaphore = tokio::sync::Semaphore::const_new(1);
426 Ok(CAPTURE_PIPELINE.acquire().await?)
427}
428
429pub(crate) async fn run_image_task<T: Send + 'static>(
430 work: impl FnOnce() -> Result<T> + Send + 'static,
431) -> Result<T> {
432 static IMAGE_WORKER: tokio::sync::Semaphore = tokio::sync::Semaphore::const_new(1);
433 let permit = IMAGE_WORKER.acquire().await?;
434 tokio::task::spawn_blocking(move || {
435 let _permit = permit;
436 work()
437 })
438 .await
439 .context("screenshot processing task failed")?
440}
441
442pub fn prepare_screenshot_payload(
443 raw: RawScreenshotCapture,
444 options: ScreenshotPayloadOptions,
445) -> Result<ScreenshotCapture> {
446 if raw.bytes.is_empty() {
447 bail!("screenshot file was empty");
448 }
449 let (coordinate_width, coordinate_height) = png_dimensions(&raw.bytes)?;
450 let original_bytes = raw.bytes.len();
451 let options = options.resolve();
452 let (target_width, target_height) =
453 target_dimensions(coordinate_width, coordinate_height, options);
454
455 let (bytes, width, height) = if options.format == ScreenshotOutputFormat::Png
456 && target_width == coordinate_width
457 && target_height == coordinate_height
458 && original_bytes <= options.max_bytes
459 {
460 (raw.bytes, coordinate_width, coordinate_height)
461 } else {
462 encode_screenshot_to_fit_bytes(
463 &raw.bytes,
464 coordinate_width,
465 coordinate_height,
466 target_width,
467 target_height,
468 options,
469 )?
470 };
471
472 let encoded = STANDARD.encode(&bytes);
473 let scale = if coordinate_width == 0 {
474 1.0
475 } else {
476 width as f32 / coordinate_width as f32
477 };
478
479 Ok(ScreenshotCapture {
480 mime_type: options.format.mime_type().to_string(),
481 data_url: format!("data:{};base64,{encoded}", options.format.mime_type()),
482 source: raw.source,
483 width,
484 height,
485 coordinate_width,
486 coordinate_height,
487 scale,
488 resized: width != coordinate_width || height != coordinate_height,
489 bytes: bytes.len(),
490 original_bytes,
491 max_bytes: options.max_bytes,
492 format: options.format,
493 quality: (options.format == ScreenshotOutputFormat::Jpeg).then_some(options.quality),
494 })
495}
496
497async fn capture_with_gnome_shell() -> Result<RawScreenshotCapture> {
498 let connection = zbus::Connection::session()
499 .await
500 .context("failed to connect to session bus")?;
501 let proxy = Proxy::new(
502 &connection,
503 "org.gnome.Shell.Screenshot",
504 "/org/gnome/Shell/Screenshot",
505 "org.gnome.Shell.Screenshot",
506 )
507 .await
508 .context("failed to create GNOME Shell screenshot proxy")?;
509 let path = temp_png_path("gnome-shell");
510 let filename = path
511 .to_str()
512 .context("temporary screenshot path is not valid UTF-8")?;
513 let result = proxy.call("Screenshot", &(false, false, filename)).await;
514 let (success, filename_used): (bool, String) = match result {
515 Ok(result) => result,
516 Err(error) => {
517 cleanup_gnome_requested_path(&path);
518 return Err(error).context("GNOME Shell Screenshot call failed");
519 }
520 };
521
522 if !success {
523 cleanup_gnome_requested_path(&path);
524 bail!("GNOME Shell reported screenshot failure");
525 }
526
527 read_png_as_capture(
528 PathBuf::from(filename_used),
529 "gnome-shell",
530 ScreenshotCleanup::DeletePath(path),
531 )
532 .await
533}
534
535async fn capture_with_portal() -> Result<RawScreenshotCapture> {
536 let connection = zbus::Connection::session()
537 .await
538 .context("failed to connect to session bus")?;
539 let portal_owner = portal_desktop_owner(&connection).await?;
540 let token = request_token();
541 let mut response_stream = portal_response_stream(&connection, &portal_owner).await?;
544
545 let portal_proxy = Proxy::new(
546 &connection,
547 portal_owner.clone(),
548 PORTAL_DESKTOP_PATH,
549 PORTAL_SCREENSHOT_INTERFACE,
550 )
551 .await
552 .context("failed to create XDG portal screenshot proxy")?;
553 let mut options: HashMap<&str, Value<'_>> = HashMap::new();
554 options.insert("handle_token", Value::from(token.as_str()));
555 options.insert("interactive", Value::from(false));
556 let mut pictures_command = Command::new("xdg-user-dir");
557 pictures_command.arg("PICTURES");
558 let pictures = crate::command_runner::output(pictures_command, "locate Pictures directory")
559 .await
560 .ok()
561 .filter(|output| output.status.success())
562 .and_then(|output| String::from_utf8(output.stdout).ok())
563 .map(|path| PathBuf::from(path.trim_end_matches('\n')))
564 .filter(|path| path.is_absolute());
565 let pictures = match pictures {
566 Some(path) => match tokio::time::timeout(
569 Duration::from_secs(1),
570 tokio::task::spawn_blocking(move || PortalPicturesSnapshot::new(path)),
571 )
572 .await
573 {
574 Ok(Ok(Ok(snapshot))) => Some(snapshot),
575 _ => None,
576 },
577 None => None,
578 };
579 let request = PortalScreenshotRequest {
580 requested_at: SystemTime::now(),
581 pictures,
582 };
583 let handle: OwnedObjectPath = portal_proxy
584 .call("Screenshot", &("", options))
585 .await
586 .context("XDG portal Screenshot call failed")?;
587
588 let (response_code, results) = tokio::time::timeout(
589 Duration::from_secs(20),
590 wait_for_portal_response(&mut response_stream, handle.as_str(), portal_owner.as_str()),
591 )
592 .await
593 .context("timed out waiting for XDG portal screenshot response")??;
594
595 match response_code {
596 0 => {}
597 1 => bail!("XDG portal screenshot was cancelled by the user (response code 1)"),
598 2 => bail!(
599 "XDG portal screenshot ended without success (response code 2); \
600 check the desktop portal logs for permission or backend errors"
601 ),
602 _ => bail!("XDG portal screenshot returned unknown response code {response_code}"),
603 }
604
605 let uri_value = results
606 .get("uri")
607 .context("XDG portal screenshot response did not include a uri")?;
608 let uri: String = uri_value
609 .try_clone()
610 .context("failed to clone XDG portal screenshot uri")?
611 .try_into()
612 .context("XDG portal screenshot uri was not a string")?;
613 let path = file_uri_to_path(&uri)?;
614
615 read_portal_png_as_capture(path, request).await
616}
617
618async fn read_portal_png_as_capture(
619 path: PathBuf,
620 request: PortalScreenshotRequest,
621) -> Result<RawScreenshotCapture> {
622 let cleanup = ScreenshotCleanup::PortalPath {
625 path: path.clone(),
626 request,
627 };
628 read_png_as_capture(path, "xdg-desktop-portal", cleanup).await
629}
630
631const GNOME_SCREENSHOT_TIMEOUT: Duration = Duration::from_secs(20);
634
635async fn capture_with_gnome_screenshot() -> Result<RawScreenshotCapture> {
636 let path = temp_png_path("gnome-screenshot");
637 let filename = path
638 .to_str()
639 .context("temporary screenshot path is not valid UTF-8")?;
640
641 let mut command = Command::new("gnome-screenshot");
645 command
646 .args(["-f", filename])
647 .stdin(Stdio::null())
648 .stdout(Stdio::null())
649 .stderr(Stdio::null());
650 let mut child = match crate::command_runner::spawn_retrying_busy(&mut command).await {
651 Ok(child) => child,
652 Err(error) => {
653 cleanup_gnome_requested_path(&path);
654 return Err(error).context("failed to spawn gnome-screenshot");
655 }
656 };
657
658 let status = match tokio::time::timeout(GNOME_SCREENSHOT_TIMEOUT, child.wait()).await {
661 Ok(Ok(status)) => status,
662 Ok(Err(error)) => {
663 cleanup_gnome_requested_path(&path);
664 return Err(error).context("failed to wait for gnome-screenshot");
665 }
666 Err(_) => {
667 let _ = child.kill().await;
668 cleanup_gnome_requested_path(&path);
669 bail!("gnome-screenshot timed out");
670 }
671 };
672
673 if !status.success() {
674 cleanup_gnome_requested_path(&path);
675 bail!("gnome-screenshot exited with {status}");
676 }
677
678 read_png_as_capture(
679 path.clone(),
680 "gnome-screenshot",
681 ScreenshotCleanup::DeletePath(path),
682 )
683 .await
684}
685
686async fn portal_desktop_owner(connection: &zbus::Connection) -> Result<OwnedUniqueName> {
687 let dbus = zbus::fdo::DBusProxy::new(connection)
688 .await
689 .context("failed to create session-bus identity proxy")?;
690 match dbus
693 .get_name_owner(BusName::try_from(PORTAL_DESKTOP_SERVICE)?)
694 .await
695 {
696 Ok(owner) => return Ok(owner),
697 Err(zbus::fdo::Error::NameHasNoOwner(_)) => {}
698 Err(error) => return Err(error).context("failed to resolve the XDG desktop portal owner"),
699 }
700 let _: u32 = dbus
701 .start_service_by_name(WellKnownName::try_from(PORTAL_DESKTOP_SERVICE)?, 0)
702 .await
703 .context("failed to activate the XDG desktop portal")?;
704 dbus.get_name_owner(BusName::try_from(PORTAL_DESKTOP_SERVICE)?)
705 .await
706 .context("failed to resolve the XDG desktop portal owner")
707}
708
709async fn portal_response_stream(
710 connection: &zbus::Connection,
711 portal_owner: &OwnedUniqueName,
712) -> Result<MessageStream> {
713 let response_rule = MatchRule::builder()
714 .msg_type(MessageType::Signal)
715 .sender(portal_owner.clone())?
716 .interface(PORTAL_REQUEST_INTERFACE)?
717 .member("Response")?
718 .path_namespace(PORTAL_REQUEST_PATH_NAMESPACE)?
719 .build();
720
721 MessageStream::for_match_rule(response_rule, connection, None)
722 .await
723 .context("failed to subscribe to XDG portal screenshot responses")
724}
725
726async fn wait_for_portal_response(
727 response_stream: &mut MessageStream,
728 request_path: &str,
729 portal_owner: &str,
730) -> Result<(u32, HashMap<String, OwnedValue>)> {
731 loop {
732 let response = response_stream
733 .next()
734 .await
735 .context("XDG portal screenshot response stream ended")?
736 .context("XDG portal screenshot response stream failed")?;
737
738 if !portal_response_matches(&response, request_path, portal_owner) {
739 continue;
740 }
741
742 return response
743 .body()
744 .deserialize()
745 .context("failed to decode XDG portal screenshot response");
746 }
747}
748
749fn portal_response_matches(response: &Message, request_path: &str, portal_owner: &str) -> bool {
750 let header = response.header();
751 header
752 .sender()
753 .is_some_and(|sender| sender.as_str() == portal_owner)
754 && header
755 .path()
756 .is_some_and(|path| path.as_str() == request_path)
757}
758
759async fn read_png_as_capture(
760 path: PathBuf,
761 source: &str,
762 cleanup: ScreenshotCleanup,
763) -> Result<RawScreenshotCapture> {
764 let source = source.to_string();
765 let (path, mut cleanup) = if matches!(cleanup, ScreenshotCleanup::PortalPath { .. }) {
766 tokio::time::timeout(
767 Duration::from_secs(1),
768 tokio::task::spawn_blocking(move || {
769 let path = path
771 .parent()
772 .and_then(|parent| fs::canonicalize(parent).ok())
773 .zip(path.file_name())
774 .map(|(parent, name)| parent.join(name))
775 .unwrap_or(path);
776 let cleanup = match cleanup {
777 ScreenshotCleanup::PortalPath { request, .. } => {
778 ScreenshotCleanup::PortalPath {
779 path: path.clone(),
780 request,
781 }
782 }
783 other => other,
784 };
785 (path, OwnedScreenshotCleanup::new(cleanup))
786 }),
787 )
788 .await
789 .context("timed out preparing portal screenshot file")?
790 .context("portal screenshot preparation task failed")?
791 } else {
792 (path, OwnedScreenshotCleanup::new(cleanup))
793 };
794 run_image_task(move || {
795 let result = read_png_as_capture_inner(&path, &source);
796 cleanup.cleanup_now();
797 result
798 })
799 .await
800}
801
802fn read_png_as_capture_inner(path: &Path, source: &str) -> Result<RawScreenshotCapture> {
803 let file = OpenOptions::new()
804 .read(true)
805 .custom_flags(libc::O_NOFOLLOW | libc::O_NONBLOCK)
806 .open(path)
807 .with_context(|| format!("failed to open screenshot file {}", path.display()))?;
808 let metadata = file
809 .metadata()
810 .with_context(|| format!("failed to inspect screenshot file {}", path.display()))?;
811 if !metadata.file_type().is_file() {
812 bail!("screenshot path is not a regular file: {}", path.display());
813 }
814 if metadata.len() > MAX_SCREENSHOT_SOURCE_BYTES {
815 bail!(
816 "screenshot file is {} bytes, over the {}-byte source limit",
817 metadata.len(),
818 MAX_SCREENSHOT_SOURCE_BYTES
819 );
820 }
821 let mut bytes = Vec::with_capacity(metadata.len() as usize);
822 file.take(MAX_SCREENSHOT_SOURCE_BYTES + 1)
823 .read_to_end(&mut bytes)
824 .with_context(|| format!("failed to read screenshot file {}", path.display()))?;
825 if bytes.len() as u64 > MAX_SCREENSHOT_SOURCE_BYTES {
826 bail!(
827 "screenshot file grew beyond the {}-byte source limit",
828 MAX_SCREENSHOT_SOURCE_BYTES
829 );
830 }
831 if bytes.is_empty() {
832 bail!("screenshot file was empty: {}", path.display());
833 }
834 let (width, height) = png_dimensions(&bytes)?;
835 Ok(RawScreenshotCapture {
836 mime_type: "image/png".to_string(),
837 bytes,
838 source: source.to_string(),
839 width,
840 height,
841 })
842}
843
844fn target_dimensions(
845 width: u32,
846 height: u32,
847 options: ResolvedScreenshotPayloadOptions,
848) -> (u32, u32) {
849 let width_scale = options.max_width as f64 / width as f64;
850 let height_scale = options.max_height as f64 / height as f64;
851 let scale = f64::from(options.scale)
852 .min(width_scale)
853 .min(height_scale)
854 .min(1.0);
855
856 let target_width = ((width as f64 * scale).round() as u32).clamp(1, width);
857 let target_height = ((height as f64 * scale).round() as u32).clamp(1, height);
858 (target_width, target_height)
859}
860
861fn encode_screenshot_to_fit_bytes(
862 raw: &[u8],
863 original_width: u32,
864 original_height: u32,
865 mut target_width: u32,
866 mut target_height: u32,
867 options: ResolvedScreenshotPayloadOptions,
868) -> Result<(Vec<u8>, u32, u32)> {
869 let img = image::load_from_memory_with_format(raw, image::ImageFormat::Png)
870 .context("failed to decode screenshot PNG for encoding")?;
871
872 loop {
873 let bytes = if options.format == ScreenshotOutputFormat::Png
874 && target_width == original_width
875 && target_height == original_height
876 {
877 raw.to_vec()
878 } else {
879 let output = if target_width == original_width && target_height == original_height {
880 img.clone()
881 } else {
882 img.resize_exact(target_width, target_height, FilterType::Lanczos3)
883 };
884 encode_image(&output, options)?
885 };
886
887 if bytes.len() <= options.max_bytes {
888 return Ok((bytes, target_width, target_height));
889 }
890
891 if target_width == 1 && target_height == 1 {
892 bail!(
893 "screenshot payload is {} bytes at 1x1, over max_bytes {}",
894 bytes.len(),
895 options.max_bytes
896 );
897 }
898
899 (target_width, target_height) = next_dimensions_for_byte_cap(
900 target_width,
901 target_height,
902 bytes.len(),
903 options.max_bytes,
904 );
905 }
906}
907
908fn encode_image(
909 img: &image::DynamicImage,
910 options: ResolvedScreenshotPayloadOptions,
911) -> Result<Vec<u8>> {
912 let mut out = Vec::new();
913 match options.format {
914 ScreenshotOutputFormat::Png => {
915 img.write_to(&mut Cursor::new(&mut out), image::ImageFormat::Png)
916 .context("failed to encode screenshot PNG")?;
917 }
918 ScreenshotOutputFormat::Jpeg => {
919 let rgb = img.to_rgb8();
920 JpegEncoder::new_with_quality(&mut out, options.quality)
921 .encode_image(&rgb)
922 .context("failed to encode screenshot JPEG")?;
923 }
924 }
925 Ok(out)
926}
927
928fn next_dimensions_for_byte_cap(
929 width: u32,
930 height: u32,
931 encoded_bytes: usize,
932 max_bytes: usize,
933) -> (u32, u32) {
934 let shrink = ((max_bytes as f64 / encoded_bytes as f64).sqrt() * 0.9).clamp(0.1, 0.95);
935 let mut next_width = ((width as f64 * shrink).floor() as u32).max(1);
936 let mut next_height = ((height as f64 * shrink).floor() as u32).max(1);
937
938 if next_width >= width && width > 1 {
939 next_width = width - 1;
940 }
941 if next_height >= height && height > 1 {
942 next_height = height - 1;
943 }
944
945 (next_width, next_height)
946}
947
948fn cleanup_gnome_requested_path(path: &Path) {
949 let _ = fs::remove_file(path);
950}
951
952fn png_dimensions(bytes: &[u8]) -> Result<(u32, u32)> {
953 const PNG_SIGNATURE: &[u8; 8] = b"\x89PNG\r\n\x1a\n";
954 if bytes.len() < 24 || &bytes[..8] != PNG_SIGNATURE || &bytes[12..16] != b"IHDR" {
955 bail!("screenshot file was not a valid PNG");
956 }
957 let width = u32::from_be_bytes(bytes[16..20].try_into().unwrap());
958 let height = u32::from_be_bytes(bytes[20..24].try_into().unwrap());
959 if width == 0 || height == 0 {
960 bail!("screenshot PNG had invalid dimensions {width}x{height}");
961 }
962 let pixels = u64::from(width) * u64::from(height);
963 if pixels > MAX_SCREENSHOT_SOURCE_PIXELS {
964 bail!(
965 "screenshot PNG has {pixels} pixels, over the {MAX_SCREENSHOT_SOURCE_PIXELS}-pixel source limit"
966 );
967 }
968 Ok((width, height))
969}
970
971fn file_uri_to_path(uri: &str) -> Result<PathBuf> {
972 let Some(rest) = uri.strip_prefix("file://") else {
973 bail!("unsupported screenshot uri: {uri}");
974 };
975 let path = PathBuf::from(percent_decode(rest));
976 if !path.is_absolute() {
977 bail!("screenshot uri did not contain an absolute file path: {uri}");
978 }
979 Ok(path)
980}
981
982fn percent_decode(value: &str) -> String {
983 let bytes = value.as_bytes();
984 let mut decoded = Vec::with_capacity(bytes.len());
985 let mut index = 0;
986
987 while index < bytes.len() {
988 if bytes[index] == b'%' && index + 2 < bytes.len() {
989 if let Ok(hex) = std::str::from_utf8(&bytes[index + 1..index + 3]) {
990 if let Ok(byte) = u8::from_str_radix(hex, 16) {
991 decoded.push(byte);
992 index += 3;
993 continue;
994 }
995 }
996 }
997
998 decoded.push(bytes[index]);
999 index += 1;
1000 }
1001
1002 String::from_utf8_lossy(&decoded).into_owned()
1003}
1004
1005fn temp_png_path(source: &str) -> PathBuf {
1006 std::env::temp_dir().join(format!(
1007 "computer-use-linux-{source}-{}.png",
1008 unique_suffix()
1009 ))
1010}
1011
1012fn request_token() -> String {
1013 format!("computer_use_linux_{}", unique_suffix().replace('-', "_"))
1014}
1015
1016fn unique_suffix() -> String {
1017 let nanos = SystemTime::now()
1018 .duration_since(UNIX_EPOCH)
1019 .map(|duration| duration.as_nanos())
1020 .unwrap_or_default();
1021 format!("{}-{nanos}", std::process::id())
1022}
1023
1024#[cfg(test)]
1025mod tests {
1026 use super::*;
1027
1028 fn test_path(name: &str) -> PathBuf {
1029 std::env::temp_dir().join(format!(
1030 "computer-use-linux-screenshot-test-{name}-{}",
1031 unique_suffix()
1032 ))
1033 }
1034
1035 fn valid_png(width: u32, height: u32) -> Vec<u8> {
1036 let mut png = Vec::new();
1037 png.extend_from_slice(b"\x89PNG\r\n\x1a\n");
1038 png.extend_from_slice(&13_u32.to_be_bytes());
1039 png.extend_from_slice(b"IHDR");
1040 png.extend_from_slice(&width.to_be_bytes());
1041 png.extend_from_slice(&height.to_be_bytes());
1042 png.extend_from_slice(&[8, 6, 0, 0, 0]);
1043 png
1044 }
1045
1046 fn solid_png(width: u32, height: u32) -> Vec<u8> {
1047 let img = image::RgbaImage::from_pixel(width, height, image::Rgba([24, 96, 160, 255]));
1048 encode_test_png(img)
1049 }
1050
1051 fn noisy_png(width: u32, height: u32) -> Vec<u8> {
1052 let mut img = image::RgbaImage::new(width, height);
1053 for (x, y, pixel) in img.enumerate_pixels_mut() {
1054 let r = ((x * 31 + y * 17) % 256) as u8;
1055 let g = ((x * 13 + y * 47) % 256) as u8;
1056 let b = ((x * 97 + y * 7) % 256) as u8;
1057 *pixel = image::Rgba([r, g, b, 255]);
1058 }
1059 encode_test_png(img)
1060 }
1061
1062 fn encode_test_png(img: image::RgbaImage) -> Vec<u8> {
1063 let mut out = Vec::new();
1064 image::DynamicImage::ImageRgba8(img)
1065 .write_to(&mut Cursor::new(&mut out), image::ImageFormat::Png)
1066 .unwrap();
1067 out
1068 }
1069
1070 fn raw_capture(bytes: Vec<u8>) -> RawScreenshotCapture {
1071 let (width, height) = png_dimensions(&bytes).unwrap();
1072 RawScreenshotCapture {
1073 mime_type: "image/png".to_string(),
1074 bytes,
1075 source: "test".to_string(),
1076 width,
1077 height,
1078 }
1079 }
1080
1081 #[test]
1082 fn decodes_file_uri_percent_escapes() {
1083 assert_eq!(
1084 file_uri_to_path("file:///tmp/Codex%20Screenshot.png").unwrap(),
1085 PathBuf::from("/tmp/Codex Screenshot.png")
1086 );
1087 }
1088
1089 #[test]
1090 fn file_uri_requires_an_absolute_path() {
1091 assert!(file_uri_to_path("file://relative/path.png").is_err());
1092 }
1093
1094 #[test]
1095 fn portal_response_requires_the_pinned_sender_and_exact_path() {
1096 let request_path = "/org/freedesktop/portal/desktop/request/1_7/token";
1097 let response = Message::signal(request_path, PORTAL_REQUEST_INTERFACE, "Response")
1098 .unwrap()
1099 .sender(":1.42")
1100 .unwrap()
1101 .build(&())
1102 .unwrap();
1103
1104 assert!(portal_response_matches(&response, request_path, ":1.42"));
1105 assert!(!portal_response_matches(&response, request_path, ":1.99"));
1106 assert!(!portal_response_matches(
1107 &response,
1108 "/org/freedesktop/portal/desktop/request/1_7/other",
1109 ":1.42"
1110 ));
1111 }
1112
1113 #[test]
1114 fn parses_known_backend_names() {
1115 assert_eq!(
1116 ScreenshotBackend::parse("gnome-shell"),
1117 Some(ScreenshotBackend::GnomeShell)
1118 );
1119 assert_eq!(
1120 ScreenshotBackend::parse(" Portal "),
1121 Some(ScreenshotBackend::Portal)
1122 );
1123 assert_eq!(
1124 ScreenshotBackend::parse("GNOME_SCREENSHOT"),
1125 Some(ScreenshotBackend::GnomeScreenshot)
1126 );
1127 assert_eq!(
1128 ScreenshotBackend::parse("x11"),
1129 Some(ScreenshotBackend::X11)
1130 );
1131 assert_eq!(
1132 ScreenshotBackend::parse(" X11-Native "),
1133 Some(ScreenshotBackend::X11)
1134 );
1135 assert_eq!(ScreenshotBackend::parse("nonsense"), None);
1136 }
1137
1138 #[test]
1139 fn forced_backend_reads_env_override() {
1140 std::env::set_var(SCREENSHOT_BACKEND_ENV, "gnome-screenshot");
1142 assert_eq!(
1143 forced_backend().unwrap(),
1144 Some(ScreenshotBackend::GnomeScreenshot)
1145 );
1146
1147 std::env::set_var(SCREENSHOT_BACKEND_ENV, " ");
1148 assert_eq!(forced_backend().unwrap(), None);
1149
1150 std::env::set_var(SCREENSHOT_BACKEND_ENV, "bogus");
1151 let error = forced_backend().unwrap_err();
1152 assert!(error.to_string().contains("not a recognized backend"));
1153
1154 std::env::remove_var(SCREENSHOT_BACKEND_ENV);
1155 assert_eq!(forced_backend().unwrap(), None);
1156 }
1157
1158 #[test]
1159 fn request_token_is_portal_safe() {
1160 let token = request_token();
1161 assert!(token.starts_with("computer_use_linux_"));
1162 assert!(token.chars().all(|c| c.is_ascii_alphanumeric() || c == '_'));
1163 }
1164
1165 #[test]
1166 fn reads_png_dimensions_from_ihdr() {
1167 let png = valid_png(3840, 1080);
1168
1169 assert_eq!(png_dimensions(&png).unwrap(), (3840, 1080));
1170 }
1171
1172 #[test]
1173 fn default_payload_downscales_long_edge() {
1174 let capture =
1175 prepare_screenshot_payload(raw_capture(solid_png(4000, 1000)), Default::default())
1176 .unwrap();
1177
1178 assert_eq!((capture.width, capture.height), (1920, 480));
1179 assert_eq!(
1180 (capture.coordinate_width, capture.coordinate_height),
1181 (4000, 1000)
1182 );
1183 assert!(capture.resized);
1184 assert!(capture.bytes <= DEFAULT_SCREENSHOT_MAX_BYTES);
1185 assert!(capture.data_url.starts_with("data:image/png;base64,"));
1186 }
1187
1188 #[test]
1189 fn larger_bounded_request_can_keep_more_detail() {
1190 let capture = prepare_screenshot_payload(
1191 raw_capture(solid_png(3000, 1000)),
1192 ScreenshotPayloadOptions {
1193 max_width: Some(3000),
1194 max_height: Some(3000),
1195 max_bytes: Some(DEFAULT_SCREENSHOT_MAX_BYTES),
1196 ..Default::default()
1197 },
1198 )
1199 .unwrap();
1200
1201 assert_eq!((capture.width, capture.height), (3000, 1000));
1202 assert_eq!(
1203 (capture.coordinate_width, capture.coordinate_height),
1204 (3000, 1000)
1205 );
1206 assert!(!capture.resized);
1207 }
1208
1209 #[test]
1210 fn byte_cap_downscales_until_payload_fits() {
1211 let capture = prepare_screenshot_payload(
1212 raw_capture(noisy_png(512, 512)),
1213 ScreenshotPayloadOptions {
1214 max_width: Some(512),
1215 max_height: Some(512),
1216 max_bytes: Some(20_000),
1217 ..Default::default()
1218 },
1219 )
1220 .unwrap();
1221
1222 assert!(capture.bytes <= 20_000);
1223 assert!(capture.width < 512);
1224 assert_eq!(
1225 (capture.coordinate_width, capture.coordinate_height),
1226 (512, 512)
1227 );
1228 assert!(capture.resized);
1229 }
1230
1231 #[test]
1232 fn jpeg_format_compresses_when_requested() {
1233 let capture = prepare_screenshot_payload(
1234 raw_capture(noisy_png(512, 512)),
1235 ScreenshotPayloadOptions {
1236 max_width: Some(512),
1237 max_height: Some(512),
1238 max_bytes: Some(DEFAULT_SCREENSHOT_MAX_BYTES),
1239 format: Some(ScreenshotOutputFormat::Jpeg),
1240 quality: Some(60),
1241 ..Default::default()
1242 },
1243 )
1244 .unwrap();
1245
1246 assert_eq!(capture.mime_type, "image/jpeg");
1247 assert_eq!(capture.format, ScreenshotOutputFormat::Jpeg);
1248 assert_eq!(capture.quality, Some(60));
1249 assert_eq!((capture.width, capture.height), (512, 512));
1250 assert_eq!(
1251 (capture.coordinate_width, capture.coordinate_height),
1252 (512, 512)
1253 );
1254 assert!(capture.bytes < capture.original_bytes);
1255 assert!(capture.data_url.starts_with("data:image/jpeg;base64,"));
1256 }
1257
1258 #[tokio::test]
1259 async fn cancelled_queued_capture_removes_its_owned_file() {
1260 let cache = std::env::var_os("XDG_CACHE_HOME")
1261 .map(PathBuf::from)
1262 .unwrap_or_else(|| PathBuf::from(std::env::var_os("HOME").unwrap()).join(".cache"));
1263 let path = cache.join(format!(
1264 "cul-cancelled-capture-{}.png",
1265 getrandom::u64().unwrap()
1266 ));
1267 fs::create_dir_all(&cache).unwrap();
1268 fs::write(&path, valid_png(1, 1)).unwrap();
1269 let (release, wait) = std::sync::mpsc::channel();
1270 let (started, ready) = tokio::sync::oneshot::channel();
1271 let busy = tokio::spawn(run_image_task(move || {
1272 started.send(()).unwrap();
1273 wait.recv().unwrap();
1274 Ok(())
1275 }));
1276 ready.await.unwrap();
1277 let mut queued = tokio::spawn(read_png_as_capture(
1278 path.clone(),
1279 "owned-test-capture",
1280 ScreenshotCleanup::DeletePath(path.clone()),
1281 ));
1282 assert!(tokio::time::timeout(Duration::from_millis(20), &mut queued)
1283 .await
1284 .is_err());
1285 queued.abort();
1286 assert!(queued.await.unwrap_err().is_cancelled());
1287 release.send(()).unwrap();
1288 busy.await.unwrap().unwrap();
1289 let removed = !path.exists();
1290 let _ = fs::remove_file(path);
1291 assert!(removed, "cancelled screenshot left its owned file behind");
1292 }
1293
1294 #[tokio::test]
1295 async fn portal_capture_removes_new_returned_path() {
1296 let path = test_path("portal-valid");
1297 let requested_at = SystemTime::now();
1298 tokio::time::sleep(Duration::from_millis(20)).await;
1301 fs::write(&path, valid_png(1, 1)).unwrap();
1302 let capture = read_portal_png_as_capture(
1303 path.clone(),
1304 PortalScreenshotRequest {
1305 requested_at,
1306 pictures: None,
1307 },
1308 )
1309 .await
1310 .unwrap();
1311 assert_eq!(capture.source, "xdg-desktop-portal");
1312 let retained = path.exists();
1313 let _ = fs::remove_file(&path);
1314 assert!(!retained, "portal screenshot was left behind");
1315 }
1316
1317 #[tokio::test]
1318 async fn portal_capture_removes_new_invalid_returned_path() {
1319 let path = test_path("portal-invalid");
1320 let requested_at = SystemTime::now();
1321 tokio::time::sleep(Duration::from_millis(20)).await;
1324 fs::write(&path, b"").unwrap();
1325 let error = read_portal_png_as_capture(
1326 path.clone(),
1327 PortalScreenshotRequest {
1328 requested_at,
1329 pictures: None,
1330 },
1331 )
1332 .await
1333 .unwrap_err();
1334 assert!(error.to_string().contains("screenshot file was empty"));
1335 let retained = path.exists();
1336 let _ = fs::remove_file(&path);
1337 assert!(!retained, "failed portal screenshot was left behind");
1338 }
1339
1340 #[tokio::test]
1341 async fn portal_capture_preserves_preexisting_returned_path() {
1342 let path = test_path("portal-existing");
1343 fs::write(&path, valid_png(1, 1)).unwrap();
1344 let requested_at = SystemTime::now();
1345 let capture = read_portal_png_as_capture(
1346 path.clone(),
1347 PortalScreenshotRequest {
1348 requested_at,
1349 pictures: None,
1350 },
1351 )
1352 .await
1353 .unwrap();
1354 assert_eq!(capture.source, "xdg-desktop-portal");
1355 assert!(path.exists(), "preexisting user image was removed");
1356 let _ = fs::remove_file(path);
1357 }
1358
1359 #[tokio::test]
1360 async fn portal_capture_cleans_same_tick_picture_and_preserves_existing_name() {
1361 let directory = test_path("portal-pictures");
1362 fs::create_dir_all(&directory).unwrap();
1363 let existing = directory.join("existing.png");
1364 fs::write(&existing, valid_png(1, 1)).unwrap();
1365 let pictures = PortalPicturesSnapshot::new(directory.clone()).unwrap();
1366 let path = directory.join("new.png");
1367 let request = PortalScreenshotRequest {
1368 requested_at: SystemTime::now(),
1369 pictures: Some(pictures),
1370 };
1371 fs::write(&path, valid_png(1, 1)).unwrap();
1372 let capture = read_portal_png_as_capture(path.clone(), request)
1373 .await
1374 .unwrap();
1375 assert_eq!(capture.source, "xdg-desktop-portal");
1376 assert!(!path.exists(), "same-tick portal capture was left behind");
1377 let pictures = PortalPicturesSnapshot::new(directory.clone()).unwrap();
1378 fs::remove_file(&existing).unwrap();
1379 fs::write(&existing, valid_png(1, 1)).unwrap();
1380 let request = PortalScreenshotRequest {
1381 requested_at: UNIX_EPOCH,
1382 pictures: Some(pictures),
1383 };
1384 read_portal_png_as_capture(existing.clone(), request)
1385 .await
1386 .unwrap();
1387 assert!(existing.exists(), "preexisting filename was removed");
1388 fs::remove_dir_all(directory).unwrap();
1389 }
1390
1391 #[test]
1392 fn portal_cleanup_preserves_a_replacement_file() {
1393 let path = test_path("portal-replaced");
1394 fs::write(&path, valid_png(1, 1)).unwrap();
1395 let created = fs::metadata(&path).unwrap().created().unwrap();
1396 let cleanup = OwnedScreenshotCleanup::new(ScreenshotCleanup::PortalPath {
1397 path: path.clone(),
1398 request: PortalScreenshotRequest {
1399 requested_at: created,
1400 pictures: None,
1401 },
1402 });
1403 fs::remove_file(&path).unwrap();
1404 fs::write(&path, b"user replacement").unwrap();
1405 drop(cleanup);
1406 let result = fs::read(&path);
1407 let _ = fs::remove_file(&path);
1408 assert_eq!(result.unwrap(), b"user replacement");
1409 }
1410
1411 #[tokio::test]
1412 async fn screenshot_reader_rejects_symlinks() {
1413 let target = test_path("symlink-target");
1414 let link = test_path("symlink-link");
1415 fs::write(&target, valid_png(1, 1)).unwrap();
1416 std::os::unix::fs::symlink(&target, &link).unwrap();
1417
1418 let error = read_png_as_capture(link.clone(), "test", ScreenshotCleanup::Preserve)
1419 .await
1420 .unwrap_err();
1421
1422 assert!(error.to_string().contains("failed to open screenshot file"));
1423 let _ = fs::remove_file(link);
1424 let _ = fs::remove_file(target);
1425 }
1426
1427 #[tokio::test]
1428 async fn screenshot_reader_rejects_oversized_source_files() {
1429 let path = test_path("oversized-source");
1430 let file = fs::File::create(&path).unwrap();
1431 file.set_len(MAX_SCREENSHOT_SOURCE_BYTES + 1).unwrap();
1432
1433 let error = read_png_as_capture(path.clone(), "test", ScreenshotCleanup::Preserve)
1434 .await
1435 .unwrap_err();
1436
1437 assert!(error.to_string().contains("over the"));
1438 let _ = fs::remove_file(path);
1439 }
1440
1441 #[test]
1442 fn png_dimensions_reject_decompression_bomb_dimensions() {
1443 let png = valid_png(u32::MAX, u32::MAX);
1444
1445 assert!(png_dimensions(&png)
1446 .unwrap_err()
1447 .to_string()
1448 .contains("pixel source limit"));
1449 }
1450
1451 #[tokio::test]
1452 async fn gnome_capture_deletes_backend_temp_path_on_success() {
1453 let path = test_path("gnome-valid");
1454 fs::write(&path, valid_png(1, 1)).unwrap();
1455
1456 let capture = read_png_as_capture(
1457 path.clone(),
1458 "gnome-shell",
1459 ScreenshotCleanup::DeletePath(path.clone()),
1460 )
1461 .await
1462 .unwrap();
1463
1464 assert_eq!(capture.source, "gnome-shell");
1465 assert!(!path.exists());
1466 }
1467
1468 #[tokio::test]
1469 async fn gnome_capture_deletes_backend_temp_path_on_parse_failure() {
1470 let path = test_path("gnome-invalid");
1471 fs::write(&path, b"").unwrap();
1472
1473 let error = read_png_as_capture(
1474 path.clone(),
1475 "gnome-shell",
1476 ScreenshotCleanup::DeletePath(path.clone()),
1477 )
1478 .await
1479 .unwrap_err();
1480
1481 assert!(error.to_string().contains("screenshot file was empty"));
1482 assert!(!path.exists());
1483 }
1484
1485 #[test]
1486 fn gnome_failure_cleanup_removes_requested_temp_path() {
1487 let path = test_path("gnome-pre-read-failure");
1488 fs::write(&path, b"partial").unwrap();
1489
1490 cleanup_gnome_requested_path(&path);
1491
1492 assert!(!path.exists());
1493 }
1494
1495 #[tokio::test]
1496 async fn gnome_deletes_requested_temp_path_and_preserves_unexpected_returned_path() {
1497 let requested = test_path("gnome-requested");
1498 let returned = test_path("gnome-returned");
1499 fs::write(&requested, b"partial").unwrap();
1500 fs::write(&returned, valid_png(1, 1)).unwrap();
1501
1502 let capture = read_png_as_capture(
1503 returned.clone(),
1504 "gnome-shell",
1505 ScreenshotCleanup::DeletePath(requested.clone()),
1506 )
1507 .await
1508 .unwrap();
1509
1510 assert_eq!(capture.source, "gnome-shell");
1511 assert!(!requested.exists());
1512 assert!(returned.exists());
1513 let _ = fs::remove_file(returned);
1514 }
1515}