systemprompt_api/services/gateway/image_fetch/
mod.rs1mod guard;
20
21use base64::Engine as _;
22use base64::engine::general_purpose::STANDARD as BASE64;
23use systemprompt_models::net::{HTTP_CONNECT_TIMEOUT, trusted_http_hosts_from_env};
24
25use super::protocol::canonical::{CanonicalContent, CanonicalRequest, ImageSource};
26
27pub const MAX_IMAGE_BYTES: usize = 5 * 1024 * 1024;
32
33pub const ACCEPTED_MIME: [&str; 5] = [
37 "image/png",
38 "image/jpeg",
39 "image/webp",
40 "image/heic",
41 "image/heif",
42];
43
44const MAX_REDIRECTS: u8 = 3;
45const FETCH_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
46
47#[derive(Debug, thiserror::Error)]
53#[error("image url {url} could not be inlined: {message}")]
54pub struct ImageFetchFailed {
55 pub url: String,
56 pub message: String,
57 pub caller_fault: bool,
58}
59
60#[derive(Debug, Clone)]
64pub struct ImageFetchPolicy {
65 pub timeout: std::time::Duration,
66 pub max_bytes: usize,
67 pub max_redirects: u8,
68 pub trusted_hosts: Vec<String>,
69}
70
71impl Default for ImageFetchPolicy {
72 fn default() -> Self {
73 Self {
74 timeout: FETCH_TIMEOUT,
75 max_bytes: MAX_IMAGE_BYTES,
76 max_redirects: MAX_REDIRECTS,
77 trusted_hosts: trusted_http_hosts_from_env(),
78 }
79 }
80}
81
82#[derive(Debug, Clone)]
84pub struct InlineImage {
85 pub media_type: String,
86 pub base64: String,
87}
88
89fn client() -> &'static reqwest::Client {
93 static CLIENT: std::sync::OnceLock<reqwest::Client> = std::sync::OnceLock::new();
94 CLIENT.get_or_init(|| {
95 reqwest::Client::builder()
96 .redirect(reqwest::redirect::Policy::none())
97 .connect_timeout(HTTP_CONNECT_TIMEOUT)
98 .build()
99 .unwrap_or_default()
100 })
101}
102
103pub async fn inline_url_images(
108 request: &mut CanonicalRequest,
109 policy: &ImageFetchPolicy,
110) -> Result<usize, ImageFetchFailed> {
111 let mut count = 0usize;
112 for message in &mut request.messages {
113 for content in &mut message.content {
114 let CanonicalContent::Image(ImageSource::Url { url, detail }) = content else {
115 continue;
116 };
117 let fetched = fetch(url, policy).await?;
118 *content = CanonicalContent::Image(ImageSource::Base64 {
119 media_type: fetched.media_type,
120 data: fetched.base64,
121 detail: *detail,
122 });
123 count += 1;
124 }
125 }
126 Ok(count)
127}
128
129pub async fn fetch(url: &str, policy: &ImageFetchPolicy) -> Result<InlineImage, ImageFetchFailed> {
133 let fail = |message: String, caller_fault: bool| ImageFetchFailed {
134 url: url.to_owned(),
135 message,
136 caller_fault,
137 };
138 tokio::time::timeout(policy.timeout, fetch_inner(url, policy))
139 .await
140 .map_or_else(
141 |_| Err(fail(format!("fetch exceeded {:?}", policy.timeout), false)),
142 |result| result.map_err(|(message, caller_fault)| fail(message, caller_fault)),
143 )
144}
145
146type FetchError = (String, bool);
147
148async fn fetch_inner(url: &str, policy: &ImageFetchPolicy) -> Result<InlineImage, FetchError> {
149 let mut next = guard::checked_url(url, &policy.trusted_hosts)
150 .await
151 .map_err(|e| (e, true))?;
152 for _ in 0..=policy.max_redirects {
153 let response = client()
154 .get(next.clone())
155 .send()
156 .await
157 .map_err(|e| (format!("request failed: {e}"), false))?;
158 if let Some(location) = redirect_target(&response) {
159 let joined = next
160 .join(&location)
161 .map_err(|e| (format!("invalid redirect target: {e}"), true))?;
162 next = guard::checked_url(joined.as_str(), &policy.trusted_hosts)
163 .await
164 .map_err(|e| (format!("redirect rejected: {e}"), true))?;
165 continue;
166 }
167 return read_image(response, policy).await;
168 }
169 Err((
170 format!("more than {} redirects", policy.max_redirects),
171 true,
172 ))
173}
174
175fn redirect_target(response: &reqwest::Response) -> Option<String> {
176 if !response.status().is_redirection() {
177 return None;
178 }
179 response
180 .headers()
181 .get(reqwest::header::LOCATION)
182 .and_then(|v| v.to_str().ok())
183 .map(ToOwned::to_owned)
184}
185
186async fn read_image(
187 mut response: reqwest::Response,
188 policy: &ImageFetchPolicy,
189) -> Result<InlineImage, FetchError> {
190 let status = response.status();
191 if !status.is_success() {
192 return Err((format!("host returned {status}"), true));
193 }
194 let media_type = declared_mime(&response)?;
195 let mut body: Vec<u8> = Vec::new();
199 while let Some(chunk) = response
200 .chunk()
201 .await
202 .map_err(|e| (format!("read failed: {e}"), false))?
203 {
204 if body.len() + chunk.len() > policy.max_bytes {
205 return Err((format!("larger than {} bytes", policy.max_bytes), true));
206 }
207 body.extend_from_slice(&chunk);
208 }
209 if body.is_empty() {
210 return Err(("empty response body".to_owned(), true));
211 }
212 Ok(InlineImage {
213 media_type,
214 base64: BASE64.encode(&body),
215 })
216}
217
218fn declared_mime(response: &reqwest::Response) -> Result<String, FetchError> {
219 let raw = response
220 .headers()
221 .get(reqwest::header::CONTENT_TYPE)
222 .and_then(|v| v.to_str().ok())
223 .ok_or_else(|| ("no content-type".to_owned(), true))?;
224 let mime = raw
225 .split(';')
226 .next()
227 .unwrap_or_default()
228 .trim()
229 .to_ascii_lowercase();
230 if ACCEPTED_MIME.contains(&mime.as_str()) {
231 return Ok(mime);
232 }
233 Err((
234 format!("content-type {mime} is not an inlineable image"),
235 true,
236 ))
237}