Skip to main content

aft/github_read/
attachments.rs

1use std::io::Read;
2use std::sync::LazyLock;
3use std::time::Duration;
4
5use reqwest::blocking::Client;
6use reqwest::header::{CONTENT_TYPE, LOCATION};
7use reqwest::redirect::Policy;
8use url::Url;
9
10/// At most this many GitHub-hosted images become attachments for one read.
11/// Keeping the count fixed prevents a long issue thread from creating an
12/// unbounded transport payload.
13pub const MAX_GITHUB_IMAGE_ATTACHMENTS: usize = 8;
14/// Attachments for one read may contain at most this many downloaded bytes in
15/// total. A candidate that exceeds the remaining budget is dropped whole, so
16/// callers never receive a partial image.
17pub const MAX_GITHUB_IMAGE_ATTACHMENT_BYTES: usize = 8 * 1024 * 1024;
18const MAX_IMAGE_REDIRECTS: usize = 5;
19const IMAGE_CONNECT_TIMEOUT: Duration = Duration::from_secs(15);
20const IMAGE_REQUEST_TIMEOUT: Duration = Duration::from_secs(45);
21
22static IMAGE_URL: LazyLock<regex::Regex> = LazyLock::new(|| {
23    regex::Regex::new(r#"https://[^\s<>\"'()\[\]]+"#)
24        .expect("image URL regular expression is valid")
25});
26
27/// One complete, validated image attachment. Text rendering retains the source
28/// URL; this is out-of-band data for a caller that explicitly supports vision.
29#[derive(Clone, Debug, Eq, PartialEq)]
30pub struct GithubImageAttachment {
31    pub source_url: String,
32    pub mime: String,
33    pub bytes: Vec<u8>,
34}
35
36/// A downloader result whose final URL is retained for redirect validation.
37#[derive(Clone, Debug, Eq, PartialEq)]
38pub struct DownloadedGithubImage {
39    pub final_url: Url,
40    pub mime: String,
41    pub bytes: Vec<u8>,
42}
43
44/// Synchronous downloader used inside a deferred worker. Tests can replace it
45/// with a deterministic fixture downloader that records URLs and byte budgets.
46pub trait GithubImageDownloader: Send + Sync {
47    fn download(
48        &self,
49        url: &Url,
50        maximum_bytes: usize,
51    ) -> Result<Option<DownloadedGithubImage>, String>;
52}
53
54/// Production GitHub-image downloader. Redirect following is manual so every
55/// target receives the same allowlist check as the original URL.
56#[derive(Default)]
57pub struct ReqwestGithubImageDownloader;
58
59impl GithubImageDownloader for ReqwestGithubImageDownloader {
60    fn download(
61        &self,
62        url: &Url,
63        maximum_bytes: usize,
64    ) -> Result<Option<DownloadedGithubImage>, String> {
65        if !is_allowed_github_image_url(url) || maximum_bytes == 0 {
66            return Ok(None);
67        }
68        let tls = crate::platform_tls::client_config()
69            .map_err(|error| format!("failed to configure GitHub image TLS: {error}"))?;
70        let client = Client::builder()
71            .redirect(Policy::none())
72            .connect_timeout(IMAGE_CONNECT_TIMEOUT)
73            .timeout(IMAGE_REQUEST_TIMEOUT)
74            .use_preconfigured_tls(tls)
75            .build()
76            .map_err(|error| format!("failed to build GitHub image client: {error}"))?;
77        let mut current = url.clone();
78
79        for redirects in 0..=MAX_IMAGE_REDIRECTS {
80            if !is_allowed_github_image_url(&current) {
81                return Ok(None);
82            }
83            let mut response = client
84                .get(current.clone())
85                .send()
86                .map_err(|error| format!("image download failed: {error}"))?;
87            if response.status().is_redirection() {
88                if redirects == MAX_IMAGE_REDIRECTS {
89                    return Ok(None);
90                }
91                let Some(location) = response
92                    .headers()
93                    .get(LOCATION)
94                    .and_then(|value| value.to_str().ok())
95                else {
96                    return Ok(None);
97                };
98                let next = current.join(location).map_err(|error| {
99                    format!("GitHub image redirect location is invalid: {error}")
100                })?;
101                if !is_allowed_github_image_url(&next) {
102                    return Ok(None);
103                }
104                current = next;
105                continue;
106            }
107            if !response.status().is_success() {
108                return Ok(None);
109            }
110            let mime = response
111                .headers()
112                .get(CONTENT_TYPE)
113                .and_then(|value| value.to_str().ok())
114                .map(|value| {
115                    value
116                        .split(';')
117                        .next()
118                        .unwrap_or(value)
119                        .trim()
120                        .to_ascii_lowercase()
121                })
122                .filter(|value| supported_image_mime(value))
123                .ok_or_else(|| "GitHub image response had an unsupported media type".to_string())?;
124            if response
125                .content_length()
126                .is_some_and(|length| length > maximum_bytes as u64)
127            {
128                return Ok(None);
129            }
130            let mut bytes = Vec::new();
131            let mut limited = response
132                .by_ref()
133                .take(maximum_bytes.saturating_add(1) as u64);
134            limited
135                .read_to_end(&mut bytes)
136                .map_err(|error| format!("failed to read GitHub image: {error}"))?;
137            if bytes.len() > maximum_bytes {
138                return Ok(None);
139            }
140            return Ok(Some(DownloadedGithubImage {
141                final_url: current,
142                mime,
143                bytes,
144            }));
145        }
146        Ok(None)
147    }
148}
149
150/// Discover eligible HTTPS GitHub image candidates in textual document order.
151/// Repeated URLs remain repeated candidates: their positions in the canonical
152/// document determine which attachment claims a count slot first.
153pub fn discover_github_image_urls(canonical_text: &str) -> Vec<Url> {
154    IMAGE_URL
155        .find_iter(canonical_text)
156        .filter_map(|capture| {
157            let raw = capture.as_str().trim_end_matches(|character: char| {
158                matches!(character, '.' | ',' | ';' | ':' | '!' | '?')
159            });
160            Url::parse(raw).ok()
161        })
162        .filter(is_allowed_github_image_url)
163        .collect()
164}
165
166/// Download complete attachment candidates without changing canonical text.
167/// This function is capability-agnostic so tests can call it directly; callers
168/// with a missing or false vision capability must not invoke it.
169pub fn download_github_image_attachments(
170    canonical_text: &str,
171    downloader: &dyn GithubImageDownloader,
172) -> Vec<GithubImageAttachment> {
173    let mut attachments = Vec::new();
174    let mut consumed_bytes = 0usize;
175    for url in discover_github_image_urls(canonical_text) {
176        if attachments.len() == MAX_GITHUB_IMAGE_ATTACHMENTS {
177            break;
178        }
179        let remaining = MAX_GITHUB_IMAGE_ATTACHMENT_BYTES.saturating_sub(consumed_bytes);
180        if remaining == 0 {
181            break;
182        }
183        let Ok(Some(downloaded)) = downloader.download(&url, remaining) else {
184            continue;
185        };
186        // Fixture downloaders are untrusted too: they must not let tests or a
187        // future implementation bypass final redirect verification or budgets.
188        if !is_allowed_github_image_url(&downloaded.final_url)
189            || !supported_image_mime(&downloaded.mime)
190            || !supported_image_bytes(&downloaded.bytes)
191            || downloaded.bytes.len() > remaining
192        {
193            continue;
194        }
195        consumed_bytes += downloaded.bytes.len();
196        attachments.push(GithubImageAttachment {
197            source_url: url.to_string(),
198            mime: downloaded.mime,
199            bytes: downloaded.bytes,
200        });
201    }
202    attachments
203}
204
205/// The only hosts allowed to become vision attachments. `github.com` is
206/// limited to `/user-attachments/`; an arbitrary GitHub page is not an image
207/// source and must not turn this feature into a generic web fetcher.
208pub fn is_allowed_github_image_url(url: &Url) -> bool {
209    url.scheme() == "https"
210        && url.username().is_empty()
211        && url.password().is_none()
212        && matches!(url.port(), None | Some(443))
213        && match url.host_str() {
214            Some("user-images.githubusercontent.com") => true,
215            Some("github.com") => url.path().starts_with("/user-attachments/"),
216            _ => false,
217        }
218}
219
220fn supported_image_mime(mime: &str) -> bool {
221    matches!(
222        mime,
223        "image/png" | "image/jpeg" | "image/gif" | "image/webp"
224    )
225}
226
227fn supported_image_bytes(bytes: &[u8]) -> bool {
228    matches!(
229        image::guess_format(bytes),
230        Ok(image::ImageFormat::Png
231            | image::ImageFormat::Jpeg
232            | image::ImageFormat::Gif
233            | image::ImageFormat::WebP)
234    )
235}
236
237#[cfg(test)]
238mod tests {
239    use std::sync::Mutex;
240
241    use super::*;
242
243    #[derive(Default)]
244    struct FixtureDownloader {
245        calls: Mutex<Vec<(String, usize)>>,
246    }
247
248    impl GithubImageDownloader for FixtureDownloader {
249        fn download(
250            &self,
251            url: &Url,
252            maximum_bytes: usize,
253        ) -> Result<Option<DownloadedGithubImage>, String> {
254            self.calls
255                .lock()
256                .unwrap()
257                .push((url.to_string(), maximum_bytes));
258            Ok(Some(DownloadedGithubImage {
259                final_url: url.clone(),
260                mime: "image/png".to_string(),
261                bytes: vec![137, 80, 78, 71, 13, 10, 26, 10],
262            }))
263        }
264    }
265
266    #[test]
267    fn attachment_urls_are_allowlisted_in_document_order() {
268        let source = concat!(
269            "https://example.test/nope.png ",
270            "https://github.com/user-attachments/files/1/a.png ",
271            "https://user-images.githubusercontent.com/2/b.png"
272        );
273        let urls = discover_github_image_urls(source);
274        assert_eq!(urls.len(), 2);
275        assert_eq!(urls[0].host_str(), Some("github.com"));
276        assert_eq!(
277            urls[1].host_str(),
278            Some("user-images.githubusercontent.com")
279        );
280    }
281
282    #[test]
283    fn downloader_receives_only_complete_allowlisted_attachments() {
284        let downloader = FixtureDownloader::default();
285        let attachments = download_github_image_attachments(
286            "https://user-images.githubusercontent.com/2/b.png",
287            &downloader,
288        );
289        assert_eq!(attachments.len(), 1);
290        assert_eq!(attachments[0].bytes, vec![137, 80, 78, 71, 13, 10, 26, 10]);
291        assert_eq!(downloader.calls.lock().unwrap().len(), 1);
292    }
293}