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
30const 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 #[doc(hidden)]
52 pub public_host_overrides: Vec<(String, Vec<IpAddr>)>,
53 #[doc(hidden)]
56 pub connect_overrides: Vec<(String, SocketAddr)>,
57 #[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
256fn 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(¤t_url, options)?;
374 let response = send_with_transient_retries(&client, ¤t_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
408fn 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 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
456fn 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 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
560pub 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 || (a == 100 && (64..=127).contains(&b))
588 || (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
618fn 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
674fn 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 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> <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}