winx_code_agent/tools/
read_image.rs

1//! Implementation of the `ReadImage` tool.
2//!
3//! This module provides the implementation for the `ReadImage` tool, which is used
4//! to read image files and return their contents as base64-encoded data with
5//! the appropriate MIME type.
6
7use base64::{engine::general_purpose, Engine};
8use mime_guess::MimeGuess;
9use std::path::{Path, PathBuf};
10use std::sync::Arc;
11use tokio::sync::Mutex;
12use tracing::{debug, error, info, instrument};
13
14use crate::errors::{Result, WinxError};
15use crate::state::bash_state::BashState;
16use crate::types::ReadImage;
17use crate::utils::path::{expand_user, validate_path_in_workspace, validate_path_with_roots};
18
19/// Supported MIME types for images
20pub const SUPPORTED_MIME_TYPES: [&str; 4] = ["image/jpeg", "image/png", "image/gif", "image/webp"];
21
22/// Refuse to read images larger than this. base64 encoding inflates the
23/// footprint ~1.37×, so a few-hundred-MB file would spike to gigabytes; under
24/// `panic = "abort"` an allocation failure aborts the whole server. Real images
25/// (screenshots, mockups, error PNGs) are orders of magnitude smaller.
26const MAX_IMAGE_BYTES: u64 = 50 * 1024 * 1024;
27
28/// Read an image from the file system
29///
30/// This function reads an image file, base64 encodes its contents, and
31/// determines the MIME type based on the file extension.
32///
33/// # Arguments
34///
35/// * `file_path` - Path to the image file
36/// * `cwd` - Current working directory for resolving relative paths
37///
38/// # Returns
39///
40/// A tuple containing:
41/// - The MIME type of the image
42/// - The base64-encoded image data
43///
44/// # Errors
45///
46/// Returns an error if the file cannot be accessed or read
47#[instrument(level = "debug", skip(file_path))]
48fn read_image_from_path(
49    file_path: &str,
50    cwd: &Path,
51    workspace_root: &Path,
52) -> Result<(String, String)> {
53    read_image_from_path_impl(file_path, cwd, workspace_root, None)
54}
55
56#[cfg(test)]
57fn read_image_from_path_with_roots(
58    file_path: &str,
59    cwd: &Path,
60    workspace_root: &Path,
61    extra_roots: &[PathBuf],
62) -> Result<(String, String)> {
63    read_image_from_path_impl(file_path, cwd, workspace_root, Some(extra_roots))
64}
65
66fn read_image_from_path_impl(
67    file_path: &str,
68    cwd: &Path,
69    workspace_root: &Path,
70    extra_roots: Option<&[PathBuf]>,
71) -> Result<(String, String)> {
72    debug!("Reading image: {}", file_path);
73
74    // Expand the path
75    let file_path = expand_user(file_path);
76
77    // Ensure path is absolute
78    let path = if Path::new(&file_path).is_absolute() {
79        PathBuf::from(&file_path)
80    } else {
81        // Use current working directory if path is relative
82        cwd.join(&file_path)
83    };
84
85    // Confine to the workspace and resolve symlinks BEFORE touching the file —
86    // the same guarantee ReadFiles/ContextSave give. Without this, ReadImage is
87    // an arbitrary-file read primitive (`/etc/shadow`, `~/.ssh/*`), and over the
88    // network-reachable HTTP transport that means remote exfiltration.
89    let validated = match extra_roots {
90        Some(extra_roots) => validate_path_with_roots(&path, workspace_root, extra_roots),
91        None => validate_path_in_workspace(&path, workspace_root),
92    };
93    let path = validated
94        .map_err(|e| WinxError::PathSecurityError { path: path.clone(), message: e.to_string() })?;
95
96    // Ensure it's a regular file (also rejects a path that doesn't exist).
97    if !path.is_file() {
98        return Err(WinxError::FileAccessError {
99            path: path.clone(),
100            message: "File does not exist or is not a regular file".to_string(),
101        });
102    }
103
104    // Cap the size before reading it into memory + base64-encoding it.
105    let size = std::fs::metadata(&path).map_or(0, |m| m.len());
106    if size > MAX_IMAGE_BYTES {
107        return Err(WinxError::FileAccessError {
108            path: path.clone(),
109            message: format!("Image too large: {size} bytes (max {MAX_IMAGE_BYTES} bytes)"),
110        });
111    }
112
113    // Read the file as bytes
114    let image_bytes = std::fs::read(&path).map_err(|e| WinxError::FileAccessError {
115        path: path.clone(),
116        message: format!("Error reading file: {e}"),
117    })?;
118
119    // Encode the bytes to base64
120    let image_b64 = general_purpose::STANDARD.encode(&image_bytes);
121
122    // Guess the MIME type from the file extension
123    let mime_type =
124        MimeGuess::from_path(&path).first_raw().unwrap_or("application/octet-stream").to_string();
125
126    // Verify the MIME type is a supported image type
127    if SUPPORTED_MIME_TYPES.contains(&mime_type.as_str()) {
128        Ok((mime_type, image_b64))
129    } else {
130        debug!("Detected MIME type '{}' is not in the supported list", mime_type);
131        // Fall back to a best effort based on common extensions
132        let extension = path.extension().and_then(|ext| ext.to_str()).unwrap_or("").to_lowercase();
133
134        let mime_type = match extension.as_str() {
135            "png" => "image/png",
136            "gif" => "image/gif",
137            "webp" => "image/webp",
138            _ => "image/jpeg", // Default fallback
139        };
140
141        debug!("Using fallback MIME type: {}", mime_type);
142        Ok((mime_type.to_string(), image_b64))
143    }
144}
145
146/// Handle the `ReadImage` tool call
147///
148/// This function processes the `ReadImage` tool call, which reads an image file
149/// and returns its contents as base64-encoded data with the appropriate MIME type.
150///
151/// # Arguments
152///
153/// * `bash_state_arc` - Shared reference to the bash state
154/// * `read_image` - The read image parameters
155///
156/// # Returns
157///
158/// A Result containing a tuple with the MIME type and base64-encoded image data
159///
160/// # Errors
161///
162/// Returns an error if the image file cannot be accessed or read
163#[instrument(level = "info", skip(bash_state_arc, read_image))]
164pub async fn handle_tool_call(
165    bash_state_arc: &Arc<Mutex<Option<BashState>>>,
166    read_image: ReadImage,
167) -> Result<(String, String)> {
168    info!("ReadImage tool called with: {:?}", read_image);
169
170    // We need to extract data from the bash state before awaiting
171    // to avoid holding the MutexGuard across await points
172    let cwd: PathBuf;
173    let workspace_root: PathBuf;
174
175    // Lock bash state to extract data
176    {
177        let bash_state_guard = bash_state_arc.lock().await;
178
179        // Ensure bash state is initialized
180        let Some(bash_state) = &*bash_state_guard else {
181            error!("BashState not initialized");
182            return Err(WinxError::BashStateNotInitialized);
183        };
184
185        // Extract needed data
186        cwd = bash_state.cwd.clone();
187        workspace_root = bash_state.workspace_root.clone();
188    }
189
190    // Filesystem reads and base64 encoding are blocking/CPU-heavy for large
191    // images, so keep them off Tokio's request workers.
192    let file_path = read_image.file_path;
193    tokio::task::spawn_blocking(move || read_image_from_path(&file_path, &cwd, &workspace_root))
194        .await
195        .map_err(|error| {
196            WinxError::CommandExecutionError(format!("ReadImage worker failed: {error}"))
197        })?
198}
199
200#[cfg(test)]
201mod tests {
202    #![allow(clippy::unwrap_used)]
203    use super::*;
204    use std::fs;
205    use tempfile::TempDir;
206
207    #[test]
208    fn reads_image_inside_workspace() {
209        let ws = TempDir::new().unwrap();
210        let img = ws.path().join("shot.png");
211        fs::write(&img, b"\x89PNG\r\n\x1a\nfake").unwrap();
212        let (mime, b64) =
213            read_image_from_path_with_roots(img.to_str().unwrap(), ws.path(), ws.path(), &[])
214                .unwrap();
215        assert_eq!(mime, "image/png");
216        assert!(!b64.is_empty());
217    }
218
219    #[test]
220    fn rejects_image_outside_workspace() {
221        // The exfil case: an absolute path outside the workspace must be refused,
222        // not read (this is what made ReadImage an arbitrary-file-read primitive).
223        let ws = TempDir::new().unwrap();
224        let outside = TempDir::new().unwrap();
225        let secret = outside.path().join("secret.png");
226        fs::write(&secret, b"\x89PNG secret").unwrap();
227        let err =
228            read_image_from_path_with_roots(secret.to_str().unwrap(), ws.path(), ws.path(), &[]);
229        assert!(matches!(err, Err(WinxError::PathSecurityError { .. })));
230    }
231}
winx_code_agent/tools/read_image.rs

winx_code_agent/tools/
read_image.rs
