Skip to main content

aft/
url_fetch.rs

1use std::error::Error;
2use std::fmt;
3use std::fs;
4use std::io::{self, Read, Write};
5use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, ToSocketAddrs};
6use std::path::{Path, PathBuf};
7use std::sync::{mpsc, Arc};
8use std::thread;
9use std::time::{Duration, SystemTime, UNIX_EPOCH};
10
11use htmd::{
12    element_handler::{HandlerResult, Handlers},
13    Element, HtmlToMarkdown,
14};
15use reqwest::blocking::{Client, Response as HttpResponse};
16use reqwest::header::{ACCEPT, CONTENT_TYPE, LOCATION, USER_AGENT};
17use reqwest::redirect::Policy;
18use serde::{Deserialize, Serialize};
19use sha2::{Digest, Sha256};
20use url::Url;
21
22use crate::parser::detect_language;
23
24const MAX_RESPONSE_BYTES: u64 = 10 * 1024 * 1024;
25const CACHE_TTL_MS: u64 = 24 * 60 * 60 * 1000;
26const CONNECT_TIMEOUT: Duration = Duration::from_millis(30_000);
27const BODY_CHUNK_TIMEOUT: Duration = Duration::from_millis(15_000);
28const MAX_REDIRECTS: usize = 5;
29
30/// Retry budget for transient connect/transport failures only. Agents
31/// shouldn't have to retry manually for a single TCP/TLS hiccup. We cap
32/// at 2 retries (= 3 total attempts) with short jittered backoff so a
33/// genuinely-broken host fails fast instead of dragging the foreground
34/// fetch out to many seconds.
35///
36/// We deliberately do NOT retry on:
37///   - HTTP error status (4xx/5xx) — the server actually answered
38///   - Redirect errors / SSRF rejections — those are deterministic
39///   - Body read stalls — already handled by BODY_CHUNK_TIMEOUT
40const TRANSIENT_RETRY_ATTEMPTS: usize = 2;
41const TRANSIENT_RETRY_BACKOFFS_MS: [u64; TRANSIENT_RETRY_ATTEMPTS] = [200, 600];
42const ACCEPT_HEADER: &str = "application/vnd.github.raw, text/markdown, text/x-markdown, text/html;q=0.9, application/json;q=0.8, text/plain;q=0.5";
43const USER_AGENT_VALUE: &str = "aft-opencode-plugin";
44const CONVERTED_MARKDOWN_CONTENT_TYPE: &str = "text/markdown; charset=utf-8";
45
46#[derive(Clone, Default)]
47pub struct UrlFetchOptions {
48    pub allow_private: bool,
49    /// Test hook: treat a hostname as resolving to these IPs during SSRF validation.
50    /// Production callers leave this empty and use `std::net::ToSocketAddrs`.
51    #[doc(hidden)]
52    pub public_host_overrides: Vec<(String, Vec<IpAddr>)>,
53    /// Test hook: force reqwest to connect a hostname to a local mock server while
54    /// SSRF validation still sees `public_host_overrides` above.
55    #[doc(hidden)]
56    pub connect_overrides: Vec<(String, SocketAddr)>,
57    /// Test hook: observes the temp path immediately before the atomic rename.
58    #[doc(hidden)]
59    pub atomic_write_observer: Option<Arc<dyn Fn(&Path, &Path) + Send + Sync>>,
60}
61
62#[derive(Debug, Clone)]
63pub struct UrlFetchError {
64    message: String,
65}
66
67impl UrlFetchError {
68    fn new(message: impl Into<String>) -> Self {
69        Self {
70            message: message.into(),
71        }
72    }
73}
74
75impl fmt::Display for UrlFetchError {
76    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
77        f.write_str(&self.message)
78    }
79}
80
81impl std::error::Error for UrlFetchError {}
82
83#[derive(Debug, Serialize, Deserialize)]
84struct CacheMeta {
85    url: String,
86    #[serde(rename = "contentType")]
87    content_type: String,
88    extension: String,
89    #[serde(rename = "fetchedAt")]
90    fetched_at: u64,
91}
92
93pub fn is_http_url(value: &str) -> bool {
94    value.starts_with("http://") || value.starts_with("https://")
95}
96
97pub fn fetch_url_to_cache(
98    url: &str,
99    storage_dir: &Path,
100    options: UrlFetchOptions,
101) -> Result<PathBuf, UrlFetchError> {
102    let parsed = Url::parse(url).map_err(|_| UrlFetchError::new(format!("Invalid URL: {url}")))?;
103    let fetch_url = rewrite_url_for_fetch(&parsed);
104    validate_public_url(&fetch_url, &options)?;
105
106    let dir = cache_dir(storage_dir);
107    fs::create_dir_all(&dir).map_err(|error| {
108        UrlFetchError::new(format!(
109            "Failed to create URL cache directory {}: {error}",
110            dir.display()
111        ))
112    })?;
113
114    let hash = hash_url(url);
115    let meta_file = meta_path(storage_dir, &hash);
116    if let Some(cached) = fresh_cached_path(storage_dir, &hash, &meta_file, &fetch_url)? {
117        return Ok(cached);
118    }
119
120    let response = fetch_with_redirects(&fetch_url, url, &options)?;
121    if !response.status().is_success() {
122        return Err(UrlFetchError::new(format!(
123            "HTTP {} {} fetching {url}",
124            response.status().as_u16(),
125            response.status().canonical_reason().unwrap_or("")
126        )));
127    }
128
129    let content_type = response
130        .headers()
131        .get(CONTENT_TYPE)
132        .and_then(|value| value.to_str().ok())
133        .unwrap_or("text/plain")
134        .to_string();
135    let (extension, from_source_path) =
136        resolve_fetch_extension(&fetch_url, &content_type).ok_or_else(|| {
137            UrlFetchError::new(format!(
138                "Unsupported content type '{content_type}' for {url}. Supported: text/html, text/markdown, application/json, text/plain; source files via URL path extension (e.g. .rs, .ts, .mjs)"
139            ))
140        })?;
141
142    if let Some(length) = response.content_length() {
143        if length > MAX_RESPONSE_BYTES {
144            return Err(UrlFetchError::new(format!(
145                "Response too large: {length} bytes (max {MAX_RESPONSE_BYTES})"
146            )));
147        }
148    }
149
150    let body = read_response_body(response, url)?;
151    if from_source_path && body_contains_nul_in_prefix(&body) {
152        return Err(UrlFetchError::new(format!(
153            "Binary content detected for source URL {url}"
154        )));
155    }
156    let (body, content_type, extension) = if extension == ".html" {
157        (
158            convert_html_body_to_markdown(&body, url)?,
159            CONVERTED_MARKDOWN_CONTENT_TYPE.to_string(),
160            ".md",
161        )
162    } else {
163        (body, content_type, extension)
164    };
165
166    let content_file = content_path(storage_dir, &hash, extension);
167    atomic_write(&content_file, &body, &options)?;
168
169    let meta = CacheMeta {
170        url: url.to_string(),
171        content_type,
172        extension: extension.to_string(),
173        fetched_at: now_ms(),
174    };
175    let meta_bytes = serde_json::to_vec(&meta).map_err(|error| {
176        UrlFetchError::new(format!("Failed to encode URL cache metadata: {error}"))
177    })?;
178    atomic_write(&meta_file, &meta_bytes, &options)?;
179
180    Ok(content_file)
181}
182
183pub fn cleanup_url_cache(storage_dir: &Path) -> Result<usize, UrlFetchError> {
184    let dir = cache_dir(storage_dir);
185    if !dir.exists() {
186        return Ok(0);
187    }
188
189    let entries = fs::read_dir(&dir).map_err(|error| {
190        UrlFetchError::new(format!(
191            "URL cache cleanup failed reading {}: {error}",
192            dir.display()
193        ))
194    })?;
195    let mut removed = 0usize;
196    let now = now_ms();
197
198    for entry in entries.flatten() {
199        let path = entry.path();
200        let Some(name) = path.file_name().and_then(|name| name.to_str()) else {
201            continue;
202        };
203        if !name.ends_with(".meta.json") {
204            continue;
205        }
206
207        let meta = fs::read_to_string(&path)
208            .ok()
209            .and_then(|content| serde_json::from_str::<CacheMeta>(&content).ok());
210        let Some(meta) = meta else {
211            if fs::remove_file(&path).is_ok() {
212                removed += 1;
213            }
214            continue;
215        };
216
217        if now.saturating_sub(meta.fetched_at) <= CACHE_TTL_MS {
218            continue;
219        }
220
221        let hash = name.trim_end_matches(".meta.json");
222        let content = content_path(storage_dir, hash, &meta.extension);
223        let _ = fs::remove_file(content);
224        if fs::remove_file(&path).is_ok() {
225            removed += 1;
226        }
227    }
228
229    Ok(removed)
230}
231
232#[doc(hidden)]
233pub fn cache_content_path_for_url(storage_dir: &Path, url: &str, extension: &str) -> PathBuf {
234    content_path(storage_dir, &hash_url(url), extension)
235}
236
237#[doc(hidden)]
238pub fn cache_meta_path_for_url(storage_dir: &Path, url: &str) -> PathBuf {
239    meta_path(storage_dir, &hash_url(url))
240}
241
242#[doc(hidden)]
243pub fn is_private_ip_for_test(ip: IpAddr) -> bool {
244    is_private_ip(ip)
245}
246
247fn cache_dir(storage_dir: &Path) -> PathBuf {
248    storage_dir.join("url_cache")
249}
250
251fn hash_url(url: &str) -> String {
252    let digest = Sha256::digest(url.as_bytes());
253    format!("{digest:x}").chars().take(16).collect()
254}
255
256/// Rewrite well-known web-view links before fetching because those pages serve
257/// an HTML application shell instead of the file contents agents need to read.
258fn rewrite_url_for_fetch(url: &Url) -> Url {
259    let original = url.clone();
260    let mut rewritten = url.clone();
261    rewritten.set_fragment(None);
262
263    let Some(host) = url.host_str() else {
264        return rewritten;
265    };
266    if !matches!(url.scheme(), "http" | "https") {
267        return rewritten;
268    }
269
270    let segments: Vec<&str> = url
271        .path_segments()
272        .map(|segments| segments.collect())
273        .unwrap_or_default();
274    let (raw_host, raw_path) = if matches!(host, "github.com" | "www.github.com") {
275        if segments.len() >= 5 && matches!(segments[2], "blob" | "raw") {
276            (
277                "raw.githubusercontent.com",
278                format!(
279                    "/{}/{}/{}/{}",
280                    segments[0],
281                    segments[1],
282                    segments[3],
283                    segments[4..].join("/")
284                ),
285            )
286        } else {
287            return rewritten;
288        }
289    } else if host == "gitlab.com"
290        && segments.len() >= 6
291        && segments[2] == "-"
292        && segments[3] == "blob"
293    {
294        (
295            "gitlab.com",
296            format!(
297                "/{}/{}/-/raw/{}/{}",
298                segments[0],
299                segments[1],
300                segments[4],
301                segments[5..].join("/")
302            ),
303        )
304    } else {
305        return rewritten;
306    };
307
308    rewritten
309        .set_host(Some(raw_host))
310        .expect("raw content host should be a valid URL host");
311    rewritten.set_path(&raw_path);
312    log::debug!(
313        "url_fetch: rewriting web-view URL to raw content: {} -> {}",
314        original,
315        rewritten
316    );
317    rewritten
318}
319
320fn meta_path(storage_dir: &Path, hash: &str) -> PathBuf {
321    cache_dir(storage_dir).join(format!("{hash}.meta.json"))
322}
323
324fn content_path(storage_dir: &Path, hash: &str, extension: &str) -> PathBuf {
325    cache_dir(storage_dir).join(format!("{hash}{extension}"))
326}
327
328fn fresh_cached_path(
329    storage_dir: &Path,
330    hash: &str,
331    meta_file: &Path,
332    url: &Url,
333) -> Result<Option<PathBuf>, UrlFetchError> {
334    if !meta_file.exists() {
335        return Ok(None);
336    }
337
338    let meta = match fs::read_to_string(meta_file)
339        .ok()
340        .and_then(|content| serde_json::from_str::<CacheMeta>(&content).ok())
341    {
342        Some(meta) => meta,
343        None => return Ok(None),
344    };
345    let age = now_ms().saturating_sub(meta.fetched_at);
346    if meta.extension == ".html" {
347        return Ok(None);
348    }
349
350    let content_type = meta.content_type.as_str();
351    let current = resolve_fetch_extension(url, content_type);
352    let expected_ext = current.map(|(ext, _)| ext);
353    if expected_ext != Some(meta.extension.as_str()) {
354        return Ok(None);
355    }
356
357    let cached = content_path(storage_dir, hash, &meta.extension);
358    if age < CACHE_TTL_MS && cached.exists() {
359        return Ok(Some(cached));
360    }
361    Ok(None)
362}
363
364fn fetch_with_redirects(
365    start_url: &Url,
366    original_url: &str,
367    options: &UrlFetchOptions,
368) -> Result<HttpResponse, UrlFetchError> {
369    let client = build_client(options)?;
370    let mut current_url = start_url.clone();
371
372    for redirect_count in 0..=MAX_REDIRECTS {
373        validate_public_url(&current_url, options)?;
374        let response = send_with_transient_retries(&client, &current_url)?;
375
376        if !response.status().is_redirection() {
377            return Ok(response);
378        }
379        if redirect_count == MAX_REDIRECTS {
380            return Err(UrlFetchError::new(format!(
381                "Too many redirects fetching {original_url}"
382            )));
383        }
384
385        let location = response
386            .headers()
387            .get(LOCATION)
388            .and_then(|value| value.to_str().ok())
389            .ok_or_else(|| {
390                UrlFetchError::new(format!(
391                    "Redirect from {} missing Location header",
392                    current_url.as_str()
393                ))
394            })?;
395        current_url = current_url.join(location).map_err(|error| {
396            UrlFetchError::new(format!(
397                "Invalid redirect Location '{location}' from {}: {error}",
398                current_url.as_str()
399            ))
400        })?;
401    }
402
403    Err(UrlFetchError::new(format!(
404        "Too many redirects fetching {original_url}"
405    )))
406}
407
408/// Issue a single GET with the configured User-Agent + Accept headers and
409/// transparently retry only on transient connect/transport failures.
410///
411/// Returns the response (including 4xx/5xx — caller decides how to treat
412/// those). On a non-transient reqwest error (e.g. an HTTP-shaped reply that
413/// reqwest still surfaces as Err, or a TLS handshake fault that doesn't read
414/// as `is_connect`), the original error is returned immediately so the user
415/// sees the real failure without an artificial 800ms-plus delay.
416fn send_with_transient_retries(
417    client: &Client,
418    target: &Url,
419) -> Result<HttpResponse, UrlFetchError> {
420    let mut last_error: Option<reqwest::Error> = None;
421    for attempt in 0..=TRANSIENT_RETRY_ATTEMPTS {
422        let result = client
423            .get(target.clone())
424            .header(USER_AGENT, USER_AGENT_VALUE)
425            .header(ACCEPT, ACCEPT_HEADER)
426            .send();
427        match result {
428            Ok(response) => return Ok(response),
429            Err(error) => {
430                if attempt < TRANSIENT_RETRY_ATTEMPTS && is_transient_reqwest_error(&error) {
431                    thread::sleep(Duration::from_millis(TRANSIENT_RETRY_BACKOFFS_MS[attempt]));
432                    last_error = Some(error);
433                    continue;
434                }
435                return Err(UrlFetchError::new(format!(
436                    "Failed to fetch {}: {}",
437                    target.as_str(),
438                    reqwest_error_detail(&error)
439                )));
440            }
441        }
442    }
443    // Loop fell through after the last allowed retry exhausted — surface the
444    // most recent transient error rather than swallowing it.
445    Err(UrlFetchError::new(format!(
446        "Failed to fetch {} after {} retries: {}",
447        target.as_str(),
448        TRANSIENT_RETRY_ATTEMPTS,
449        last_error
450            .as_ref()
451            .map(reqwest_error_detail)
452            .unwrap_or_else(|| "unknown transient error".to_string())
453    )))
454}
455
456/// Classify a reqwest error as transient (worth a quick retry) vs terminal.
457///
458/// Transient: TCP connect failures, request-build/send TCP-level failures
459/// that don't carry status, and timeouts. These typically clear on a single
460/// retry — agents shouldn't have to ask twice for a momentary blip.
461///
462/// Terminal: anything where reqwest got far enough to decode an HTTP-shaped
463/// reply (`is_status()`, `is_body()`, `is_decode()`). Retrying those would
464/// just hammer a server that already answered.
465fn is_transient_reqwest_error(error: &reqwest::Error) -> bool {
466    error.is_connect() || error.is_timeout() || error.is_request()
467}
468
469fn build_client(options: &UrlFetchOptions) -> Result<Client, UrlFetchError> {
470    let mut builder = Client::builder()
471        .redirect(Policy::none())
472        .connect_timeout(CONNECT_TIMEOUT);
473
474    for (host, address) in &options.connect_overrides {
475        builder = builder.resolve(host, *address);
476    }
477
478    builder
479        .build()
480        .map_err(|error| UrlFetchError::new(format!("Failed to build URL fetch client: {error}")))
481}
482
483fn validate_public_url(url: &Url, options: &UrlFetchOptions) -> Result<(), UrlFetchError> {
484    if url.scheme() != "http" && url.scheme() != "https" {
485        return Err(UrlFetchError::new(format!(
486            "Only http:// and https:// URLs are supported, got: {}:",
487            url.scheme()
488        )));
489    }
490    if options.allow_private {
491        return Ok(());
492    }
493
494    let host = url
495        .host_str()
496        .ok_or_else(|| UrlFetchError::new(format!("URL missing host: {url}")))?;
497    let host_for_parse = host
498        .trim_matches(['[', ']'])
499        .split('%')
500        .next()
501        .unwrap_or(host);
502
503    if let Ok(ip) = host_for_parse.parse::<IpAddr>() {
504        reject_private_ip(host, ip)?;
505        return Ok(());
506    }
507    if host_for_parse.contains(':') {
508        return Err(UrlFetchError::new(format!(
509            "Blocked private URL host {host} ({host_for_parse})"
510        )));
511    }
512
513    let addresses = resolve_host_ips(host_for_parse, url.port_or_known_default(), options)?;
514    if addresses.is_empty() {
515        return Err(UrlFetchError::new(format!(
516            "Failed to resolve URL host {host}"
517        )));
518    }
519    for ip in addresses {
520        reject_private_ip(host, ip)?;
521    }
522
523    // We validate all resolved addresses before issuing the request. Reqwest's
524    // default resolver runs again during TCP connect, leaving the same small
525    // DNS-rebinding window the old Bun fallback accepted. A custom per-request
526    // resolver hook would close that window but adds complexity for marginal
527    // value in this opt-in agent-tooling surface.
528    Ok(())
529}
530
531fn resolve_host_ips(
532    host: &str,
533    port: Option<u16>,
534    options: &UrlFetchOptions,
535) -> Result<Vec<IpAddr>, UrlFetchError> {
536    if let Some((_, ips)) = options
537        .public_host_overrides
538        .iter()
539        .find(|(override_host, _)| override_host == host)
540    {
541        return Ok(ips.clone());
542    }
543
544    let port = port.unwrap_or(80);
545    let addrs = (host, port).to_socket_addrs().map_err(|error| {
546        UrlFetchError::new(format!("Failed to resolve URL host {host}: {error}"))
547    })?;
548    Ok(addrs.map(|addr| addr.ip()).collect())
549}
550
551fn reject_private_ip(host: &str, ip: IpAddr) -> Result<(), UrlFetchError> {
552    if is_private_ip(ip) {
553        return Err(UrlFetchError::new(format!(
554            "Blocked private URL host {host} ({ip})"
555        )));
556    }
557    Ok(())
558}
559
560/// True for any private/link-local/CGNAT/benchmark/multicast/reserved/loopback
561/// address (the full set this module refuses to fetch). Exposed so the semantic
562/// embedding SSRF guard shares one authoritative range list instead of keeping a
563/// drifting copy. Note: this INCLUDES loopback — callers that intentionally
564/// allow loopback (e.g. a local Ollama endpoint) must exclude it themselves.
565pub fn is_private_or_reserved_ip(ip: IpAddr) -> bool {
566    is_private_ip(ip)
567}
568
569fn is_private_ip(ip: IpAddr) -> bool {
570    match ip {
571        IpAddr::V4(ipv4) => is_private_ipv4(ipv4),
572        IpAddr::V6(ipv6) => is_private_ipv6(ipv6),
573    }
574}
575
576fn is_private_ipv4(ip: Ipv4Addr) -> bool {
577    let [a, b, _, _] = ip.octets();
578    a == 0
579        || a == 10
580        || a == 127
581        || (a == 172 && (16..=31).contains(&b))
582        || (a == 192 && b == 168)
583        || (a == 169 && b == 254)
584        // RFC 6598 Shared Address Space (CGNAT): 100.64.0.0/10. Not globally
585        // routable; used for provider/VPC-internal endpoints — must not be
586        // reachable via SSRF.
587        || (a == 100 && (64..=127).contains(&b))
588        // RFC 2544 benchmark subnet: 198.18.0.0/15. Reserved, non-routable.
589        || (a == 198 && (18..=19).contains(&b))
590        || a >= 224
591}
592
593fn is_private_ipv6(ip: Ipv6Addr) -> bool {
594    let segments = ip.segments();
595    let top_six_zero = segments[..6].iter().all(|segment| *segment == 0);
596    let is_mapped = segments[..5].iter().all(|segment| *segment == 0) && segments[5] == 0xffff;
597    if is_mapped || top_six_zero {
598        let embedded = Ipv4Addr::new(
599            (segments[6] >> 8) as u8,
600            (segments[6] & 0xff) as u8,
601            (segments[7] >> 8) as u8,
602            (segments[7] & 0xff) as u8,
603        );
604        return is_private_ipv4(embedded);
605    }
606
607    let first = segments[0];
608    (0xfe80..=0xfebf).contains(&first) || (0xfc00..=0xfdff).contains(&first) || first >= 0xff00
609}
610
611const BINARY_SNIFF_PREFIX: usize = 8 * 1024;
612
613fn body_contains_nul_in_prefix(body: &[u8]) -> bool {
614    let end = body.len().min(BINARY_SNIFF_PREFIX);
615    body[..end].contains(&0)
616}
617
618/// Resolve cache extension: URL path (parsed, percent-decoded) first for source
619/// languages AFT parses; otherwise content-type mapping.
620fn resolve_fetch_extension(url: &Url, content_type: &str) -> Option<(&'static str, bool)> {
621    if let Some(ext) = extension_from_url_path(url) {
622        return Some((ext, true));
623    }
624    resolve_extension_from_content_type(content_type).map(|ext| (ext, false))
625}
626
627fn extension_from_url_path(url: &Url) -> Option<&'static str> {
628    let path = url.path();
629    if path.is_empty() || path == "/" {
630        return None;
631    }
632    let segment = path.rsplit('/').next().unwrap_or(path);
633    let file_name = percent_decode_path_segment(segment);
634    let dot = file_name.rfind('.')?;
635    let ext = &file_name[dot + 1..];
636    if ext.is_empty() {
637        return None;
638    }
639    let probe = Path::new("file").with_extension(ext);
640    if detect_language(&probe).is_some() {
641        static_extension_for_lang_ext(ext)
642    } else {
643        None
644    }
645}
646
647fn percent_decode_path_segment(segment: &str) -> String {
648    let mut out = String::with_capacity(segment.len());
649    let bytes = segment.as_bytes();
650    let mut i = 0;
651    while i < bytes.len() {
652        if bytes[i] == b'%' && i + 2 < bytes.len() {
653            if let (Some(h1), Some(h2)) = (from_hex(bytes[i + 1]), from_hex(bytes[i + 2])) {
654                out.push(char::from(h1 << 4 | h2));
655                i += 3;
656                continue;
657            }
658        }
659        out.push(bytes[i] as char);
660        i += 1;
661    }
662    out
663}
664
665fn from_hex(byte: u8) -> Option<u8> {
666    match byte {
667        b'0'..=b'9' => Some(byte - b'0'),
668        b'a'..=b'f' => Some(byte - b'a' + 10),
669        b'A'..=b'F' => Some(byte - b'A' + 10),
670        _ => None,
671    }
672}
673
674/// Dotted cache extension for a path extension that `detect_language` accepts.
675fn static_extension_for_lang_ext(ext: &str) -> Option<&'static str> {
676    match ext.to_ascii_lowercase().as_str() {
677        "ts" | "mts" | "cts" => Some(".ts"),
678        "tsx" => Some(".tsx"),
679        "js" => Some(".js"),
680        "jsx" => Some(".jsx"),
681        "mjs" => Some(".mjs"),
682        "cjs" => Some(".cjs"),
683        "py" | "pyi" => Some(".py"),
684        "rs" => Some(".rs"),
685        "go" => Some(".go"),
686        "c" | "h" => Some(".c"),
687        "cc" | "cpp" | "cxx" | "hpp" | "hh" => Some(".cpp"),
688        "zig" => Some(".zig"),
689        "cs" => Some(".cs"),
690        "sh" | "bash" | "zsh" => Some(".sh"),
691        "html" | "htm" => Some(".html"),
692        "md" | "markdown" | "mdx" => Some(".md"),
693        "sol" => Some(".sol"),
694        "scss" => Some(".scss"),
695        "vue" => Some(".vue"),
696        "json" | "jsonc" => Some(".json"),
697        "scala" | "sc" => Some(".scala"),
698        "java" => Some(".java"),
699        "rb" => Some(".rb"),
700        "kt" | "kts" => Some(".kt"),
701        "swift" => Some(".swift"),
702        "inc" | "php" => Some(".php"),
703        "lua" => Some(".lua"),
704        "pl" | "pm" | "t" => Some(".pl"),
705        "yaml" | "yml" => Some(".yaml"),
706        _ => None,
707    }
708}
709
710fn resolve_extension_from_content_type(content_type: &str) -> Option<&'static str> {
711    let lower = content_type.to_ascii_lowercase();
712    let media_type = lower
713        .split(';')
714        .next()
715        .unwrap_or("")
716        .split(',')
717        .next()
718        .unwrap_or("")
719        .trim();
720
721    match media_type {
722        "text/html"
723        | "application/xhtml+xml"
724        | "application/vnd.github.html"
725        | "application/vnd.github+html" => Some(".html"),
726        "text/markdown"
727        | "text/x-markdown"
728        | "application/markdown"
729        | "application/vnd.github.raw"
730        | "application/vnd.github+raw"
731        | "application/vnd.github.v3.raw"
732        | "text/plain" => Some(".md"),
733        "application/json" | "application/ld+json" => Some(".json"),
734        other if other.ends_with("+json") => Some(".json"),
735        "text/javascript" | "application/javascript" | "application/ecmascript" => Some(".js"),
736        "text/typescript" | "application/typescript" => Some(".ts"),
737        _ => None,
738    }
739}
740
741fn convert_html_body_to_markdown(body: &[u8], url: &str) -> Result<Vec<u8>, UrlFetchError> {
742    let html = String::from_utf8_lossy(body);
743    let mut markdown = html_to_markdown_converter()
744        .convert(&html)
745        .map_err(|error| {
746            UrlFetchError::new(format!(
747                "Failed to convert HTML from {url} to Markdown: {error}"
748            ))
749        })?;
750    if !markdown.ends_with('\n') {
751        markdown.push('\n');
752    }
753    Ok(markdown.into_bytes())
754}
755
756fn html_to_markdown_converter() -> HtmlToMarkdown {
757    HtmlToMarkdown::builder()
758        .skip_tags(vec![
759            "head", "script", "style", "nav", "footer", "aside", "noscript", "button",
760        ])
761        .add_handler(
762            vec!["a"],
763            |handlers: &dyn Handlers, element: Element| -> Option<HandlerResult> {
764                if is_permalink_anchor(&element)
765                    || handlers.walk_children(element.node).content.trim() == "§"
766                {
767                    None
768                } else if element_attr_value(&element, "href")
769                    .is_some_and(|href| href.trim() == "#")
770                {
771                    Some(handlers.walk_children(element.node).content.into())
772                } else {
773                    handlers.fallback(element)
774                }
775            },
776        )
777        .add_handler(
778            vec!["header"],
779            |handlers: &dyn Handlers, element: Element| -> Option<HandlerResult> {
780                if should_skip_header(&element) {
781                    None
782                } else {
783                    handlers.fallback(element)
784                }
785            },
786        )
787        .add_handler(
788            vec!["span"],
789            |handlers: &dyn Handlers, element: Element| -> Option<HandlerResult> {
790                if element_has_class_token(&element, "token-line") {
791                    let mut content = handlers.walk_children(element.node).content;
792                    content.push('\n');
793                    Some(content.into())
794                } else {
795                    handlers.fallback(element)
796                }
797            },
798        )
799        .build()
800}
801
802fn is_permalink_anchor(element: &Element<'_>) -> bool {
803    element_has_class_token(element, "hash-link")
804        || element_has_class_token(element, "doc-anchor")
805        || element_attr_value(element, "aria-label")
806            .is_some_and(|value| value.to_ascii_lowercase().starts_with("direct link to"))
807}
808
809fn should_skip_header(element: &Element<'_>) -> bool {
810    element_has_class_token(element, "navbar")
811        || element_has_class_token(element, "site-header")
812        || element_has_class_token(element, "site-nav")
813        || element_has_class_token(element, "topbar")
814        || element_attr_value(element, "role")
815            .is_some_and(|value| value.eq_ignore_ascii_case("banner"))
816        || element_attr_value(element, "id").is_some_and(|value| {
817            let value = value.to_ascii_lowercase();
818            value.contains("navbar") || value.contains("site-header") || value.contains("site-nav")
819        })
820}
821
822fn element_has_class_token(element: &Element<'_>, token: &str) -> bool {
823    element_attr_value(element, "class")
824        .is_some_and(|value| value.split_ascii_whitespace().any(|class| class == token))
825}
826
827fn element_attr_value<'a>(element: &'a Element<'_>, name: &str) -> Option<&'a str> {
828    element
829        .attrs
830        .iter()
831        .find(|attr| attr.name.local.as_ref() == name)
832        .map(|attr| attr.value.as_ref())
833}
834
835enum BodyReadEvent {
836    Chunk(Vec<u8>),
837    Done,
838    Error(io::ErrorKind, String),
839}
840
841fn read_response_body(mut response: HttpResponse, url: &str) -> Result<Vec<u8>, UrlFetchError> {
842    let (tx, rx) = mpsc::channel();
843    thread::spawn(move || {
844        let mut buffer = [0u8; 16 * 1024];
845        loop {
846            match response.read(&mut buffer) {
847                Ok(0) => {
848                    let _ = tx.send(BodyReadEvent::Done);
849                    break;
850                }
851                Ok(n) => {
852                    if tx.send(BodyReadEvent::Chunk(buffer[..n].to_vec())).is_err() {
853                        break;
854                    }
855                }
856                Err(error) => {
857                    let kind = error.kind();
858                    let message = error.to_string();
859                    let _ = tx.send(BodyReadEvent::Error(kind, message));
860                    break;
861                }
862            }
863        }
864    });
865
866    let mut chunks = Vec::new();
867    let mut total = 0u64;
868    loop {
869        match rx.recv_timeout(BODY_CHUNK_TIMEOUT) {
870            Ok(BodyReadEvent::Chunk(chunk)) => {
871                total += chunk.len() as u64;
872                if total > MAX_RESPONSE_BYTES {
873                    return Err(UrlFetchError::new(format!(
874                        "Response exceeded {MAX_RESPONSE_BYTES} bytes, aborted"
875                    )));
876                }
877                chunks.extend_from_slice(&chunk);
878            }
879            Ok(BodyReadEvent::Done) => return Ok(chunks),
880            Ok(BodyReadEvent::Error(kind, _message)) if is_body_stall_kind(kind) => {
881                return Err(body_stall_error(url));
882            }
883            Ok(BodyReadEvent::Error(_, message)) => {
884                return Err(UrlFetchError::new(format!(
885                    "Failed to read response body for {url}: {message}"
886                )));
887            }
888            Err(mpsc::RecvTimeoutError::Timeout) => return Err(body_stall_error(url)),
889            Err(mpsc::RecvTimeoutError::Disconnected) => {
890                return Err(UrlFetchError::new(format!(
891                    "Failed to read response body for {url}: body reader stopped unexpectedly"
892                )));
893            }
894        }
895    }
896}
897
898fn body_stall_error(url: &str) -> UrlFetchError {
899    UrlFetchError::new(format!(
900        "Body read stalled (no data for {}ms) fetching {url}",
901        BODY_CHUNK_TIMEOUT.as_millis()
902    ))
903}
904
905fn is_body_stall_kind(kind: io::ErrorKind) -> bool {
906    matches!(kind, io::ErrorKind::TimedOut | io::ErrorKind::WouldBlock)
907}
908
909fn atomic_write(
910    final_path: &Path,
911    bytes: &[u8],
912    options: &UrlFetchOptions,
913) -> Result<(), UrlFetchError> {
914    let parent = final_path.parent().unwrap_or_else(|| Path::new("."));
915    fs::create_dir_all(parent).map_err(|error| {
916        UrlFetchError::new(format!(
917            "Failed to create URL cache parent {}: {error}",
918            parent.display()
919        ))
920    })?;
921
922    let file_name = final_path
923        .file_name()
924        .and_then(|name| name.to_str())
925        .ok_or_else(|| {
926            UrlFetchError::new(format!("Invalid cache path: {}", final_path.display()))
927        })?;
928    let tmp_path = final_path.with_file_name(format!(
929        "{file_name}.tmp-{}-{}",
930        std::process::id(),
931        random_nonce()
932    ));
933
934    let write_result = (|| -> io::Result<()> {
935        let mut file = fs::File::create(&tmp_path)?;
936        file.write_all(bytes)?;
937        file.flush()?;
938        Ok(())
939    })();
940    if let Err(error) = write_result {
941        let _ = fs::remove_file(&tmp_path);
942        return Err(UrlFetchError::new(format!(
943            "Failed to write URL cache temp file {}: {error}",
944            tmp_path.display()
945        )));
946    }
947
948    if let Some(observer) = &options.atomic_write_observer {
949        observer(&tmp_path, final_path);
950    }
951
952    fs::rename(&tmp_path, final_path).map_err(|error| {
953        let _ = fs::remove_file(&tmp_path);
954        UrlFetchError::new(format!(
955            "Failed to finalize URL cache file {}: {error}",
956            final_path.display()
957        ))
958    })
959}
960
961fn random_nonce() -> String {
962    let mut bytes = [0u8; 8];
963    if getrandom::fill(&mut bytes).is_err() {
964        let fallback = now_ms() ^ u64::from(std::process::id());
965        bytes = fallback.to_le_bytes();
966    }
967    let mut out = String::with_capacity(bytes.len() * 2);
968    for byte in bytes {
969        use std::fmt::Write as _;
970        let _ = write!(out, "{byte:02x}");
971    }
972    out
973}
974
975fn now_ms() -> u64 {
976    SystemTime::now()
977        .duration_since(UNIX_EPOCH)
978        .unwrap_or_default()
979        .as_millis()
980        .try_into()
981        .unwrap_or(u64::MAX)
982}
983
984fn reqwest_error_detail(error: &reqwest::Error) -> String {
985    if error.is_timeout() {
986        return format!("timeout: {error}");
987    }
988    if let Some(source) = error.source() {
989        return format!("{source}");
990    }
991    error.to_string()
992}
993
994#[cfg(test)]
995mod tests {
996    use super::*;
997    use url::Url;
998
999    #[test]
1000    fn web_view_urls_rewrite_to_raw_content() {
1001        let cases = [
1002            (
1003                "https://github.com/owner/repo/blob/main/docs/guide.md?download=1#L42",
1004                "https://raw.githubusercontent.com/owner/repo/main/docs/guide.md?download=1",
1005            ),
1006            (
1007                "https://github.com/owner/repo/raw/v1.2/src/lib.rs?raw=1#L7",
1008                "https://raw.githubusercontent.com/owner/repo/v1.2/src/lib.rs?raw=1",
1009            ),
1010            (
1011                "https://www.github.com/owner/repo/blob/main/README.md",
1012                "https://raw.githubusercontent.com/owner/repo/main/README.md",
1013            ),
1014            (
1015                "https://gitlab.com/owner/repo/-/blob/main/docs/guide.md?plain=1#L4",
1016                "https://gitlab.com/owner/repo/-/raw/main/docs/guide.md?plain=1",
1017            ),
1018        ];
1019
1020        for (input, expected) in cases {
1021            let url = Url::parse(input).unwrap();
1022            assert_eq!(rewrite_url_for_fetch(&url).as_str(), expected);
1023        }
1024    }
1025
1026    #[test]
1027    fn non_web_view_urls_are_not_rewritten() {
1028        for input in [
1029            "https://gist.github.com/owner/abcdef/blob/main/file.md",
1030            "https://github.com/owner/repo",
1031            "https://github.com/owner/repo/blob",
1032            "https://api.github.com/repos/owner/repo",
1033            "https://notgithub.com/owner/repo/blob/main/file.md",
1034            "https://gitlab.example.com/owner/repo/-/blob/main/file.md",
1035        ] {
1036            let url = Url::parse(input).unwrap();
1037            assert_eq!(rewrite_url_for_fetch(&url).as_str(), input);
1038        }
1039    }
1040
1041    #[test]
1042    fn extension_from_path_uses_parser_mapping() {
1043        let url = Url::parse("https://example.com/pkg/index.mjs").unwrap();
1044        let (ext, from_path) = resolve_fetch_extension(&url, "text/javascript").unwrap();
1045        assert_eq!(ext, ".mjs");
1046        assert!(from_path);
1047    }
1048
1049    #[test]
1050    fn text_plain_rs_url_ignores_content_type_gate() {
1051        let url = Url::parse("https://raw.githubusercontent.com/o/r/main/lib.rs").unwrap();
1052        let (ext, from_path) = resolve_fetch_extension(&url, "text/plain").unwrap();
1053        assert_eq!(ext, ".rs");
1054        assert!(from_path);
1055    }
1056
1057    #[test]
1058    fn extensionless_javascript_maps_to_js() {
1059        let url = Url::parse("https://cdn.example/bundle").unwrap();
1060        let (ext, from_path) = resolve_fetch_extension(&url, "text/javascript").unwrap();
1061        assert_eq!(ext, ".js");
1062        assert!(!from_path);
1063    }
1064
1065    #[test]
1066    fn extensionless_plain_stays_md() {
1067        let url = Url::parse("https://example.com/readme").unwrap();
1068        let (ext, _) = resolve_fetch_extension(&url, "text/plain").unwrap();
1069        assert_eq!(ext, ".md");
1070    }
1071
1072    #[test]
1073    fn query_and_fragment_do_not_break_path_extension() {
1074        let url = Url::parse("https://example.com/src/file.ts?v=2#L10").unwrap();
1075        let (ext, from_path) = resolve_fetch_extension(&url, "text/plain").unwrap();
1076        assert_eq!(ext, ".ts");
1077        assert!(from_path);
1078    }
1079
1080    #[test]
1081    fn percent_encoded_path_segment() {
1082        let url = Url::parse("https://example.com/foo%2Fbar.rs").unwrap();
1083        // path is /foo%2Fbar.rs - segment is foo%2Fbar.rs -> decode to foo/bar.rs -> ext rs
1084        let (ext, _) = resolve_fetch_extension(&url, "text/plain").unwrap();
1085        assert_eq!(ext, ".rs");
1086    }
1087
1088    #[test]
1089    fn binary_sniff_detects_nul() {
1090        let mut body = vec![b'f', b'n', 0, b' '];
1091        assert!(body_contains_nul_in_prefix(&body));
1092        body = vec![b'h'; 9000];
1093        assert!(!body_contains_nul_in_prefix(&body));
1094    }
1095
1096    #[test]
1097    fn unsupported_pdf_still_errors_via_resolve() {
1098        let url = Url::parse("https://example.com/doc.pdf").unwrap();
1099        assert!(resolve_fetch_extension(&url, "application/pdf").is_none());
1100    }
1101
1102    #[test]
1103    fn is_permalink_anchor_doc_anchor_class_drops_anchor() {
1104        let md = html_to_markdown_converter()
1105            .convert("<a class=\"doc-anchor\" href=\"#stdio-transport\">x</a>")
1106            .unwrap();
1107        assert_eq!(
1108            md.trim(),
1109            "",
1110            "doc-anchor permalink should be dropped: {md:?}"
1111        );
1112    }
1113
1114    #[test]
1115    fn html_to_markdown_rustdoc_heading_cleanup() {
1116        let html = "<main><h1>Function <span class=\"fn\">stdio</span>&nbsp;<button id=\"copy-path\">Copy item path</button></h1><h2 class=\"location\"><a href=\"#\">stdio</a></h2><h2><a class=\"doc-anchor\" href=\"#stdio-transport\">§</a>StdIO Transport</h2></main>";
1117        let md = html_to_markdown_converter().convert(html).unwrap();
1118        assert!(md.contains("# Function stdio"), "{md:?}");
1119        assert!(!md.contains("Copy item path"), "{md:?}");
1120        assert!(md.contains("## stdio"), "{md:?}");
1121        assert!(!md.contains("[stdio](#)"), "{md:?}");
1122        assert!(md.contains("## StdIO Transport"), "{md:?}");
1123        assert!(!md.contains('§'), "{md:?}");
1124        assert!(!md.contains("](#stdio-transport)"), "{md:?}");
1125    }
1126
1127    #[test]
1128    fn html_to_markdown_docusaurus_hash_link_still_dropped() {
1129        let html = "<h2>The Retain Pipeline<a class=\"hash-link\" href=\"#x\" aria-label=\"Direct link to The Retain Pipeline\">\u{200b}</a></h2>";
1130        let md = html_to_markdown_converter().convert(html).unwrap();
1131        assert!(md.contains("## The Retain Pipeline"), "{md:?}");
1132        assert!(!md.contains('\u{200b}'), "{md:?}");
1133    }
1134}