Skip to main content

computer_use_linux/
screenshot.rs

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