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