construct/tools/
image_info.rs

1use super::traits::{Tool, ToolResult};
2use crate::security::SecurityPolicy;
3use async_trait::async_trait;
4use serde_json::json;
5use std::fmt::Write;
6use std::path::Path;
7use std::sync::Arc;
8
9/// Maximum file size we will read and base64-encode (5 MB).
10const MAX_IMAGE_BYTES: u64 = 5_242_880;
11
12/// Tool to read image metadata and optionally return base64-encoded data.
13///
14/// Since providers are currently text-only, this tool extracts what it can
15/// (file size, format, dimensions from header bytes) and provides base64
16/// data for future multimodal provider support.
17pub struct ImageInfoTool {
18    security: Arc<SecurityPolicy>,
19}
20
21impl ImageInfoTool {
22    pub fn new(security: Arc<SecurityPolicy>) -> Self {
23        Self { security }
24    }
25
26    /// Detect image format from first few bytes (magic numbers).
27    fn detect_format(bytes: &[u8]) -> &'static str {
28        if bytes.len() < 4 {
29            return "unknown";
30        }
31        if bytes.starts_with(b"\x89PNG") {
32            "png"
33        } else if bytes.starts_with(b"\xFF\xD8\xFF") {
34            "jpeg"
35        } else if bytes.starts_with(b"GIF8") {
36            "gif"
37        } else if bytes.starts_with(b"RIFF") && bytes.len() >= 12 && &bytes[8..12] == b"WEBP" {
38            "webp"
39        } else if bytes.starts_with(b"BM") {
40            "bmp"
41        } else {
42            "unknown"
43        }
44    }
45
46    /// Try to extract dimensions from image header bytes.
47    /// Returns (width, height) if detectable.
48    fn extract_dimensions(bytes: &[u8], format: &str) -> Option<(u32, u32)> {
49        match format {
50            "png" => {
51                // PNG IHDR chunk: bytes 16-19 = width, 20-23 = height (big-endian)
52                if bytes.len() >= 24 {
53                    let w = u32::from_be_bytes([bytes[16], bytes[17], bytes[18], bytes[19]]);
54                    let h = u32::from_be_bytes([bytes[20], bytes[21], bytes[22], bytes[23]]);
55                    Some((w, h))
56                } else {
57                    None
58                }
59            }
60            "gif" => {
61                // GIF: bytes 6-7 = width, 8-9 = height (little-endian)
62                if bytes.len() >= 10 {
63                    let w = u32::from(u16::from_le_bytes([bytes[6], bytes[7]]));
64                    let h = u32::from(u16::from_le_bytes([bytes[8], bytes[9]]));
65                    Some((w, h))
66                } else {
67                    None
68                }
69            }
70            "bmp" => {
71                // BMP: bytes 18-21 = width, 22-25 = height (little-endian, signed)
72                if bytes.len() >= 26 {
73                    let w = u32::from_le_bytes([bytes[18], bytes[19], bytes[20], bytes[21]]);
74                    let h_raw = i32::from_le_bytes([bytes[22], bytes[23], bytes[24], bytes[25]]);
75                    let h = h_raw.unsigned_abs();
76                    Some((w, h))
77                } else {
78                    None
79                }
80            }
81            "jpeg" => Self::jpeg_dimensions(bytes),
82            _ => None,
83        }
84    }
85
86    /// Parse JPEG SOF markers to extract dimensions.
87    fn jpeg_dimensions(bytes: &[u8]) -> Option<(u32, u32)> {
88        let mut i = 2; // skip SOI marker
89        while i + 1 < bytes.len() {
90            if bytes[i] != 0xFF {
91                return None;
92            }
93            let marker = bytes[i + 1];
94            i += 2;
95
96            // SOF0..SOF3 markers contain dimensions
97            if (0xC0..=0xC3).contains(&marker) {
98                if i + 7 <= bytes.len() {
99                    let h = u32::from(u16::from_be_bytes([bytes[i + 3], bytes[i + 4]]));
100                    let w = u32::from(u16::from_be_bytes([bytes[i + 5], bytes[i + 6]]));
101                    return Some((w, h));
102                }
103                return None;
104            }
105
106            // Skip this segment
107            if i + 1 < bytes.len() {
108                let seg_len = u16::from_be_bytes([bytes[i], bytes[i + 1]]) as usize;
109                if seg_len < 2 {
110                    return None; // Malformed segment (valid segments have length >= 2)
111                }
112                i += seg_len;
113            } else {
114                return None;
115            }
116        }
117        None
118    }
119}
120
121#[async_trait]
122impl Tool for ImageInfoTool {
123    fn name(&self) -> &str {
124        "image_info"
125    }
126
127    fn description(&self) -> &str {
128        "Read image file metadata (format, dimensions, size) and optionally return base64-encoded data."
129    }
130
131    fn parameters_schema(&self) -> serde_json::Value {
132        json!({
133            "type": "object",
134            "properties": {
135                "path": {
136                    "type": "string",
137                    "description": "Path to the image file (absolute or relative to workspace)"
138                },
139                "include_base64": {
140                    "type": "boolean",
141                    "description": "Include base64-encoded image data in output (default: false)"
142                }
143            },
144            "required": ["path"]
145        })
146    }
147
148    async fn execute(&self, args: serde_json::Value) -> anyhow::Result<ToolResult> {
149        let path_str = args
150            .get("path")
151            .and_then(|v| v.as_str())
152            .ok_or_else(|| anyhow::anyhow!("Missing 'path' parameter"))?;
153
154        let include_base64 = args
155            .get("include_base64")
156            .and_then(serde_json::Value::as_bool)
157            .unwrap_or(false);
158
159        let path = Path::new(path_str);
160
161        // Restrict reads to workspace directory to prevent arbitrary file exfiltration
162        if !self.security.is_path_allowed(path_str) {
163            return Ok(ToolResult {
164                success: false,
165                output: String::new(),
166                error: Some(format!(
167                    "Path not allowed: {path_str} (must be within workspace)"
168                )),
169            });
170        }
171
172        if !path.exists() {
173            return Ok(ToolResult {
174                success: false,
175                output: String::new(),
176                error: Some(format!("File not found: {path_str}")),
177            });
178        }
179
180        let metadata = tokio::fs::metadata(path)
181            .await
182            .map_err(|e| anyhow::anyhow!("Failed to read file metadata: {e}"))?;
183
184        let file_size = metadata.len();
185
186        if file_size > MAX_IMAGE_BYTES {
187            return Ok(ToolResult {
188                success: false,
189                output: String::new(),
190                error: Some(format!(
191                    "Image too large: {file_size} bytes (max {MAX_IMAGE_BYTES} bytes)"
192                )),
193            });
194        }
195
196        let bytes = tokio::fs::read(path)
197            .await
198            .map_err(|e| anyhow::anyhow!("Failed to read image file: {e}"))?;
199
200        let format = Self::detect_format(&bytes);
201        let dimensions = Self::extract_dimensions(&bytes, format);
202
203        let mut output = format!("File: {path_str}\nFormat: {format}\nSize: {file_size} bytes");
204
205        if let Some((w, h)) = dimensions {
206            let _ = write!(output, "\nDimensions: {w}x{h}");
207        }
208
209        if include_base64 {
210            use base64::Engine;
211            let encoded = base64::engine::general_purpose::STANDARD.encode(&bytes);
212            let mime = match format {
213                "png" => "image/png",
214                "jpeg" => "image/jpeg",
215                "gif" => "image/gif",
216                "webp" => "image/webp",
217                "bmp" => "image/bmp",
218                _ => "application/octet-stream",
219            };
220            let _ = write!(output, "\ndata:{mime};base64,{encoded}");
221        }
222
223        Ok(ToolResult {
224            success: true,
225            output,
226            error: None,
227        })
228    }
229}
230
231#[cfg(test)]
232mod tests {
233    use super::*;
234    use crate::security::{AutonomyLevel, SecurityPolicy};
235
236    fn test_security() -> Arc<SecurityPolicy> {
237        Arc::new(SecurityPolicy {
238            autonomy: AutonomyLevel::Full,
239            workspace_dir: std::env::temp_dir(),
240            workspace_only: false,
241            forbidden_paths: vec![],
242            ..SecurityPolicy::default()
243        })
244    }
245
246    #[test]
247    fn image_info_tool_name() {
248        let tool = ImageInfoTool::new(test_security());
249        assert_eq!(tool.name(), "image_info");
250    }
251
252    #[test]
253    fn image_info_tool_description() {
254        let tool = ImageInfoTool::new(test_security());
255        assert!(!tool.description().is_empty());
256        assert!(tool.description().contains("image"));
257    }
258
259    #[test]
260    fn image_info_tool_schema() {
261        let tool = ImageInfoTool::new(test_security());
262        let schema = tool.parameters_schema();
263        assert!(schema["properties"]["path"].is_object());
264        assert!(schema["properties"]["include_base64"].is_object());
265        let required = schema["required"].as_array().unwrap();
266        assert!(required.contains(&json!("path")));
267    }
268
269    #[test]
270    fn image_info_tool_spec() {
271        let tool = ImageInfoTool::new(test_security());
272        let spec = tool.spec();
273        assert_eq!(spec.name, "image_info");
274        assert!(spec.parameters.is_object());
275    }
276
277    // ── Format detection ────────────────────────────────────────
278
279    #[test]
280    fn detect_png() {
281        let bytes = b"\x89PNG\r\n\x1a\n";
282        assert_eq!(ImageInfoTool::detect_format(bytes), "png");
283    }
284
285    #[test]
286    fn detect_jpeg() {
287        let bytes = b"\xFF\xD8\xFF\xE0";
288        assert_eq!(ImageInfoTool::detect_format(bytes), "jpeg");
289    }
290
291    #[test]
292    fn detect_gif() {
293        let bytes = b"GIF89a";
294        assert_eq!(ImageInfoTool::detect_format(bytes), "gif");
295    }
296
297    #[test]
298    fn detect_webp() {
299        let bytes = b"RIFF\x00\x00\x00\x00WEBP";
300        assert_eq!(ImageInfoTool::detect_format(bytes), "webp");
301    }
302
303    #[test]
304    fn detect_bmp() {
305        let bytes = b"BM\x00\x00";
306        assert_eq!(ImageInfoTool::detect_format(bytes), "bmp");
307    }
308
309    #[test]
310    fn detect_unknown_short() {
311        let bytes = b"\x00\x01";
312        assert_eq!(ImageInfoTool::detect_format(bytes), "unknown");
313    }
314
315    #[test]
316    fn detect_unknown_garbage() {
317        let bytes = b"this is not an image";
318        assert_eq!(ImageInfoTool::detect_format(bytes), "unknown");
319    }
320
321    // ── Dimension extraction ────────────────────────────────────
322
323    #[test]
324    fn png_dimensions() {
325        // Minimal PNG IHDR: 8-byte signature + 4-byte length + 4-byte IHDR + 4-byte width + 4-byte height
326        let mut bytes = vec![
327            0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, // PNG signature
328            0x00, 0x00, 0x00, 0x0D, // IHDR length
329            0x49, 0x48, 0x44, 0x52, // "IHDR"
330            0x00, 0x00, 0x03, 0x20, // width: 800
331            0x00, 0x00, 0x02, 0x58, // height: 600
332        ];
333        bytes.extend_from_slice(&[0u8; 10]); // padding
334        let dims = ImageInfoTool::extract_dimensions(&bytes, "png");
335        assert_eq!(dims, Some((800, 600)));
336    }
337
338    #[test]
339    fn gif_dimensions() {
340        let bytes = [
341            0x47, 0x49, 0x46, 0x38, 0x39, 0x61, // GIF89a
342            0x40, 0x01, // width: 320 (LE)
343            0xF0, 0x00, // height: 240 (LE)
344        ];
345        let dims = ImageInfoTool::extract_dimensions(&bytes, "gif");
346        assert_eq!(dims, Some((320, 240)));
347    }
348
349    #[test]
350    fn bmp_dimensions() {
351        let mut bytes = vec![0u8; 26];
352        bytes[0] = b'B';
353        bytes[1] = b'M';
354        // width at offset 18 (LE): 1024
355        bytes[18] = 0x00;
356        bytes[19] = 0x04;
357        bytes[20] = 0x00;
358        bytes[21] = 0x00;
359        // height at offset 22 (LE): 768
360        bytes[22] = 0x00;
361        bytes[23] = 0x03;
362        bytes[24] = 0x00;
363        bytes[25] = 0x00;
364        let dims = ImageInfoTool::extract_dimensions(&bytes, "bmp");
365        assert_eq!(dims, Some((1024, 768)));
366    }
367
368    #[test]
369    fn jpeg_dimensions() {
370        // Minimal JPEG-like byte sequence with SOF0 marker
371        let mut bytes: Vec<u8> = vec![
372            0xFF, 0xD8, // SOI
373            0xFF, 0xE0, // APP0 marker
374            0x00, 0x10, // APP0 length = 16
375        ];
376        bytes.extend_from_slice(&[0u8; 14]); // APP0 payload
377        bytes.extend_from_slice(&[
378            0xFF, 0xC0, // SOF0 marker
379            0x00, 0x11, // SOF0 length
380            0x08, // precision
381            0x01, 0xE0, // height: 480
382            0x02, 0x80, // width: 640
383        ]);
384        let dims = ImageInfoTool::extract_dimensions(&bytes, "jpeg");
385        assert_eq!(dims, Some((640, 480)));
386    }
387
388    #[test]
389    fn jpeg_malformed_zero_length_segment() {
390        // Zero-length segment should return None instead of looping forever
391        let bytes: Vec<u8> = vec![
392            0xFF, 0xD8, // SOI
393            0xFF, 0xE0, // APP0 marker
394            0x00, 0x00, // length = 0 (malformed)
395        ];
396        let dims = ImageInfoTool::extract_dimensions(&bytes, "jpeg");
397        assert!(dims.is_none());
398    }
399
400    #[test]
401    fn unknown_format_no_dimensions() {
402        let bytes = b"random data here";
403        let dims = ImageInfoTool::extract_dimensions(bytes, "unknown");
404        assert!(dims.is_none());
405    }
406
407    // ── Execute tests ───────────────────────────────────────────
408
409    #[tokio::test]
410    async fn execute_missing_path() {
411        let tool = ImageInfoTool::new(test_security());
412        let result = tool.execute(json!({})).await;
413        assert!(result.is_err());
414    }
415
416    #[tokio::test]
417    async fn execute_nonexistent_file() {
418        let tool = ImageInfoTool::new(test_security());
419        let result = tool
420            .execute(json!({"path": "/tmp/nonexistent_image_xyz.png"}))
421            .await
422            .unwrap();
423        assert!(!result.success);
424        assert!(result.error.as_ref().unwrap().contains("not found"));
425    }
426
427    #[tokio::test]
428    async fn execute_real_file() {
429        // Create a minimal valid PNG
430        let dir = std::env::temp_dir().join("construct_image_info_test");
431        let _ = tokio::fs::create_dir_all(&dir).await;
432        let png_path = dir.join("test.png");
433
434        // Minimal 1x1 red PNG (67 bytes)
435        let png_bytes: Vec<u8> = vec![
436            0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, // signature
437            0x00, 0x00, 0x00, 0x0D, // IHDR length
438            0x49, 0x48, 0x44, 0x52, // IHDR
439            0x00, 0x00, 0x00, 0x01, // width: 1
440            0x00, 0x00, 0x00, 0x01, // height: 1
441            0x08, 0x02, 0x00, 0x00, 0x00, // bit depth, color type, etc.
442            0x90, 0x77, 0x53, 0xDE, // CRC
443            0x00, 0x00, 0x00, 0x0C, // IDAT length
444            0x49, 0x44, 0x41, 0x54, // IDAT
445            0x08, 0xD7, 0x63, 0xF8, 0xCF, 0xC0, 0x00, 0x00, 0x00, 0x02, 0x00, 0x01, 0xE2, 0x21,
446            0xBC, 0x33, // CRC
447            0x00, 0x00, 0x00, 0x00, // IEND length
448            0x49, 0x45, 0x4E, 0x44, // IEND
449            0xAE, 0x42, 0x60, 0x82, // CRC
450        ];
451        tokio::fs::write(&png_path, &png_bytes).await.unwrap();
452
453        let tool = ImageInfoTool::new(test_security());
454        let result = tool
455            .execute(json!({"path": png_path.to_string_lossy()}))
456            .await
457            .unwrap();
458        assert!(result.success);
459        assert!(result.output.contains("Format: png"));
460        assert!(result.output.contains("Dimensions: 1x1"));
461        assert!(!result.output.contains("data:"));
462
463        // Clean up
464        let _ = tokio::fs::remove_dir_all(&dir).await;
465    }
466
467    #[tokio::test]
468    async fn execute_with_base64() {
469        let dir = std::env::temp_dir().join("construct_image_info_b64");
470        let _ = tokio::fs::create_dir_all(&dir).await;
471        let png_path = dir.join("test_b64.png");
472
473        // Minimal 1x1 PNG
474        let png_bytes: Vec<u8> = vec![
475            0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, 0x00, 0x00, 0x00, 0x0D, 0x49, 0x48,
476            0x44, 0x52, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x08, 0x02, 0x00, 0x00,
477            0x00, 0x90, 0x77, 0x53, 0xDE, 0x00, 0x00, 0x00, 0x0C, 0x49, 0x44, 0x41, 0x54, 0x08,
478            0xD7, 0x63, 0xF8, 0xCF, 0xC0, 0x00, 0x00, 0x00, 0x02, 0x00, 0x01, 0xE2, 0x21, 0xBC,
479            0x33, 0x00, 0x00, 0x00, 0x00, 0x49, 0x45, 0x4E, 0x44, 0xAE, 0x42, 0x60, 0x82,
480        ];
481        tokio::fs::write(&png_path, &png_bytes).await.unwrap();
482
483        let tool = ImageInfoTool::new(test_security());
484        let result = tool
485            .execute(json!({"path": png_path.to_string_lossy(), "include_base64": true}))
486            .await
487            .unwrap();
488        assert!(result.success);
489        assert!(result.output.contains("data:image/png;base64,"));
490
491        let _ = tokio::fs::remove_dir_all(&dir).await;
492    }
493}
construct/tools/image_info.rs

construct/tools/
image_info.rs
