imagegen_bridge_artifacts/
input.rs1use 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#[derive(Debug, Clone, PartialEq, Eq)]
16pub struct LoadedImage {
17 pub bytes: Vec<u8>,
19 pub metadata: ImageMetadata,
21 pub filename: Option<String>,
23}
24
25struct AllowedRoot {
26 canonical_path: PathBuf,
27 dir: Dir,
28}
29
30pub 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 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 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}