Skip to main content

imagegen_bridge_artifacts/
input.rs

1//! Capability-scoped local and inline image input loading.
2
3use std::{
4    io::Read,
5    path::{Component, Path, PathBuf},
6};
7
8use base64::{Engine as _, engine::general_purpose::STANDARD};
9use cap_std::{ambient_authority, fs::Dir};
10use imagegen_bridge_core::{BridgeError, ErrorCode, ImageInput, ImageSource};
11
12use crate::{ImageLimits, ImageMetadata, inspect_image};
13
14/// Loaded and verified image input.
15#[derive(Debug, Clone, PartialEq, Eq)]
16pub struct LoadedImage {
17    /// Encoded image body.
18    pub bytes: Vec<u8>,
19    /// Verified metadata.
20    pub metadata: ImageMetadata,
21    /// Optional safe logical filename supplied by the caller.
22    pub filename: Option<String>,
23}
24
25struct AllowedRoot {
26    canonical_path: PathBuf,
27    dir: Dir,
28}
29
30/// Loads inputs without granting arbitrary filesystem authority.
31pub struct InputLoader {
32    roots: Vec<AllowedRoot>,
33    limits: ImageLimits,
34}
35
36impl std::fmt::Debug for InputLoader {
37    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
38        formatter
39            .debug_struct("InputLoader")
40            .field("root_count", &self.roots.len())
41            .field("limits", &self.limits)
42            .finish()
43    }
44}
45
46impl InputLoader {
47    /// Opens configured roots as directory capabilities.
48    pub fn new(
49        roots: impl IntoIterator<Item = PathBuf>,
50        limits: ImageLimits,
51    ) -> Result<Self, BridgeError> {
52        let roots = roots
53            .into_iter()
54            .map(|root| {
55                let canonical_path = std::fs::canonicalize(&root).map_err(|error| {
56                    input_error(format!("could not open an allowed input root: {error}"))
57                })?;
58                let dir = Dir::open_ambient_dir(&canonical_path, ambient_authority()).map_err(
59                    |error| input_error(format!("could not open an allowed input root: {error}")),
60                )?;
61                Ok(AllowedRoot {
62                    canonical_path,
63                    dir,
64                })
65            })
66            .collect::<Result<Vec<_>, BridgeError>>()?;
67        Ok(Self { roots, limits })
68    }
69
70    /// Loads and verifies one image input.
71    pub fn load(&self, input: &ImageInput) -> Result<LoadedImage, BridgeError> {
72        let bytes = match &input.source {
73            ImageSource::File { path } => self.load_file(path)?,
74            ImageSource::Base64 { data } => self.decode_base64(data)?,
75            ImageSource::DataUrl { data_url } => self.decode_data_url(data_url)?,
76            ImageSource::Url { .. } => {
77                return Err(input_error(
78                    "remote URL input requires an explicitly configured remote fetcher",
79                ));
80            }
81        };
82        let metadata = inspect_image(&bytes, self.limits)?;
83        if let Some(expected) = input.media_type.as_deref() {
84            let actual = media_type(metadata.format);
85            if expected != actual {
86                return Err(input_error(format!(
87                    "declared media type does not match image bytes; detected {actual}"
88                )));
89            }
90        }
91        Ok(LoadedImage {
92            bytes,
93            metadata,
94            filename: input.filename.clone(),
95        })
96    }
97
98    fn load_file(&self, path: &Path) -> Result<Vec<u8>, BridgeError> {
99        if self.roots.is_empty() {
100            return Err(input_error("local file inputs are disabled"));
101        }
102        validate_lexical_path(path)?;
103
104        let mut last_error = None;
105        for root in &self.roots {
106            let relative = if path.is_absolute() {
107                match path.strip_prefix(&root.canonical_path) {
108                    Ok(relative) => relative,
109                    Err(_) => continue,
110                }
111            } else {
112                path
113            };
114            match root.dir.open(relative) {
115                Ok(mut file) => {
116                    let metadata = file.metadata().map_err(|error| {
117                        input_error(format!("could not inspect input: {error}"))
118                    })?;
119                    if !metadata.is_file() {
120                        return Err(input_error("local input is not a regular file"));
121                    }
122                    if metadata.len() > self.limits.max_encoded_bytes {
123                        return Err(input_error("local input exceeds the configured byte limit"));
124                    }
125                    return read_bounded(&mut file, self.limits.max_encoded_bytes);
126                }
127                Err(error) => last_error = Some(error),
128            }
129        }
130        let suffix = last_error
131            .map(|error| format!(": {error}"))
132            .unwrap_or_default();
133        Err(input_error(format!(
134            "local input is outside allowed roots or cannot be opened{suffix}"
135        )))
136    }
137
138    fn decode_base64(&self, data: &str) -> Result<Vec<u8>, BridgeError> {
139        let maximum_encoded = self
140            .limits
141            .max_encoded_bytes
142            .saturating_mul(4)
143            .div_ceil(3)
144            .saturating_add(4);
145        if u64::try_from(data.len()).unwrap_or(u64::MAX) > maximum_encoded {
146            return Err(input_error(
147                "base64 input exceeds the configured byte limit",
148            ));
149        }
150        let decoded = STANDARD
151            .decode(data)
152            .map_err(|_| input_error("base64 input is malformed"))?;
153        if u64::try_from(decoded.len()).unwrap_or(u64::MAX) > self.limits.max_encoded_bytes {
154            return Err(input_error(
155                "decoded input exceeds the configured byte limit",
156            ));
157        }
158        Ok(decoded)
159    }
160
161    fn decode_data_url(&self, data_url: &str) -> Result<Vec<u8>, BridgeError> {
162        let (header, encoded) = data_url
163            .split_once(',')
164            .ok_or_else(|| input_error("data URL is malformed"))?;
165        if !header.starts_with("data:image/") || !header.ends_with(";base64") {
166            return Err(input_error(
167                "data URL must contain a supported image media type and base64 encoding",
168            ));
169        }
170        self.decode_base64(encoded)
171    }
172}
173
174fn validate_lexical_path(path: &Path) -> Result<(), BridgeError> {
175    if path.as_os_str().is_empty()
176        || path
177            .components()
178            .any(|component| matches!(component, Component::ParentDir | Component::Prefix(_)))
179    {
180        return Err(input_error(
181            "local input path contains a forbidden component",
182        ));
183    }
184    Ok(())
185}
186
187fn read_bounded(reader: &mut impl Read, maximum: u64) -> Result<Vec<u8>, BridgeError> {
188    let mut bytes = Vec::new();
189    reader
190        .take(maximum.saturating_add(1))
191        .read_to_end(&mut bytes)
192        .map_err(|error| input_error(format!("could not read local input: {error}")))?;
193    if u64::try_from(bytes.len()).unwrap_or(u64::MAX) > maximum {
194        return Err(input_error("input grew beyond the configured byte limit"));
195    }
196    Ok(bytes)
197}
198
199const fn media_type(format: imagegen_bridge_core::OutputFormat) -> &'static str {
200    match format {
201        imagegen_bridge_core::OutputFormat::Png => "image/png",
202        imagegen_bridge_core::OutputFormat::Jpeg => "image/jpeg",
203        imagegen_bridge_core::OutputFormat::Webp => "image/webp",
204    }
205}
206
207fn input_error(message: impl Into<String>) -> BridgeError {
208    BridgeError::new(ErrorCode::Input, message)
209}
210
211#[cfg(test)]
212mod tests {
213    #![allow(clippy::unwrap_used)]
214
215    use std::fs;
216
217    use super::*;
218    use crate::inspect::test_png;
219
220    fn input(source: ImageSource) -> ImageInput {
221        ImageInput {
222            source,
223            media_type: Some("image/png".to_owned()),
224            filename: Some("input.png".to_owned()),
225        }
226    }
227
228    #[test]
229    fn loads_bounded_base64_and_verifies_type() {
230        let bytes = test_png(2, 2);
231        let loader = InputLoader::new(Vec::new(), ImageLimits::default()).unwrap();
232        let image = loader
233            .load(&input(ImageSource::Base64 {
234                data: STANDARD.encode(&bytes),
235            }))
236            .unwrap();
237        assert_eq!(image.bytes, bytes);
238        assert_eq!(image.metadata.width, 2);
239    }
240
241    #[test]
242    fn loads_file_beneath_allowed_root() {
243        let root = tempfile::tempdir().unwrap();
244        fs::write(root.path().join("input.png"), test_png(2, 3)).unwrap();
245        let loader = InputLoader::new([root.path().to_path_buf()], ImageLimits::default()).unwrap();
246        let image = loader
247            .load(&input(ImageSource::File {
248                path: PathBuf::from("input.png"),
249            }))
250            .unwrap();
251        assert_eq!((image.metadata.width, image.metadata.height), (2, 3));
252    }
253
254    #[test]
255    fn rejects_parent_directory_traversal() {
256        let root = tempfile::tempdir().unwrap();
257        let loader = InputLoader::new([root.path().to_path_buf()], ImageLimits::default()).unwrap();
258        assert!(
259            loader
260                .load(&input(ImageSource::File {
261                    path: PathBuf::from("../escape.png"),
262                }))
263                .is_err()
264        );
265    }
266
267    #[cfg(unix)]
268    #[test]
269    fn rejects_symlink_that_escapes_allowed_root() {
270        use std::os::unix::fs::symlink;
271
272        let root = tempfile::tempdir().unwrap();
273        let outside = tempfile::tempdir().unwrap();
274        fs::write(outside.path().join("outside.png"), test_png(1, 1)).unwrap();
275        symlink(outside.path(), root.path().join("escape")).unwrap();
276        let loader = InputLoader::new([root.path().to_path_buf()], ImageLimits::default()).unwrap();
277        assert!(
278            loader
279                .load(&input(ImageSource::File {
280                    path: PathBuf::from("escape/outside.png"),
281                }))
282                .is_err()
283        );
284    }
285}