vtcode_commons/
image.rs

1//! Image processing utilities
2
3use anyhow::{Context, Result};
4use base64::Engine;
5use std::path::Path;
6
7/// Represents the data from an image file ready for LLM consumption
8#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
9pub struct ImageData {
10    /// Base64-encoded image data
11    pub base64_data: String,
12
13    /// MIME type of the image (e.g., "image/png", "image/jpeg")
14    pub mime_type: String,
15
16    /// Original file path or URL
17    pub file_path: String,
18
19    /// File size in bytes
20    pub size: u64,
21}
22
23/// Detects MIME type from Content-Type header
24pub fn detect_mime_type_from_content_type(content_type: &str) -> Option<String> {
25    let content_type = content_type.to_lowercase();
26    if content_type.starts_with("image/png") {
27        Some("image/png".to_string())
28    } else if content_type.starts_with("image/jpeg") || content_type.starts_with("image/jpg") {
29        Some("image/jpeg".to_string())
30    } else if content_type.starts_with("image/gif") {
31        Some("image/gif".to_string())
32    } else if content_type.starts_with("image/webp") {
33        Some("image/webp".to_string())
34    } else if content_type.starts_with("image/bmp") {
35        Some("image/bmp".to_string())
36    } else if content_type.starts_with("image/tiff") || content_type.starts_with("image/tif") {
37        Some("image/tiff".to_string())
38    } else if content_type.starts_with("image/svg") {
39        Some("image/svg+xml".to_string())
40    } else {
41        None
42    }
43}
44
45/// Detects MIME type from file data (magic bytes)
46pub fn detect_mime_type_from_data(data: &[u8]) -> String {
47    // JPEG magic bytes: starts with FF D8
48    if data.len() >= 2 && data[0] == 0xFF && data[1] == 0xD8 {
49        return "image/jpeg".to_string();
50    }
51
52    // Need at least 8 bytes for other formats
53    if data.len() < 8 {
54        return "image/png".to_string();
55    }
56
57    match &data[..8] {
58        [0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A] => "image/png".to_string(),
59        [0x47, 0x49, 0x46, 0x38, _, _, _, _] => {
60            if data.len() >= 12 && &data[8..12] == b"WEBP" {
61                "image/webp".to_string()
62            } else {
63                "image/gif".to_string()
64            }
65        }
66        [0x52, 0x49, 0x46, 0x46, _, _, _, _] => {
67            if data.len() >= 12 && &data[8..12] == b"WEBP" {
68                "image/webp".to_string()
69            } else {
70                "image/png".to_string()
71            }
72        }
73        [0x42, 0x4D, _, _] => "image/bmp".to_string(),
74        _ => "image/png".to_string(),
75    }
76}
77
78/// Detects the MIME type based on file extension
79pub fn detect_mime_type_from_extension(path: &Path) -> Result<String> {
80    let extension = path
81        .extension()
82        .and_then(|ext| ext.to_str())
83        .unwrap_or("")
84        .to_lowercase();
85
86    let mime_type = match extension.as_str() {
87        "png" => "image/png",
88        "jpg" | "jpeg" => "image/jpeg",
89        "gif" => "image/gif",
90        "webp" => "image/webp",
91        "bmp" => "image/bmp",
92        "tiff" | "tif" => "image/tiff",
93        "svg" => "image/svg+xml",
94        _ => return Err(anyhow::anyhow!("Unsupported image format: {extension}")),
95    };
96
97    Ok(mime_type.to_string())
98}
99
100/// Validates that the image file path has a supported extension
101pub fn has_supported_image_extension(path: &Path) -> bool {
102    let extension = path
103        .extension()
104        .and_then(|ext| ext.to_str())
105        .unwrap_or("")
106        .to_lowercase();
107
108    const VALID_EXTENSIONS: &[&str] = &["png", "jpg", "jpeg", "gif", "webp", "bmp", "tiff", "svg"];
109    VALID_EXTENSIONS.contains(&extension.as_str())
110}
111
112/// Encodes binary data to base64
113pub fn encode_to_base64(data: &[u8]) -> String {
114    base64::engine::general_purpose::STANDARD.encode(data)
115}
116
117/// Reads an image file from the local filesystem and converts it to base64 format.
118///
119/// Validates the path for traversal attacks and checks the file extension
120/// against a supported set. Max file size is 20 MB.
121pub async fn read_image_file<P: AsRef<Path>>(file_path: P) -> Result<ImageData> {
122    use crate::paths::is_safe_relative_path;
123
124    let path = file_path.as_ref();
125
126    if !is_safe_relative_path(&path.to_string_lossy()) {
127        return Err(anyhow::anyhow!(
128            "Unsafe or traversal detected in image path: {}",
129            path.display()
130        ));
131    }
132
133    if !has_supported_image_extension(path) {
134        return Err(anyhow::anyhow!(
135            "Unsupported image extension for path: {}",
136            path.display()
137        ));
138    }
139
140    let file_contents = tokio::fs::read(path)
141        .await
142        .with_context(|| format!("Failed to read image file: {}", path.display()))?;
143
144    if file_contents.len() > 20 * 1024 * 1024 {
145        return Err(anyhow::anyhow!(
146            "Image file too large: {} bytes (max 20MB)",
147            file_contents.len()
148        ));
149    }
150
151    let mime_type = detect_mime_type_from_extension(path)?;
152    let base64_data = encode_to_base64(&file_contents);
153
154    Ok(ImageData {
155        base64_data,
156        mime_type,
157        file_path: path.display().to_string(),
158        size: file_contents.len() as u64,
159    })
160}
161
162/// Reads an image file from an absolute path (or already validated path) and
163/// converts it to base64.
164///
165/// This skips relative-path safety checks and should only be used when the
166/// caller has already validated the path scope and intent.
167pub async fn read_image_file_any_path<P: AsRef<Path>>(file_path: P) -> Result<ImageData> {
168    let path = file_path.as_ref();
169
170    if !has_supported_image_extension(path) {
171        return Err(anyhow::anyhow!(
172            "Unsupported image extension for path: {}",
173            path.display()
174        ));
175    }
176
177    let file_contents = tokio::fs::read(path)
178        .await
179        .with_context(|| format!("Failed to read image file: {}", path.display()))?;
180
181    if file_contents.len() > 20 * 1024 * 1024 {
182        return Err(anyhow::anyhow!(
183            "Image file too large: {} bytes (max 20MB)",
184            file_contents.len()
185        ));
186    }
187
188    let mime_type = detect_mime_type_from_extension(path)?;
189    let base64_data = encode_to_base64(&file_contents);
190
191    Ok(ImageData {
192        base64_data,
193        mime_type,
194        file_path: path.display().to_string(),
195        size: file_contents.len() as u64,
196    })
197}
vtcode_commons/image.rs

vtcode_commons/
image.rs
