atman_runtime/
attachment_store.rs

1use std::io::Write;
2use std::path::{Path, PathBuf};
3
4use base64::Engine;
5
6use crate::error::RuntimeError;
7use crate::message::{ImageData, ImageSource};
8
9const ATTACHMENTS_DIR: &str = "attachments";
10const MAX_IMAGE_BYTES: usize = 20 * 1024 * 1024;
11
12#[derive(Debug, Clone)]
13pub struct AttachmentStore {
14    root: PathBuf,
15    persistent: bool,
16}
17
18impl AttachmentStore {
19    pub fn at(session_dir: impl AsRef<Path>) -> Self {
20        let session_dir = session_dir.as_ref();
21        Self {
22            root: session_dir.join(ATTACHMENTS_DIR),
23            persistent: !session_dir.as_os_str().is_empty(),
24        }
25    }
26
27    pub fn import_path(&self, path: impl AsRef<Path>) -> Result<ImageSource, RuntimeError> {
28        let path = path.as_ref();
29        let bytes = std::fs::read(path).map_err(|error| RuntimeError::AttachmentError {
30            reason: format!("cannot read {}: {error}", path.display()),
31        })?;
32        self.import_bytes(&bytes, path.file_name().and_then(|name| name.to_str()))
33    }
34
35    pub fn import_base64(
36        &self,
37        data: &str,
38        name: Option<&str>,
39    ) -> Result<ImageSource, RuntimeError> {
40        let bytes = base64::engine::general_purpose::STANDARD
41            .decode(data)
42            .map_err(|error| RuntimeError::AttachmentError {
43                reason: format!("invalid base64 image: {error}"),
44            })?;
45        self.import_bytes(&bytes, name)
46    }
47
48    pub fn import_bytes(
49        &self,
50        bytes: &[u8],
51        name: Option<&str>,
52    ) -> Result<ImageSource, RuntimeError> {
53        validate_size(bytes)?;
54        let (media_type, extension) = detect_image_type(bytes)?;
55        if !self.persistent {
56            return Ok(ImageSource {
57                media_type: media_type.into(),
58                data: ImageData::Base64 {
59                    data: base64::engine::general_purpose::STANDARD.encode(bytes),
60                },
61                detail: crate::provider::ImageDetail::Auto,
62            });
63        }
64        let id = blake3::hash(bytes).to_hex().to_string();
65        let path = self.root.join(format!("{id}.{extension}"));
66        if !path.is_file() {
67            std::fs::create_dir_all(&self.root).map_err(|error| RuntimeError::AttachmentError {
68                reason: format!("cannot create attachment store: {error}"),
69            })?;
70            let temp_path = self
71                .root
72                .join(format!(".{id}.{}.tmp", uuid::Uuid::new_v4()));
73            let write_result = (|| -> std::io::Result<()> {
74                let mut file = std::fs::OpenOptions::new()
75                    .create_new(true)
76                    .write(true)
77                    .open(&temp_path)?;
78                file.write_all(bytes)?;
79                file.sync_all()?;
80                std::fs::rename(&temp_path, &path)
81            })();
82            if let Err(error) = write_result {
83                let _ = std::fs::remove_file(&temp_path);
84                if !path.is_file() {
85                    return Err(RuntimeError::AttachmentError {
86                        reason: format!("cannot persist attachment: {error}"),
87                    });
88                }
89            }
90        }
91        Ok(ImageSource {
92            media_type: media_type.into(),
93            data: ImageData::Artifact {
94                id,
95                path,
96                name: name.map(ToOwned::to_owned),
97            },
98            detail: crate::provider::ImageDetail::Auto,
99        })
100    }
101}
102
103pub fn image_bytes(source: &ImageSource) -> Result<Vec<u8>, RuntimeError> {
104    let bytes = match &source.data {
105        ImageData::Base64 { data } => base64::engine::general_purpose::STANDARD
106            .decode(data)
107            .map_err(|error| RuntimeError::AttachmentError {
108                reason: format!("invalid base64 image: {error}"),
109            })?,
110        ImageData::Path { path } | ImageData::Artifact { path, .. } => std::fs::read(path)
111            .map_err(|error| RuntimeError::AttachmentError {
112                reason: format!("cannot read {}: {error}", path.display()),
113            })?,
114    };
115    validate_size(&bytes)?;
116    let (actual_media_type, _) = detect_image_type(&bytes)?;
117    if source.media_type != actual_media_type {
118        return Err(RuntimeError::AttachmentError {
119            reason: format!(
120                "image media type mismatch: declared {}, detected {actual_media_type}",
121                source.media_type
122            ),
123        });
124    }
125    if let ImageData::Artifact { id, .. } = &source.data {
126        let actual_id = blake3::hash(&bytes).to_hex().to_string();
127        if actual_id != *id {
128            return Err(RuntimeError::AttachmentError {
129                reason: format!("attachment integrity check failed for {id}"),
130            });
131        }
132    }
133    Ok(bytes)
134}
135
136pub fn image_base64(source: &ImageSource) -> Result<String, RuntimeError> {
137    if let ImageData::Base64 { data } = &source.data {
138        image_bytes(source)?;
139        return Ok(data.clone());
140    }
141    Ok(base64::engine::general_purpose::STANDARD.encode(image_bytes(source)?))
142}
143
144pub fn display_name(source: &ImageSource) -> String {
145    match &source.data {
146        ImageData::Artifact { id, name, .. } => name.clone().unwrap_or_else(|| id.clone()),
147        ImageData::Path { path } => path
148            .file_name()
149            .and_then(|name| name.to_str())
150            .unwrap_or("image")
151            .to_string(),
152        ImageData::Base64 { .. } => "image".into(),
153    }
154}
155
156fn validate_size(bytes: &[u8]) -> Result<(), RuntimeError> {
157    if bytes.is_empty() {
158        return Err(RuntimeError::AttachmentError {
159            reason: "image is empty".into(),
160        });
161    }
162    if bytes.len() > MAX_IMAGE_BYTES {
163        return Err(RuntimeError::AttachmentError {
164            reason: format!(
165                "image is too large: {} bytes exceeds the {} byte limit",
166                bytes.len(),
167                MAX_IMAGE_BYTES
168            ),
169        });
170    }
171    Ok(())
172}
173
174fn detect_image_type(bytes: &[u8]) -> Result<(&'static str, &'static str), RuntimeError> {
175    let detected = if bytes.starts_with(b"\x89PNG\r\n\x1a\n") {
176        Some(("image/png", "png"))
177    } else if bytes.starts_with(&[0xff, 0xd8, 0xff]) {
178        Some(("image/jpeg", "jpg"))
179    } else if bytes.starts_with(b"GIF87a") || bytes.starts_with(b"GIF89a") {
180        Some(("image/gif", "gif"))
181    } else if bytes.len() >= 12 && &bytes[..4] == b"RIFF" && &bytes[8..12] == b"WEBP" {
182        Some(("image/webp", "webp"))
183    } else {
184        None
185    };
186    detected.ok_or_else(|| RuntimeError::AttachmentError {
187        reason: "unsupported image format; expected PNG, JPEG, GIF, or WebP".into(),
188    })
189}
190
191#[cfg(test)]
192mod tests {
193    use super::*;
194
195    const PNG_1X1: &[u8] = &[
196        0x89, b'P', b'N', b'G', 0x0d, 0x0a, 0x1a, 0x0a, 0x00, 0x00, 0x00, 0x0d,
197    ];
198
199    #[test]
200    fn import_is_content_addressed_and_deduplicated() {
201        let session = tempfile::tempdir().unwrap();
202        let store = AttachmentStore::at(session.path());
203        let first = store.import_bytes(PNG_1X1, Some("first.png")).unwrap();
204        let second = store.import_bytes(PNG_1X1, Some("second.png")).unwrap();
205
206        let ImageData::Artifact {
207            id: first_id,
208            path: first_path,
209            ..
210        } = first.data
211        else {
212            panic!("expected artifact")
213        };
214        let ImageData::Artifact {
215            id: second_id,
216            path: second_path,
217            ..
218        } = second.data
219        else {
220            panic!("expected artifact")
221        };
222        assert_eq!(first_id, second_id);
223        assert_eq!(first_path, second_path);
224        assert_eq!(std::fs::read(first_path).unwrap(), PNG_1X1);
225    }
226
227    #[test]
228    fn integrity_mismatch_is_rejected() {
229        let session = tempfile::tempdir().unwrap();
230        let store = AttachmentStore::at(session.path());
231        let source = store.import_bytes(PNG_1X1, None).unwrap();
232        let ImageData::Artifact { path, .. } = &source.data else {
233            panic!("expected artifact")
234        };
235        std::fs::write(path, b"not an image").unwrap();
236        assert!(image_bytes(&source).is_err());
237    }
238}
atman_runtime/attachment_store.rs

atman_runtime/
attachment_store.rs
