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