1use anyhow::{Result, anyhow};
2use async_trait::async_trait;
3use encoding_rs::{Encoding, UTF_8, UTF_16BE, UTF_16LE};
4use mermaid_model::models::retry::retry_transient_http;
5use mermaid_model::utils::{classify_host, truncate_content};
6use reqwest::Client;
7use serde::{Deserialize, Serialize};
8use std::collections::HashMap;
9use std::sync::{Arc, Mutex, OnceLock, Weak};
10use std::time::Duration;
11use tokio::sync::Semaphore;
12
13use crate::providers::ctx::WebByteBudget;
14
15const MAX_REDIRECTS: usize = 5;
16
17static EXTRACTION_SEMAPHORE: OnceLock<Arc<Semaphore>> = OnceLock::new();
18static DOWNLOAD_LIMITER: OnceLock<DownloadLimiter> = OnceLock::new();
19
20#[derive(Debug, Clone, Serialize, Deserialize)]
22pub struct SearchResult {
23 pub title: String,
24 pub url: String,
25 pub snippet: String,
26 pub full_content: String,
27}
28
29#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
31#[serde(rename_all = "snake_case")]
32pub enum FetchBackend {
33 Native,
34 OllamaCloud,
35}
36
37impl FetchBackend {
38 #[must_use]
39 pub fn as_str(self) -> &'static str {
40 match self {
41 Self::Native => "native",
42 Self::OllamaCloud => "ollama_cloud",
43 }
44 }
45}
46
47#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
49#[serde(rename_all = "snake_case")]
50pub enum ExtractionMode {
51 Readability,
52 HtmlToMarkdown,
53 PlainText,
54 Markdown,
55 Json,
56 Xml,
57 Cloud,
58}
59
60impl ExtractionMode {
61 #[must_use]
62 pub fn as_str(self) -> &'static str {
63 match self {
64 Self::Readability => "readability",
65 Self::HtmlToMarkdown => "html_to_markdown",
66 Self::PlainText => "plain_text",
67 Self::Markdown => "markdown",
68 Self::Json => "json",
69 Self::Xml => "xml",
70 Self::Cloud => "cloud",
71 }
72 }
73}
74
75#[derive(Debug, Clone, Serialize, Deserialize)]
78pub struct WebFetchResult {
79 pub requested_url: String,
80 pub final_url: Option<String>,
84 pub status: Option<u16>,
86 pub media_type: Option<String>,
87 pub charset: Option<String>,
88 pub backend: FetchBackend,
89 pub extraction: ExtractionMode,
90 pub source_bytes: usize,
91 pub output_bytes: usize,
92 pub truncated: bool,
93 pub title: String,
94 pub content: String,
95}
96
97#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
101pub enum WebFetchError {
102 #[error("invalid web URL: {0}")]
103 InvalidUrl(String),
104 #[error("web destination denied: {0}")]
105 DestinationDenied(String),
106 #[error("web redirect denied: {0}")]
107 RedirectDenied(String),
108 #[error("HTTP {status} fetching {url}")]
109 HttpStatus { status: u16, url: String },
110 #[error("invalid response media type: {0}")]
111 InvalidMedia(String),
112 #[error("unsupported response media type '{0}'")]
113 UnsupportedMedia(String),
114 #[error("response body exceeded the {limit} byte per-request limit")]
115 BodyTooLarge { limit: usize },
116 #[error("response body exceeded the {limit} byte aggregate web budget for this turn")]
117 TurnBudgetExceeded { limit: usize },
118 #[error("unsupported response charset '{0}'")]
119 UnsupportedCharset(String),
120 #[error("response body is not valid {0} text")]
121 Decode(String),
122 #[error("web fetch returned no extractable content")]
123 EmptyContent,
124 #[error("web transport failed: {0}")]
125 Transport(String),
126 #[error("web extraction failed: {0}")]
127 Extraction(String),
128 #[error("web backend failed: {0}")]
129 Backend(String),
130}
131
132impl WebFetchError {
133 #[must_use]
134 pub fn kind(&self) -> &'static str {
135 match self {
136 Self::InvalidUrl(_) => "invalid_url",
137 Self::DestinationDenied(_) => "destination_denied",
138 Self::RedirectDenied(_) => "redirect_denied",
139 Self::HttpStatus { .. } => "http_status",
140 Self::InvalidMedia(_) => "invalid_media",
141 Self::UnsupportedMedia(_) => "unsupported_media",
142 Self::BodyTooLarge { .. } => "body_too_large",
143 Self::TurnBudgetExceeded { .. } => "turn_budget_exceeded",
144 Self::UnsupportedCharset(_) => "unsupported_charset",
145 Self::Decode(_) => "decode",
146 Self::EmptyContent => "empty_content",
147 Self::Transport(_) => "transport",
148 Self::Extraction(_) => "extraction",
149 Self::Backend(_) => "backend",
150 }
151 }
152
153 #[must_use]
154 pub fn status(&self) -> Option<u16> {
155 match self {
156 Self::HttpStatus { status, .. } => Some(*status),
157 _ => None,
158 }
159 }
160}
161
162type FetchResult<T> = std::result::Result<T, WebFetchError>;
163
164#[derive(Debug, Clone, PartialEq, Eq)]
168pub struct ValidatedWebUrl(reqwest::Url);
169
170impl ValidatedWebUrl {
171 pub fn parse(raw: &str) -> FetchResult<Self> {
181 let url = reqwest::Url::parse(raw)
182 .map_err(|error| WebFetchError::InvalidUrl(error.to_string()))?;
183 Self::from_url(url)
184 }
185
186 fn from_url(url: reqwest::Url) -> FetchResult<Self> {
187 Self::from_url_with_policy(url, is_blocked_web_host)
188 }
189
190 fn from_url_with_policy(
191 mut url: reqwest::Url,
192 is_blocked: fn(&str) -> bool,
193 ) -> FetchResult<Self> {
194 match url.scheme() {
195 "http" | "https" => {},
196 scheme => {
197 return Err(WebFetchError::InvalidUrl(format!(
198 "unsupported scheme '{scheme}' (only http/https allowed)"
199 )));
200 },
201 }
202 if !url.username().is_empty() || url.password().is_some() {
203 return Err(WebFetchError::InvalidUrl(
204 "userinfo credentials are not allowed".to_string(),
205 ));
206 }
207 if url.as_str().len() > 8192 {
208 return Err(WebFetchError::InvalidUrl(
209 "URL exceeds the 8192 byte limit".to_string(),
210 ));
211 }
212 let host = url
213 .host_str()
214 .ok_or_else(|| WebFetchError::InvalidUrl("URL has no host".to_string()))?;
215 if is_blocked(host) {
216 return Err(WebFetchError::DestinationDenied(format!(
217 "non-public or metadata host '{host}'"
218 )));
219 }
220 url.set_fragment(None);
221 Ok(Self(url))
222 }
223
224 #[cfg(test)]
225 fn from_fixture_url(url: reqwest::Url) -> FetchResult<Self> {
226 fn fixture_blocked(host: &str) -> bool {
227 if classify_host(host).is_loopback() {
228 return false;
229 }
230 is_blocked_web_host(host)
231 }
232 Self::from_url_with_policy(url, fixture_blocked)
233 }
234
235 #[must_use]
236 pub fn as_url(&self) -> &reqwest::Url {
237 &self.0
238 }
239
240 #[must_use]
241 pub fn as_str(&self) -> &str {
242 self.0.as_str()
243 }
244}
245
246const METADATA_HOSTNAMES: &[&str] = &[
247 "metadata.google.internal",
248 "metadata.goog",
249 "metadata",
250 "instance-data",
251 "instance-data.ec2.internal",
252];
253
254fn is_blocked_web_host(host: &str) -> bool {
255 let normalized = host
256 .trim_start_matches('[')
257 .trim_end_matches(']')
258 .trim_end_matches('.')
259 .to_ascii_lowercase();
260 METADATA_HOSTNAMES.contains(&normalized.as_str()) || classify_host(host).is_internal()
261}
262
263#[async_trait]
266pub trait SearchProvider: Send + Sync {
267 async fn search(
268 &self,
269 query: &str,
270 count: usize,
271 budget: WebByteBudget,
272 ) -> Result<Vec<SearchResult>>;
273}
274
275#[async_trait]
279pub trait FetchProvider: Send + Sync {
280 async fn fetch(&self, url: &str, budget: WebByteBudget) -> FetchResult<WebFetchResult>;
281}
282
283#[derive(Debug, Deserialize)]
285struct OllamaSearchResponse {
286 results: Vec<OllamaSearchResult>,
287}
288
289#[derive(Debug, Deserialize)]
290struct OllamaSearchResult {
291 title: String,
292 url: String,
293 content: String,
294}
295
296#[derive(Debug, Deserialize)]
298struct OllamaFetchResponse {
299 title: Option<String>,
300 content: Option<String>,
301}
302
303#[derive(Debug, Deserialize)]
306struct SearxngResponse {
307 #[serde(default)]
308 results: Vec<SearxngResult>,
309}
310
311#[derive(Debug, Deserialize)]
312struct SearxngResult {
313 #[serde(default)]
314 title: String,
315 url: String,
316 #[serde(default)]
317 content: String,
318}
319
320const OLLAMA_API_BASE: &str = "https://ollama.com/api";
321
322const NATIVE_FETCH_UA: &str =
325 "Mozilla/5.0 (compatible; MermaidBot/1.0; +https://github.com/noahsabaj/mermaid-cli)";
326
327#[derive(Debug)]
331pub struct HttpStatusError {
332 status: u16,
333}
334
335impl HttpStatusError {
336 #[must_use]
337 pub fn status(&self) -> u16 {
338 self.status
339 }
340}
341
342impl std::fmt::Display for HttpStatusError {
343 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
344 write!(f, "HTTP {}", self.status)
345 }
346}
347
348impl std::error::Error for HttpStatusError {}
349
350#[derive(Clone)]
353pub struct OllamaWebClient {
354 client: Client,
355 api_key: String,
356}
357
358impl OllamaWebClient {
359 pub fn new(api_key: String) -> Result<Self> {
368 let client = Client::builder()
369 .redirect(reqwest::redirect::Policy::none())
370 .referer(false)
371 .build()
372 .map_err(|error| anyhow!("failed to build Ollama web client: {error}"))?;
373 Ok(Self { client, api_key })
374 }
375
376 async fn search_impl(
382 &self,
383 query: &str,
384 count: usize,
385 budget: WebByteBudget,
386 ) -> Result<Vec<SearchResult>> {
387 if count == 0 || count > 10 {
388 return Err(anyhow!(
389 "Result count must be between 1 and 10, got {count}"
390 ));
391 }
392
393 let client = self.client.clone();
394 let api_key = self.api_key.clone();
395 let query_owned = query.to_string();
396 let budget = budget.clone();
397 let endpoint = reqwest::Url::parse(&format!("{OLLAMA_API_BASE}/web_search"))
404 .map_err(|e| anyhow!("invalid Ollama web search endpoint: {e}"))?;
405 let download_permits = acquire_download_permits(&endpoint).await?;
406 let response = retry_transient_http(|| {
407 let client = client.clone();
408 let api_key = api_key.clone();
409 let query = query_owned.clone();
410 let endpoint = endpoint.clone();
411 async move {
412 client
413 .post(endpoint)
414 .header("Authorization", format!("Bearer {api_key}"))
415 .json(&serde_json::json!({
416 "query": query,
417 "max_results": count,
418 }))
419 .timeout(Duration::from_secs(30))
420 .send()
421 .await
422 .map_err(mermaid_model::models::ModelError::from)
426 }
427 })
428 .await
429 .map_err(|e| anyhow!("Failed to reach Ollama web search API: {e}"))?;
430
431 if !response.status().is_success() {
432 let status = response.status();
433 let body = read_body_capped(
434 response,
435 mermaid_model::constants::MAX_WEB_BODY_BYTES.min(64 * 1024),
436 &budget,
437 )
438 .await
439 .map(|body| String::from_utf8_lossy(&body).into_owned())
440 .unwrap_or_else(|_| "<unavailable>".to_string());
441 let body = mermaid_model::utils::redact_secrets(&body);
442 return Err(anyhow::Error::new(HttpStatusError {
443 status: status.as_u16(),
444 })
445 .context(format!(
446 "Ollama web search API returned error {status}: {body}"
447 )));
448 }
449
450 let body = read_body_capped(
451 response,
452 mermaid_model::constants::MAX_WEB_BODY_BYTES,
453 &budget,
454 )
455 .await?;
456 drop(download_permits);
457 let ollama_response: OllamaSearchResponse = serde_json::from_slice(&body)
458 .map_err(|e| anyhow!("Failed to parse Ollama search response: {e}"))?;
459
460 let search_results = map_search_results(
461 ollama_response
462 .results
463 .into_iter()
464 .map(|r| (r.title, r.url, r.content)),
465 count,
466 );
467
468 Ok(search_results)
470 }
471
472 async fn fetch_impl(&self, url: &str, budget: WebByteBudget) -> FetchResult<WebFetchResult> {
474 let requested = ValidatedWebUrl::parse(url)?;
475
476 let client = self.client.clone();
477 let api_key = self.api_key.clone();
478 let url_owned = requested.as_str().to_string();
479 let budget = budget.clone();
480 let safe_url = mermaid_model::utils::sanitize_url_for_display(&url_owned);
483 let endpoint = reqwest::Url::parse(&format!("{OLLAMA_API_BASE}/web_fetch"))
484 .map_err(|e| anyhow!("invalid Ollama web fetch endpoint: {e}"))
485 .map_err(|error| map_cloud_fetch_error(error, requested.as_str()))?;
486 let download_permits = acquire_download_permits(&endpoint)
487 .await
488 .map_err(|error| map_cloud_fetch_error(error, requested.as_str()))?;
489 let raw = retry_transient_http(|| {
490 let client = client.clone();
491 let api_key = api_key.clone();
492 let url = url_owned.clone();
493 let endpoint = endpoint.clone();
494 async move {
495 client
496 .post(endpoint)
497 .header("Authorization", format!("Bearer {api_key}"))
498 .json(&serde_json::json!({ "url": url }))
499 .timeout(Duration::from_secs(15))
500 .send()
501 .await
502 .map_err(mermaid_model::models::ModelError::from)
503 }
504 })
505 .await
506 .map_err(|e| anyhow!("Failed to fetch {safe_url}: {e}"))
507 .map_err(|error| map_cloud_fetch_error(error, requested.as_str()))?;
508
509 if !raw.status().is_success() {
510 let status = raw.status();
511 let error = anyhow::Error::new(HttpStatusError {
512 status: status.as_u16(),
513 })
514 .context(format!("Failed to fetch {safe_url}: HTTP {status}"));
515 return Err(map_cloud_fetch_error(error, requested.as_str()));
516 }
517
518 let body =
519 read_body_capped(raw, mermaid_model::constants::MAX_WEB_BODY_BYTES, &budget).await?;
520 drop(download_permits);
521 let source_bytes = body.len();
522 let parsed: OllamaFetchResponse = serde_json::from_slice(&body)
523 .map_err(|e| anyhow!("Failed to parse fetch response: {e}"))
524 .map_err(|error| map_cloud_fetch_error(error, requested.as_str()))?;
525 let response = (parsed, source_bytes);
526
527 let source_bytes = response.1;
528 let title = response.0.title.unwrap_or_default();
529 let content = response.0.content.unwrap_or_default();
530 if content.trim().is_empty() {
531 return Err(WebFetchError::EmptyContent);
532 }
533 let output_bytes = content.len();
534 Ok(WebFetchResult {
535 requested_url: requested.as_str().to_string(),
536 final_url: None,
538 status: None,
539 media_type: None,
540 charset: None,
541 backend: FetchBackend::OllamaCloud,
542 extraction: ExtractionMode::Cloud,
543 source_bytes,
544 output_bytes,
545 truncated: false,
546 title,
547 content,
548 })
549 }
550}
551
552#[async_trait]
553impl SearchProvider for OllamaWebClient {
554 async fn search(
555 &self,
556 query: &str,
557 count: usize,
558 budget: WebByteBudget,
559 ) -> Result<Vec<SearchResult>> {
560 self.search_impl(query, count, budget).await
561 }
562}
563
564#[async_trait]
565impl FetchProvider for OllamaWebClient {
566 async fn fetch(&self, url: &str, budget: WebByteBudget) -> FetchResult<WebFetchResult> {
567 self.fetch_impl(url, budget).await
568 }
569}
570
571fn map_cloud_fetch_error(error: anyhow::Error, requested_url: &str) -> WebFetchError {
572 if let Some(error) = error.downcast_ref::<WebFetchError>() {
573 return error.clone();
574 }
575 if let Some(status) = error.downcast_ref::<HttpStatusError>() {
576 return WebFetchError::HttpStatus {
577 status: status.status(),
578 url: mermaid_model::utils::sanitize_url_for_display(requested_url),
579 };
580 }
581 if let Some(error) = error.downcast_ref::<reqwest::Error>() {
582 return WebFetchError::Transport(mermaid_model::utils::redact_secrets(&error.to_string()));
583 }
584 WebFetchError::Backend(mermaid_model::utils::redact_secrets(&format!("{error:#}")))
585}
586
587#[derive(Clone)]
591pub struct SearxngClient {
592 client: Client,
593 base_url: String,
594}
595
596impl SearxngClient {
597 pub fn new(base_url: String) -> Result<Self> {
607 Self::build(base_url, false)
608 }
609
610 fn managed(base_url: String) -> Result<Self> {
611 Self::build(base_url, true)
612 }
613
614 fn build(base_url: String, managed_local: bool) -> Result<Self> {
615 let mut parsed = reqwest::Url::parse(&base_url)
616 .map_err(|error| anyhow!("invalid SearXNG URL: {error}"))?;
617 if !matches!(parsed.scheme(), "http" | "https") {
618 return Err(anyhow!("SearXNG URL must use http or https"));
619 }
620 if !parsed.username().is_empty() || parsed.password().is_some() {
621 return Err(anyhow!("SearXNG URL must not contain userinfo credentials"));
622 }
623 if parsed.query().is_some() {
624 return Err(anyhow!("SearXNG base URL must not contain a query"));
625 }
626 parsed.set_fragment(None);
627 let mut builder = Client::builder()
628 .redirect(reqwest::redirect::Policy::none())
629 .referer(false);
630 if managed_local {
631 builder = builder.no_proxy();
634 }
635 let client = builder
636 .build()
637 .map_err(|error| anyhow!("failed to build SearXNG client: {error}"))?;
638 Ok(Self {
639 client,
640 base_url: parsed.as_str().trim_end_matches('/').to_string(),
641 })
642 }
643}
644
645#[async_trait]
646impl SearchProvider for SearxngClient {
647 async fn search(
648 &self,
649 query: &str,
650 count: usize,
651 budget: WebByteBudget,
652 ) -> Result<Vec<SearchResult>> {
653 let safe_base = mermaid_model::utils::sanitize_url_for_display(&self.base_url);
654 let request_url = reqwest::Url::parse_with_params(
655 &format!("{}/search", self.base_url),
656 &[("q", query), ("format", "json")],
657 )
658 .map_err(|e| anyhow!("invalid SearXNG URL {safe_base}: {e}"))?;
659
660 let download_permits = acquire_download_permits(&request_url).await?;
661
662 let response = self
663 .client
664 .get(request_url)
665 .timeout(Duration::from_secs(30))
666 .send()
667 .await
668 .map_err(|e| {
669 anyhow::Error::new(e).context(format!(
670 "Failed to reach SearXNG at {safe_base} — is it running?"
671 ))
672 })?;
673
674 if !response.status().is_success() {
675 let status = response.status();
676 return Err(anyhow!(
677 "SearXNG at {safe_base} returned {status}. A 403 usually means the JSON format is \
678 disabled — add `json` to `search.formats` in its settings.yml."
679 ));
680 }
681
682 let body = read_body_capped(
683 response,
684 mermaid_model::constants::MAX_WEB_BODY_BYTES,
685 &budget,
686 )
687 .await?;
688 drop(download_permits);
689 let parsed: SearxngResponse = serde_json::from_slice(&body).map_err(|e| {
690 anyhow!("Failed to parse SearXNG response (is `format=json` enabled?): {e}")
691 })?;
692
693 let results = map_search_results(
694 parsed
695 .results
696 .into_iter()
697 .map(|r| (r.title, r.url, r.content)),
698 count,
699 );
700
701 Ok(results)
704 }
705}
706
707pub struct ManagedSearxngBackend;
712
713#[async_trait]
714impl SearchProvider for ManagedSearxngBackend {
715 async fn search(
716 &self,
717 query: &str,
718 count: usize,
719 budget: WebByteBudget,
720 ) -> Result<Vec<SearchResult>> {
721 let base_url = crate::searxng::manager().ensure_running().await?;
722 SearxngClient::managed(base_url)?
723 .search(query, count, budget)
724 .await
725 }
726}
727
728#[derive(Clone)]
731pub struct NativeFetchClient {
732 client: Client,
733}
734
735impl NativeFetchClient {
736 pub fn new() -> Result<Self> {
746 let client = native_client_builder()
747 .dns_resolver(std::sync::Arc::new(VettingResolver))
748 .build()
749 .map_err(|e| anyhow!("failed to build hardened native web client: {e}"))?;
750 Ok(Self { client })
751 }
752
753 async fn fetch_validated(
754 &self,
755 requested: ValidatedWebUrl,
756 budget: WebByteBudget,
757 ) -> FetchResult<WebFetchResult> {
758 self.fetch_with_validator(
759 requested,
760 ValidatedWebUrl::from_url,
761 mermaid_model::constants::MAX_WEB_BODY_BYTES,
762 budget,
763 )
764 .await
765 }
766
767 async fn fetch_with_validator(
768 &self,
769 requested: ValidatedWebUrl,
770 validate: fn(reqwest::Url) -> FetchResult<ValidatedWebUrl>,
771 max_body_bytes: usize,
772 budget: WebByteBudget,
773 ) -> FetchResult<WebFetchResult> {
774 let requested_url = requested.as_str().to_string();
775 let mut current = requested;
776
777 for redirect_count in 0..=MAX_REDIRECTS {
778 let download_permits = acquire_download_permits(current.as_url())
779 .await
780 .map_err(|error| WebFetchError::Transport(error.to_string()))?;
781 let response = self
782 .client
783 .get(current.as_url().clone())
784 .send()
785 .await
786 .map_err(|error| map_native_transport_error(error, current.as_str()))?;
787 let status = response.status();
788 if is_followable_redirect(status) {
789 if redirect_count == MAX_REDIRECTS {
790 return Err(WebFetchError::RedirectDenied(format!(
791 "exceeded the {MAX_REDIRECTS}-redirect limit"
792 )));
793 }
794 let location = response
795 .headers()
796 .get(reqwest::header::LOCATION)
797 .ok_or_else(|| {
798 WebFetchError::RedirectDenied(format!(
799 "redirect response {status} has no Location header"
800 ))
801 })?
802 .to_str()
803 .map_err(|_| {
804 WebFetchError::RedirectDenied(
805 "redirect Location header is not valid text".to_string(),
806 )
807 })?;
808 current = validated_redirect_destination(¤t, location, validate)?;
809 continue;
810 }
811
812 if !status.is_success() {
813 return Err(WebFetchError::HttpStatus {
814 status: status.as_u16(),
815 url: mermaid_model::utils::sanitize_url_for_display(current.as_str()),
816 });
817 }
818
819 let media = ResponseMedia::from_headers(response.headers())?;
821 let final_url = validate(response.url().clone())?;
822 let body = read_body_capped(response, max_body_bytes, &budget).await?;
823 drop(download_permits);
824
825 let source_bytes = body.len();
826 let extraction_permit = extraction_semaphore().acquire_owned().await.map_err(|_| {
827 WebFetchError::Extraction("web extraction limiter is closed".to_string())
828 })?;
829 let final_url_for_extract = final_url.as_str().to_string();
830 let media_for_extract = media.clone();
831 let (title, content, extraction, charset) = tokio::task::spawn_blocking(move || {
832 let _extraction_permit = extraction_permit;
836 decode_and_extract(body, &final_url_for_extract, &media_for_extract)
837 })
838 .await
839 .map_err(|error| {
840 WebFetchError::Extraction(format!("content extraction task failed: {error}"))
841 })??;
842
843 let output_bytes = content.len();
844 return Ok(WebFetchResult {
845 requested_url,
846 final_url: Some(final_url.as_str().to_string()),
847 status: Some(status.as_u16()),
848 media_type: media.media_type,
849 charset: Some(charset),
850 backend: FetchBackend::Native,
851 extraction,
852 source_bytes,
853 output_bytes,
854 truncated: false,
855 title,
856 content,
857 });
858 }
859
860 Err(WebFetchError::RedirectDenied(
861 "redirect state was exhausted".to_string(),
862 ))
863 }
864}
865
866fn native_client_builder() -> reqwest::ClientBuilder {
867 Client::builder()
868 .user_agent(NATIVE_FETCH_UA)
869 .timeout(Duration::from_secs(20))
870 .redirect(reqwest::redirect::Policy::none())
871 .referer(false)
872 .no_proxy()
876}
877
878fn is_followable_redirect(status: reqwest::StatusCode) -> bool {
879 matches!(
880 status,
881 reqwest::StatusCode::MOVED_PERMANENTLY
882 | reqwest::StatusCode::FOUND
883 | reqwest::StatusCode::SEE_OTHER
884 | reqwest::StatusCode::TEMPORARY_REDIRECT
885 | reqwest::StatusCode::PERMANENT_REDIRECT
886 )
887}
888
889fn validated_redirect_destination(
890 current: &ValidatedWebUrl,
891 location: &str,
892 validate: fn(reqwest::Url) -> FetchResult<ValidatedWebUrl>,
893) -> FetchResult<ValidatedWebUrl> {
894 let next = current
895 .as_url()
896 .join(location)
897 .map_err(|error| WebFetchError::RedirectDenied(error.to_string()))?;
898 let next = validate(next)?;
899 if current.as_url().scheme() == "https" && next.as_url().scheme() != "https" {
900 return Err(WebFetchError::RedirectDenied(format!(
901 "HTTPS-to-HTTP redirect from {} to {}",
902 mermaid_model::utils::sanitize_url_for_display(current.as_str()),
903 mermaid_model::utils::sanitize_url_for_display(next.as_str())
904 )));
905 }
906 Ok(next)
907}
908
909#[async_trait]
910impl FetchProvider for NativeFetchClient {
911 async fn fetch(&self, url: &str, budget: WebByteBudget) -> FetchResult<WebFetchResult> {
912 self.fetch_validated(ValidatedWebUrl::parse(url)?, budget)
913 .await
914 }
915}
916
917struct VettingResolver;
921
922#[derive(Debug, thiserror::Error)]
923#[error("{0}")]
924struct DestinationPolicyError(String);
925
926impl reqwest::dns::Resolve for VettingResolver {
927 fn resolve(&self, name: reqwest::dns::Name) -> reqwest::dns::Resolving {
928 Box::pin(async move {
929 let host = name.as_str().to_string();
930 let addrs: Vec<std::net::SocketAddr> =
933 tokio::net::lookup_host((host.as_str(), 0)).await?.collect();
934 vet_resolved_addresses(&host, &addrs)
935 .map_err(|error| -> Box<dyn std::error::Error + Send + Sync> { Box::new(error) })?;
936 Ok(Box::new(addrs.into_iter()) as reqwest::dns::Addrs)
937 })
938 }
939}
940
941fn vet_resolved_addresses(
942 host: &str,
943 addrs: &[std::net::SocketAddr],
944) -> std::result::Result<(), DestinationPolicyError> {
945 if addrs.is_empty() {
946 return Err(DestinationPolicyError(format!(
947 "'{host}' resolved to no addresses"
948 )));
949 }
950 if addrs
951 .iter()
952 .any(|addr| classify_host(&addr.ip().to_string()).is_internal())
953 {
954 return Err(DestinationPolicyError(format!(
955 "refusing to connect to '{host}' — its DNS answer contains a non-public address"
956 )));
957 }
958 Ok(())
959}
960
961fn map_native_transport_error(error: reqwest::Error, requested_url: &str) -> WebFetchError {
962 let mut source = std::error::Error::source(&error);
963 while let Some(current) = source {
964 if let Some(policy) = current.downcast_ref::<DestinationPolicyError>() {
965 return WebFetchError::DestinationDenied(policy.0.clone());
966 }
967 source = current.source();
968 }
969 WebFetchError::Transport(mermaid_model::utils::redact_secrets(&format!(
970 "{}: {error}",
971 mermaid_model::utils::sanitize_url_for_display(requested_url)
972 )))
973}
974
975struct DownloadPermits {
976 _global: tokio::sync::OwnedSemaphorePermit,
977 _origin: tokio::sync::OwnedSemaphorePermit,
978}
979
980async fn acquire_download_permits(url: &reqwest::Url) -> Result<DownloadPermits> {
981 download_limiter().acquire(url).await
982}
983
984struct DownloadLimiter {
985 global: Arc<Semaphore>,
986 origins: Mutex<HashMap<String, Weak<Semaphore>>>,
987 per_origin: usize,
988}
989
990impl DownloadLimiter {
991 fn new(global: usize, per_origin: usize) -> Self {
992 Self {
993 global: Arc::new(Semaphore::new(global)),
994 origins: Mutex::new(HashMap::new()),
995 per_origin,
996 }
997 }
998
999 async fn acquire(&self, url: &reqwest::Url) -> Result<DownloadPermits> {
1000 let origin = self
1001 .origin_semaphore(url)
1002 .acquire_owned()
1003 .await
1004 .map_err(|_| anyhow!("web origin limiter is closed"))?;
1005 let global = self
1006 .global
1007 .clone()
1008 .acquire_owned()
1009 .await
1010 .map_err(|_| anyhow!("web download limiter is closed"))?;
1011 Ok(DownloadPermits {
1012 _global: global,
1013 _origin: origin,
1014 })
1015 }
1016
1017 fn origin_semaphore(&self, url: &reqwest::Url) -> Arc<Semaphore> {
1018 let key = format!(
1019 "{}://{}:{}",
1020 url.scheme(),
1021 url.host_str().unwrap_or_default().to_ascii_lowercase(),
1022 url.port_or_known_default().unwrap_or_default()
1023 );
1024 let mut semaphores = self
1025 .origins
1026 .lock()
1027 .unwrap_or_else(std::sync::PoisonError::into_inner);
1028 semaphores.retain(|_, semaphore| semaphore.strong_count() > 0);
1029 if let Some(semaphore) = semaphores.get(&key).and_then(Weak::upgrade) {
1030 return semaphore;
1031 }
1032 let semaphore = Arc::new(Semaphore::new(self.per_origin));
1033 semaphores.insert(key, Arc::downgrade(&semaphore));
1034 semaphore
1035 }
1036}
1037
1038fn download_limiter() -> &'static DownloadLimiter {
1039 DOWNLOAD_LIMITER.get_or_init(|| {
1040 DownloadLimiter::new(
1041 mermaid_model::constants::MAX_WEB_DOWNLOAD_CONCURRENCY,
1042 mermaid_model::constants::MAX_WEB_PER_ORIGIN_CONCURRENCY,
1043 )
1044 })
1045}
1046
1047pub(super) fn extraction_semaphore() -> Arc<Semaphore> {
1048 EXTRACTION_SEMAPHORE
1049 .get_or_init(|| {
1050 Arc::new(Semaphore::new(
1051 mermaid_model::constants::MAX_WEB_EXTRACTION_CONCURRENCY,
1052 ))
1053 })
1054 .clone()
1055}
1056
1057#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1058enum MediaKind {
1059 Html,
1060 Xhtml,
1061 PlainText,
1062 Markdown,
1063 Json,
1064 Xml,
1065}
1066
1067#[derive(Debug, Clone)]
1068struct ResponseMedia {
1069 media_type: Option<String>,
1070 charset: Option<String>,
1071 kind: MediaKind,
1072}
1073
1074impl ResponseMedia {
1075 fn from_headers(headers: &reqwest::header::HeaderMap) -> FetchResult<Self> {
1076 let Some(raw) = headers.get(reqwest::header::CONTENT_TYPE) else {
1077 return Ok(Self {
1078 media_type: None,
1079 charset: None,
1080 kind: MediaKind::PlainText,
1081 });
1082 };
1083 let raw = raw.to_str().map_err(|_| {
1084 WebFetchError::InvalidMedia("Content-Type header is not valid text".to_string())
1085 })?;
1086 let parsed: mime::Mime = raw.parse().map_err(|error| {
1087 WebFetchError::InvalidMedia(format!("Content-Type '{raw}': {error}"))
1088 })?;
1089 let media_type = parsed.essence_str().to_ascii_lowercase();
1090 let charset = parsed
1091 .get_param(mime::CHARSET)
1092 .map(|value| value.as_str().to_string());
1093 let kind = match media_type.as_str() {
1094 "text/html" => MediaKind::Html,
1095 "application/xhtml+xml" => MediaKind::Xhtml,
1096 "text/markdown" | "text/x-markdown" | "application/markdown" => MediaKind::Markdown,
1097 "application/json" | "text/json" => MediaKind::Json,
1098 "application/xml" | "text/xml" => MediaKind::Xml,
1099 _ if media_type.ends_with("+json") => MediaKind::Json,
1100 _ if media_type.ends_with("+xml") => MediaKind::Xml,
1101 _ if media_type.starts_with("text/") => MediaKind::PlainText,
1102 _ => {
1103 return Err(WebFetchError::UnsupportedMedia(media_type));
1104 },
1105 };
1106 Ok(Self {
1107 media_type: Some(media_type),
1108 charset,
1109 kind,
1110 })
1111 }
1112}
1113
1114fn decode_and_extract(
1115 body: Vec<u8>,
1116 final_url: &str,
1117 media: &ResponseMedia,
1118) -> FetchResult<(String, String, ExtractionMode, String)> {
1119 if media.media_type.is_none() && body.contains(&0) {
1120 return Err(WebFetchError::UnsupportedMedia(
1121 "unlabeled binary data".to_string(),
1122 ));
1123 }
1124 let (decoded, charset) = decode_body(&body, media)?;
1125 let (title, content, extraction) = match media.kind {
1126 MediaKind::Html | MediaKind::Xhtml => extract_readable(&decoded, final_url)?,
1127 MediaKind::PlainText => (
1128 String::new(),
1129 decoded.trim().to_string(),
1130 ExtractionMode::PlainText,
1131 ),
1132 MediaKind::Markdown => (
1133 String::new(),
1134 decoded.trim().to_string(),
1135 ExtractionMode::Markdown,
1136 ),
1137 MediaKind::Json => {
1138 let _: serde_json::Value = serde_json::from_str(&decoded).map_err(|error| {
1139 WebFetchError::Extraction(format!("invalid JSON response body: {error}"))
1140 })?;
1141 (String::new(), decoded, ExtractionMode::Json)
1142 },
1143 MediaKind::Xml => (
1144 String::new(),
1145 decoded.trim().to_string(),
1146 ExtractionMode::Xml,
1147 ),
1148 };
1149 if content.trim().is_empty() {
1150 return Err(WebFetchError::EmptyContent);
1151 }
1152 Ok((title, content, extraction, charset))
1153}
1154
1155fn decode_body(body: &[u8], media: &ResponseMedia) -> FetchResult<(String, String)> {
1156 let signature_encoding = xml_signature_encoding(body, media.kind)?;
1160 let (encoding, bom_len) = if let Some(encoding) = signature_encoding {
1161 (encoding, 0)
1162 } else if let Some((encoding, bom_len)) = Encoding::for_bom(body) {
1163 (encoding, bom_len)
1164 } else {
1165 let declared = media.charset.as_deref().or_else(|| match media.kind {
1166 MediaKind::Html => find_ascii_assignment(body, b"charset"),
1167 MediaKind::Xhtml => find_ascii_assignment(body, b"charset")
1168 .or_else(|| find_ascii_assignment(body, b"encoding")),
1169 MediaKind::Xml => find_ascii_assignment(body, b"encoding"),
1170 _ => None,
1171 });
1172 let encoding = match declared {
1173 Some(label) => Encoding::for_label(label.as_bytes())
1174 .ok_or_else(|| WebFetchError::UnsupportedCharset(label.to_string()))?,
1175 None => UTF_8,
1176 };
1177 (encoding, 0)
1178 };
1179 let (decoded, had_errors) = encoding.decode_without_bom_handling(&body[bom_len..]);
1180 if had_errors {
1181 return Err(WebFetchError::Decode(encoding.name().to_string()));
1182 }
1183 Ok((decoded.into_owned(), encoding.name().to_ascii_lowercase()))
1184}
1185
1186fn xml_signature_encoding(body: &[u8], kind: MediaKind) -> FetchResult<Option<&'static Encoding>> {
1187 let Some(prefix) = body.get(..4) else {
1188 return Ok(None);
1189 };
1190 match prefix {
1193 [0x00, 0x00, 0xfe, 0xff] => {
1194 return Err(WebFetchError::UnsupportedCharset("utf-32be".to_string()));
1195 },
1196 [0xff, 0xfe, 0x00, 0x00] => {
1197 return Err(WebFetchError::UnsupportedCharset("utf-32le".to_string()));
1198 },
1199 _ => {},
1200 }
1201 if !matches!(kind, MediaKind::Xml | MediaKind::Xhtml) {
1202 return Ok(None);
1203 }
1204 match prefix {
1205 [0x00, 0x00, 0x00, 0x3c] => Err(WebFetchError::UnsupportedCharset("utf-32be".to_string())),
1208 [0x3c, 0x00, 0x00, 0x00] => Err(WebFetchError::UnsupportedCharset("utf-32le".to_string())),
1209 [0x00, 0x3c, 0x00, 0x3f] => Ok(Some(UTF_16BE)),
1210 [0x3c, 0x00, 0x3f, 0x00] => Ok(Some(UTF_16LE)),
1211 [0x4c, 0x6f, 0xa7, 0x94] => Err(WebFetchError::UnsupportedCharset("ebcdic".to_string())),
1213 _ => Ok(None),
1214 }
1215}
1216
1217fn find_ascii_assignment<'a>(body: &'a [u8], name: &[u8]) -> Option<&'a str> {
1218 let prefix = &body[..body.len().min(4096)];
1219 let lower: Vec<u8> = prefix.iter().map(u8::to_ascii_lowercase).collect();
1220 let mut search_from = 0;
1221 while search_from + name.len() <= lower.len() {
1222 let Some(relative_start) = lower[search_from..]
1223 .windows(name.len())
1224 .position(|window| window == name)
1225 else {
1226 break;
1227 };
1228 let start = search_from + relative_start + name.len();
1229 let mut cursor = start;
1230 while lower.get(cursor).is_some_and(u8::is_ascii_whitespace) {
1231 cursor += 1;
1232 }
1233 if lower.get(cursor) != Some(&b'=') {
1234 search_from = start;
1235 continue;
1236 }
1237 cursor += 1;
1238 while lower.get(cursor).is_some_and(u8::is_ascii_whitespace) {
1239 cursor += 1;
1240 }
1241 let quote = lower
1242 .get(cursor)
1243 .copied()
1244 .filter(|byte| matches!(byte, b'\'' | b'"'));
1245 if quote.is_some() {
1246 cursor += 1;
1247 }
1248 let end = lower[cursor..]
1249 .iter()
1250 .position(|byte| {
1251 quote.map_or_else(
1252 || byte.is_ascii_whitespace() || matches!(byte, b';' | b'>' | b'/'),
1253 |quote| *byte == quote,
1254 )
1255 })
1256 .map_or(lower.len(), |end| end + cursor);
1257 if end > cursor {
1258 return std::str::from_utf8(&prefix[cursor..end]).ok();
1259 }
1260 search_from = start;
1261 }
1262 None
1263}
1264
1265fn extract_readable(html: &str, url: &str) -> FetchResult<(String, String, ExtractionMode)> {
1270 use dom_smoothie::Readability;
1271
1272 if let Ok(mut readability) = Readability::new(html, Some(url), None)
1273 && let Ok(article) = readability.parse()
1274 {
1275 let content_html = article.content.to_string();
1276 let markdown = htmd::convert(&content_html).map_err(|error| {
1277 WebFetchError::Extraction(format!(
1278 "failed to convert readable HTML to markdown: {error}"
1279 ))
1280 })?;
1281 let markdown = markdown.trim();
1282 if !markdown.is_empty() {
1286 let title = if article.title.trim().is_empty() {
1287 fallback_title(html)
1288 } else {
1289 article.title
1290 };
1291 return Ok((title, markdown.to_string(), ExtractionMode::Readability));
1292 }
1293 }
1294
1295 let markdown = htmd::convert(html).map_err(|error| {
1296 WebFetchError::Extraction(format!("failed to convert HTML to markdown: {error}"))
1297 })?;
1298 let markdown = markdown.trim();
1299 if markdown.is_empty() {
1300 return Err(WebFetchError::EmptyContent);
1301 }
1302 Ok((
1303 fallback_title(html),
1304 markdown.to_string(),
1305 ExtractionMode::HtmlToMarkdown,
1306 ))
1307}
1308
1309fn fallback_title(html: &str) -> String {
1312 let lower = html.to_ascii_lowercase();
1313 let Some(open) = lower.find("<title") else {
1314 return String::new();
1315 };
1316 let after_tag = match html[open..].find('>') {
1317 Some(gt) => &html[open + gt + 1..],
1318 None => return String::new(),
1319 };
1320 match after_tag.to_ascii_lowercase().find("</title>") {
1321 Some(end) => after_tag[..end].trim().to_string(),
1322 None => String::new(),
1323 }
1324}
1325
1326fn map_search_results(
1329 hits: impl Iterator<Item = (String, String, String)>,
1330 count: usize,
1331) -> Vec<SearchResult> {
1332 hits.take(count)
1333 .map(|(title, url, content)| {
1334 let full_content =
1335 truncate_content(&content, mermaid_model::constants::WEB_CONTENT_MAX_CHARS);
1336 let snippet = content.chars().take(200).collect();
1337 SearchResult {
1338 title,
1339 url,
1340 snippet,
1341 full_content,
1342 }
1343 })
1344 .collect()
1345}
1346
1347#[must_use]
1352pub fn format_results(results: &[SearchResult]) -> String {
1353 let mut formatted = String::from("[SEARCH_RESULTS]\n");
1354
1355 for (i, result) in results.iter().enumerate() {
1356 let url = mermaid_model::utils::sanitize_url_for_display(&result.url);
1357 formatted.push_str(&format!(
1358 "[{}] Title: {}\nURL: {}\nContent:\n{}\n---\n",
1359 i + 1,
1360 result.title,
1361 url,
1362 result.full_content
1363 ));
1364 }
1365
1366 formatted.push_str("[/SEARCH_RESULTS]\n\n");
1367
1368 formatted.push_str("Sources:\n");
1370 for (i, result) in results.iter().enumerate() {
1371 let url = mermaid_model::utils::sanitize_url_for_display(&result.url);
1372 formatted.push_str(&format!("{}. {} - {}\n", i + 1, result.title, url));
1373 }
1374
1375 formatted
1376}
1377
1378async fn read_body_capped(
1385 response: reqwest::Response,
1386 max_bytes: usize,
1387 budget: &WebByteBudget,
1388) -> FetchResult<Vec<u8>> {
1389 use futures::StreamExt;
1390 if let Some(len) = response.content_length()
1391 && len > max_bytes as u64
1392 {
1393 return Err(WebFetchError::BodyTooLarge { limit: max_bytes });
1394 }
1395 if let Some(len) = response.content_length()
1396 && usize::try_from(len).is_ok_and(|len| len > budget.remaining())
1397 {
1398 return Err(WebFetchError::TurnBudgetExceeded {
1399 limit: mermaid_model::constants::MAX_WEB_TURN_BYTES,
1400 });
1401 }
1402 if budget.remaining() == 0 {
1403 return Err(WebFetchError::TurnBudgetExceeded {
1404 limit: mermaid_model::constants::MAX_WEB_TURN_BYTES,
1405 });
1406 }
1407 let mut stream = response.bytes_stream();
1408 let mut buf = Vec::new();
1409 while let Some(chunk) = stream.next().await {
1410 let chunk = chunk.map_err(|error| WebFetchError::Transport(error.to_string()))?;
1411 if budget.charge(chunk.len()).is_err() {
1415 return Err(WebFetchError::TurnBudgetExceeded {
1416 limit: mermaid_model::constants::MAX_WEB_TURN_BYTES,
1417 });
1418 }
1419 if chunk.len() > max_bytes.saturating_sub(buf.len()) {
1420 return Err(WebFetchError::BodyTooLarge { limit: max_bytes });
1421 }
1422 buf.extend_from_slice(&chunk);
1423 }
1424 Ok(buf)
1425}
1426
1427#[cfg(test)]
1428mod tests {
1429 use super::*;
1430 use tokio::io::{AsyncReadExt, AsyncWriteExt};
1431
1432 enum FixtureBodyMode {
1433 Fixed,
1434 Chunked,
1435 TruncatedChunked,
1436 }
1437
1438 struct FixtureResponse {
1439 status: &'static str,
1440 headers: Vec<(&'static str, String)>,
1441 body: Vec<u8>,
1442 body_mode: FixtureBodyMode,
1443 }
1444
1445 impl FixtureResponse {
1446 fn text(body: &str) -> Self {
1447 Self {
1448 status: "200 OK",
1449 headers: vec![("Content-Type", "text/plain; charset=utf-8".to_string())],
1450 body: body.as_bytes().to_vec(),
1451 body_mode: FixtureBodyMode::Fixed,
1452 }
1453 }
1454
1455 fn chunked_text(body: &str) -> Self {
1456 Self {
1457 status: "200 OK",
1458 headers: vec![("Content-Type", "text/plain; charset=utf-8".to_string())],
1459 body: body.as_bytes().to_vec(),
1460 body_mode: FixtureBodyMode::Chunked,
1461 }
1462 }
1463
1464 fn truncated_chunked() -> Self {
1465 Self {
1466 status: "200 OK",
1467 headers: vec![("Content-Type", "text/plain; charset=utf-8".to_string())],
1468 body: Vec::new(),
1469 body_mode: FixtureBodyMode::TruncatedChunked,
1470 }
1471 }
1472
1473 fn redirect(location: String) -> Self {
1474 Self {
1475 status: "302 Found",
1476 headers: vec![("Location", location)],
1477 body: Vec::new(),
1478 body_mode: FixtureBodyMode::Fixed,
1479 }
1480 }
1481
1482 fn status(status: &'static str) -> Self {
1483 Self {
1484 status,
1485 headers: Vec::new(),
1486 body: Vec::new(),
1487 body_mode: FixtureBodyMode::Fixed,
1488 }
1489 }
1490
1491 fn gzip(body: Vec<u8>) -> Self {
1492 Self {
1493 status: "200 OK",
1494 headers: vec![
1495 ("Content-Type", "text/plain; charset=utf-8".to_string()),
1496 ("Content-Encoding", "gzip".to_string()),
1497 ],
1498 body,
1499 body_mode: FixtureBodyMode::Fixed,
1500 }
1501 }
1502
1503 fn wire_bytes(self) -> Vec<u8> {
1504 let Self {
1505 status,
1506 headers,
1507 body,
1508 body_mode,
1509 } = self;
1510 let mut response = format!("HTTP/1.1 {status}\r\n");
1511 for (name, value) in headers {
1512 response.push_str(&format!("{name}: {value}\r\n"));
1513 }
1514 match body_mode {
1515 FixtureBodyMode::Fixed => {
1516 response.push_str(&format!(
1517 "Content-Length: {}\r\nConnection: close\r\n\r\n",
1518 body.len()
1519 ));
1520 let mut bytes = response.into_bytes();
1521 bytes.extend_from_slice(&body);
1522 bytes
1523 },
1524 FixtureBodyMode::Chunked => {
1525 response.push_str("Transfer-Encoding: chunked\r\nConnection: close\r\n\r\n");
1526 let mut bytes = response.into_bytes();
1527 bytes.extend_from_slice(format!("{:x}\r\n", body.len()).as_bytes());
1528 bytes.extend_from_slice(&body);
1529 bytes.extend_from_slice(b"\r\n0\r\n\r\n");
1530 bytes
1531 },
1532 FixtureBodyMode::TruncatedChunked => {
1533 response.push_str("Transfer-Encoding: chunked\r\nConnection: close\r\n\r\n");
1534 response.into_bytes()
1535 },
1536 }
1537 }
1538 }
1539
1540 struct FixtureServer {
1541 base_url: String,
1542 task: tokio::task::JoinHandle<Vec<String>>,
1543 }
1544
1545 impl FixtureServer {
1546 async fn spawn(responses: Vec<FixtureResponse>) -> Self {
1547 let listener = tokio::net::TcpListener::bind(("127.0.0.1", 0))
1548 .await
1549 .expect("bind fixture listener");
1550 let address = listener.local_addr().expect("fixture address");
1551 let task = tokio::spawn(async move {
1552 let mut requests = Vec::with_capacity(responses.len());
1553 for response in responses {
1554 let (mut stream, _) =
1555 tokio::time::timeout(Duration::from_secs(5), listener.accept())
1556 .await
1557 .expect("fixture request timed out")
1558 .expect("accept fixture request");
1559 requests.push(read_request_headers(&mut stream).await);
1560 stream
1561 .write_all(&response.wire_bytes())
1562 .await
1563 .expect("write fixture response");
1564 stream.shutdown().await.expect("close fixture response");
1565 }
1566 requests
1567 });
1568 Self {
1569 base_url: format!("http://{address}"),
1570 task,
1571 }
1572 }
1573
1574 async fn requests(self) -> Vec<String> {
1575 tokio::time::timeout(Duration::from_secs(5), self.task)
1576 .await
1577 .expect("fixture server did not finish")
1578 .expect("fixture server task panicked")
1579 }
1580 }
1581
1582 async fn read_request_headers(stream: &mut tokio::net::TcpStream) -> String {
1583 let mut request = Vec::new();
1584 loop {
1585 let mut chunk = [0_u8; 1024];
1586 let read = tokio::time::timeout(Duration::from_secs(5), stream.read(&mut chunk))
1587 .await
1588 .expect("request read timed out")
1589 .expect("read fixture request");
1590 if read == 0 {
1591 break;
1592 }
1593 request.extend_from_slice(&chunk[..read]);
1594 if request.windows(4).any(|window| window == b"\r\n\r\n") {
1595 break;
1596 }
1597 assert!(
1598 request.len() <= 64 * 1024,
1599 "fixture request headers too large"
1600 );
1601 }
1602 String::from_utf8(request).expect("HTTP request headers are UTF-8")
1603 }
1604
1605 fn fixture_client() -> NativeFetchClient {
1606 NativeFetchClient {
1607 client: native_client_builder()
1608 .build()
1609 .expect("build fixture client"),
1610 }
1611 }
1612
1613 async fn fixture_fetch(
1614 client: &NativeFetchClient,
1615 url: &str,
1616 max_body_bytes: usize,
1617 ) -> Result<WebFetchResult> {
1618 fixture_fetch_with_budget(client, url, max_body_bytes, WebByteBudget::isolated()).await
1619 }
1620
1621 async fn fixture_fetch_with_budget(
1622 client: &NativeFetchClient,
1623 url: &str,
1624 max_body_bytes: usize,
1625 budget: WebByteBudget,
1626 ) -> Result<WebFetchResult> {
1627 let requested = ValidatedWebUrl::from_fixture_url(
1628 reqwest::Url::parse(url).expect("valid fixture URL"),
1629 )?;
1630 Ok(client
1631 .fetch_with_validator(
1632 requested,
1633 ValidatedWebUrl::from_fixture_url,
1634 max_body_bytes,
1635 budget,
1636 )
1637 .await?)
1638 }
1639
1640 #[test]
1641 fn test_ollama_web_client_creation() {
1642 let client = OllamaWebClient::new("test-key".to_string()).unwrap();
1643 assert_eq!(client.api_key, "test-key");
1644 }
1645
1646 #[test]
1647 fn native_client_builds_fail_closed_configuration() {
1648 NativeFetchClient::new().expect("hardened client should build");
1649 }
1650
1651 #[test]
1652 fn configured_searxng_client_validates_its_trust_destination() {
1653 let client = SearxngClient::new("http://127.0.0.1:8080/base/#fragment".to_string())
1654 .expect("explicit self-hosted loopback is valid");
1655 assert_eq!(client.base_url, "http://127.0.0.1:8080/base");
1656 for invalid in [
1657 "file:///tmp/searxng",
1658 "https://user:password@example.test",
1659 "https://example.test/?token=secret",
1660 "not a URL",
1661 ] {
1662 assert!(
1663 SearxngClient::new(invalid.to_string()).is_err(),
1664 "accepted {invalid}"
1665 );
1666 }
1667 }
1668
1669 #[tokio::test]
1670 async fn configured_searxng_never_follows_query_bearing_redirects() {
1671 let target = tokio::net::TcpListener::bind(("127.0.0.1", 0))
1672 .await
1673 .unwrap();
1674 let target_url = format!("http://{}", target.local_addr().unwrap());
1675 let target_seen = tokio::spawn(async move {
1676 tokio::time::timeout(Duration::from_millis(250), target.accept())
1677 .await
1678 .is_ok()
1679 });
1680 let source = FixtureServer::spawn(vec![FixtureResponse::redirect(target_url)]).await;
1681 let client = SearxngClient::new(source.base_url.clone()).unwrap();
1682
1683 let error = client
1684 .search("credential-shaped-query", 1, WebByteBudget::isolated())
1685 .await
1686 .unwrap_err();
1687
1688 assert!(format!("{error:#}").contains("302"));
1689 assert!(!target_seen.await.unwrap(), "redirect target was contacted");
1690 let requests = source.requests().await;
1691 assert!(requests[0].contains("credential-shaped-query"));
1692 assert!(!requests[0].to_ascii_lowercase().contains("referer:"));
1693 }
1694
1695 #[tokio::test]
1696 async fn native_fetch_follows_manual_redirect_and_records_final_url() {
1697 let server = FixtureServer::spawn(vec![
1698 FixtureResponse::redirect("/final".to_string()),
1699 FixtureResponse::text("redirected content"),
1700 ])
1701 .await;
1702 let base_url = server.base_url.clone();
1703 let page = fixture_fetch(
1704 &fixture_client(),
1705 &format!("{base_url}/start#client-fragment"),
1706 mermaid_model::constants::MAX_WEB_BODY_BYTES,
1707 )
1708 .await
1709 .unwrap();
1710
1711 assert_eq!(page.requested_url, format!("{base_url}/start"));
1712 assert_eq!(
1713 page.final_url.as_deref(),
1714 Some(format!("{base_url}/final").as_str())
1715 );
1716 assert_eq!(page.status, Some(200));
1717 assert_eq!(page.content, "redirected content");
1718 assert_eq!(page.backend, FetchBackend::Native);
1719 let requests = server.requests().await;
1720 assert!(requests[0].starts_with("GET /start HTTP/1.1"));
1721 assert!(requests[1].starts_with("GET /final HTTP/1.1"));
1722 }
1723
1724 #[tokio::test]
1725 async fn native_fetch_rejects_private_redirect_before_connecting() {
1726 let server = FixtureServer::spawn(vec![FixtureResponse::redirect(
1727 "http://169.254.169.254/latest/meta-data".to_string(),
1728 )])
1729 .await;
1730 let start_url = format!("{}/start", server.base_url);
1731 let error = fixture_fetch(
1732 &fixture_client(),
1733 &start_url,
1734 mermaid_model::constants::MAX_WEB_BODY_BYTES,
1735 )
1736 .await
1737 .unwrap_err();
1738
1739 assert!(
1740 format!("{error:#}").contains("non-public"),
1741 "unexpected redirect error: {error:#}"
1742 );
1743 assert_eq!(server.requests().await.len(), 1);
1744 }
1745
1746 #[tokio::test]
1747 async fn cross_origin_redirect_does_not_send_referer() {
1748 let destination = FixtureServer::spawn(vec![FixtureResponse::text("destination")]).await;
1749 let destination_url = destination.base_url.clone();
1750 let source = FixtureServer::spawn(vec![FixtureResponse::redirect(format!(
1751 "{destination_url}/final"
1752 ))])
1753 .await;
1754 let source_url = format!("{}/start?token=source-secret", source.base_url);
1755
1756 let page = fixture_fetch(
1757 &fixture_client(),
1758 &source_url,
1759 mermaid_model::constants::MAX_WEB_BODY_BYTES,
1760 )
1761 .await
1762 .unwrap();
1763 assert_eq!(
1764 page.final_url.as_deref(),
1765 Some(format!("{destination_url}/final").as_str())
1766 );
1767 source.requests().await;
1768 let destination_requests = destination.requests().await;
1769 assert!(
1770 !destination_requests[0]
1771 .to_ascii_lowercase()
1772 .contains("\r\nreferer:"),
1773 "cross-origin request leaked Referer: {}",
1774 destination_requests[0]
1775 );
1776 }
1777
1778 #[tokio::test]
1779 async fn native_client_ignores_ambient_proxy_configuration() {
1780 let target = FixtureServer::spawn(vec![FixtureResponse::text("direct")]).await;
1781 let target_url = format!("{}/resource", target.base_url);
1782 let proxy_listener = tokio::net::TcpListener::bind(("127.0.0.1", 0))
1783 .await
1784 .expect("bind poison proxy");
1785 let proxy_url = format!(
1786 "http://{}",
1787 proxy_listener.local_addr().expect("proxy address")
1788 );
1789 let (proxy_seen_tx, proxy_seen_rx) = tokio::sync::oneshot::channel();
1790 let proxy_task = tokio::spawn(async move {
1791 let (mut stream, _) = proxy_listener.accept().await.expect("accept proxy request");
1792 let request = read_request_headers(&mut stream).await;
1793 let _ = proxy_seen_tx.send(request);
1794 stream
1795 .write_all(&FixtureResponse::status("502 Bad Gateway").wire_bytes())
1796 .await
1797 .expect("write proxy response");
1798 });
1799
1800 #[cfg(target_os = "windows")]
1801 let variables = vec![
1802 ("HTTP_PROXY", Some(proxy_url.as_str())),
1803 ("HTTPS_PROXY", Some(proxy_url.as_str())),
1804 ("ALL_PROXY", Some(proxy_url.as_str())),
1805 ("NO_PROXY", Some("")),
1806 ];
1807 #[cfg(not(target_os = "windows"))]
1808 let variables = vec![
1809 ("HTTP_PROXY", Some(proxy_url.as_str())),
1810 ("HTTPS_PROXY", Some(proxy_url.as_str())),
1811 ("ALL_PROXY", Some(proxy_url.as_str())),
1812 ("NO_PROXY", Some("")),
1813 ("http_proxy", Some(proxy_url.as_str())),
1814 ("https_proxy", Some(proxy_url.as_str())),
1815 ("all_proxy", Some(proxy_url.as_str())),
1816 ("no_proxy", Some("")),
1817 ];
1818
1819 let page = temp_env::async_with_vars(variables, async {
1820 let client = fixture_client();
1821 fixture_fetch(
1822 &client,
1823 &target_url,
1824 mermaid_model::constants::MAX_WEB_BODY_BYTES,
1825 )
1826 .await
1827 })
1828 .await
1829 .unwrap();
1830 assert_eq!(page.content, "direct");
1831 assert!(
1832 tokio::time::timeout(Duration::from_millis(150), proxy_seen_rx)
1833 .await
1834 .is_err(),
1835 "native fetch unexpectedly connected to the ambient proxy"
1836 );
1837 proxy_task.abort();
1838 target.requests().await;
1839 }
1840
1841 #[tokio::test]
1842 async fn native_fetch_preserves_typed_status_errors() {
1843 let server = FixtureServer::spawn(vec![FixtureResponse::status("404 Not Found")]).await;
1844 let url = format!("{}/missing", server.base_url);
1845 let error = fixture_fetch(
1846 &fixture_client(),
1847 &url,
1848 mermaid_model::constants::MAX_WEB_BODY_BYTES,
1849 )
1850 .await
1851 .unwrap_err();
1852
1853 assert!(matches!(
1854 error.downcast_ref::<WebFetchError>(),
1855 Some(WebFetchError::HttpStatus { status: 404, .. })
1856 ));
1857 assert!(format!("{error:#}").contains("HTTP 404"));
1858 server.requests().await;
1859 }
1860
1861 #[tokio::test]
1862 async fn decoded_gzip_body_cannot_bypass_streaming_limit() {
1863 use base64::Engine as _;
1864
1865 let compressed = base64::engine::general_purpose::STANDARD
1867 .decode("H4sIAAAAAAAACu3BAQ0AAADCoKzvX8IeDigAAACgcwNz3JmcABAAAA==")
1868 .unwrap();
1869 assert!(compressed.len() < 1024);
1870 let server = FixtureServer::spawn(vec![FixtureResponse::gzip(compressed)]).await;
1871 let url = format!("{}/compressed", server.base_url);
1872 let budget = WebByteBudget::isolated();
1873 let error = fixture_fetch_with_budget(&fixture_client(), &url, 1024, budget.clone())
1874 .await
1875 .unwrap_err();
1876
1877 assert!(
1878 matches!(
1879 error.downcast_ref::<WebFetchError>(),
1880 Some(WebFetchError::BodyTooLarge { limit: 1024 })
1881 ),
1882 "decoded body limit was not enforced: {error:#}"
1883 );
1884 assert!(
1885 mermaid_model::constants::MAX_WEB_TURN_BYTES - budget.remaining() > 1024,
1886 "the decoded overflow chunk was not charged to the turn budget"
1887 );
1888 server.requests().await;
1889 }
1890
1891 #[tokio::test]
1892 async fn decoded_chunks_charge_the_shared_turn_budget_before_buffering() {
1893 let server = FixtureServer::spawn(vec![FixtureResponse::text("two bytes")]).await;
1894 let url = format!("{}/budget", server.base_url);
1895 let budget = WebByteBudget::isolated();
1896 budget
1897 .charge(mermaid_model::constants::MAX_WEB_TURN_BYTES - 1)
1898 .unwrap();
1899
1900 let error = fixture_fetch_with_budget(
1901 &fixture_client(),
1902 &url,
1903 mermaid_model::constants::MAX_WEB_BODY_BYTES,
1904 budget,
1905 )
1906 .await
1907 .unwrap_err();
1908
1909 assert!(matches!(
1910 error.downcast_ref::<WebFetchError>(),
1911 Some(WebFetchError::TurnBudgetExceeded { .. })
1912 ));
1913 server.requests().await;
1914 }
1915
1916 #[tokio::test]
1917 async fn turn_budget_overflow_prevents_polling_later_chunked_bodies() {
1918 let server = FixtureServer::spawn(vec![
1919 FixtureResponse::chunked_text("two bytes"),
1920 FixtureResponse::truncated_chunked(),
1921 ])
1922 .await;
1923 let url = format!("{}/budget", server.base_url);
1924 let budget = WebByteBudget::isolated();
1925 budget
1926 .charge(mermaid_model::constants::MAX_WEB_TURN_BYTES - 1)
1927 .unwrap();
1928
1929 let first = fixture_fetch_with_budget(
1930 &fixture_client(),
1931 &url,
1932 mermaid_model::constants::MAX_WEB_BODY_BYTES,
1933 budget.clone(),
1934 )
1935 .await
1936 .unwrap_err();
1937 assert!(matches!(
1938 first.downcast_ref::<WebFetchError>(),
1939 Some(WebFetchError::TurnBudgetExceeded { .. })
1940 ));
1941 assert_eq!(budget.remaining(), 0);
1942
1943 let second = fixture_fetch_with_budget(
1947 &fixture_client(),
1948 &url,
1949 mermaid_model::constants::MAX_WEB_BODY_BYTES,
1950 budget.clone(),
1951 )
1952 .await
1953 .unwrap_err();
1954 assert!(
1955 matches!(
1956 second.downcast_ref::<WebFetchError>(),
1957 Some(WebFetchError::TurnBudgetExceeded { .. })
1958 ),
1959 "later body was polled after budget exhaustion: {second:#}"
1960 );
1961 assert_eq!(budget.remaining(), 0);
1962 server.requests().await;
1963 }
1964
1965 #[test]
1966 fn validated_url_normalizes_and_rejects_unsafe_destinations() {
1967 let url = ValidatedWebUrl::parse("https://example.com/path?q=1#private-state").unwrap();
1968 assert_eq!(url.as_str(), "https://example.com/path?q=1");
1969
1970 for unsafe_url in [
1971 "file:///etc/passwd",
1972 "https://user:password@example.com/",
1973 "http://127.0.0.1/",
1974 "http://2130706433/",
1975 "http://0x7f000001/",
1976 "http://169.254.169.254/latest/meta-data/",
1977 "http://198.18.0.1/",
1978 "http://[::ffff:127.0.0.1]/",
1979 "http://metadata.google.internal/",
1980 ] {
1981 assert!(
1982 ValidatedWebUrl::parse(unsafe_url).is_err(),
1983 "{unsafe_url} must be rejected"
1984 );
1985 }
1986 assert!(matches!(
1987 ValidatedWebUrl::parse("http://127.0.0.1/"),
1988 Err(WebFetchError::DestinationDenied(_))
1989 ));
1990 }
1991
1992 #[test]
1993 fn redirect_destination_is_revalidated_and_cannot_downgrade() {
1994 let https = ValidatedWebUrl::parse("https://example.com/a/start").unwrap();
1995 let relative =
1996 validated_redirect_destination(&https, "../final#section", ValidatedWebUrl::from_url)
1997 .unwrap();
1998 assert_eq!(relative.as_str(), "https://example.com/final");
1999
2000 assert!(matches!(
2001 validated_redirect_destination(
2002 &https,
2003 "http://example.com/plaintext",
2004 ValidatedWebUrl::from_url,
2005 ),
2006 Err(WebFetchError::RedirectDenied(_))
2007 ));
2008 assert!(
2009 validated_redirect_destination(
2010 &https,
2011 "http://127.0.0.1/admin",
2012 ValidatedWebUrl::from_url,
2013 )
2014 .is_err()
2015 );
2016 assert!(
2017 validated_redirect_destination(
2018 &https,
2019 "http://169.254.169.254/latest",
2020 ValidatedWebUrl::from_url,
2021 )
2022 .is_err()
2023 );
2024
2025 let http = ValidatedWebUrl::parse("http://example.com/start").unwrap();
2026 assert!(
2027 validated_redirect_destination(
2028 &http,
2029 "https://example.com/final",
2030 ValidatedWebUrl::from_url,
2031 )
2032 .is_ok()
2033 );
2034 assert!(is_followable_redirect(reqwest::StatusCode::FOUND));
2035 assert!(!is_followable_redirect(reqwest::StatusCode::NOT_MODIFIED));
2036 }
2037
2038 #[test]
2039 fn test_format_results() {
2040 let results = vec![SearchResult {
2041 title: "Test Article".to_string(),
2042 url: "https://example.com".to_string(),
2043 snippet: "This is a test".to_string(),
2044 full_content: "Full content here".to_string(),
2045 }];
2046
2047 let formatted = format_results(&results);
2048 assert!(formatted.contains("[SEARCH_RESULTS]"));
2049 assert!(formatted.contains("Test Article"));
2050 assert!(formatted.contains("https://example.com"));
2051 assert!(formatted.contains("[/SEARCH_RESULTS]"));
2052 }
2053
2054 #[test]
2055 fn format_results_sanitizes_urls_without_mutating_results() {
2056 let raw = "https://user:hunter2@example.com/path?token=opaque-secret&ok=yes#fragment";
2057 let results = vec![SearchResult {
2058 title: "Sensitive link".to_string(),
2059 url: raw.to_string(),
2060 snippet: String::new(),
2061 full_content: "content".to_string(),
2062 }];
2063
2064 let formatted = format_results(&results);
2065 assert_eq!(results[0].url, raw, "transport value must remain unchanged");
2066 for secret in ["user", "hunter2", "opaque-secret", "fragment"] {
2067 assert!(!formatted.contains(secret), "leaked {secret}: {formatted}");
2068 }
2069 assert!(formatted.contains("ok=yes"));
2070 assert!(formatted.contains("token=%5BREDACTED%5D"));
2071 }
2072
2073 #[test]
2074 fn map_search_results_truncates_and_caps_count() {
2075 let hits = (0..5).map(|i| {
2076 (
2077 format!("t{i}"),
2078 format!("https://e{i}.com"),
2079 "x".repeat(mermaid_model::constants::WEB_CONTENT_MAX_CHARS * 2),
2080 )
2081 });
2082 let out = map_search_results(hits, 3);
2083 assert_eq!(out.len(), 3, "count cap applied");
2084 assert!(
2085 out[0].full_content.len() <= mermaid_model::constants::WEB_CONTENT_MAX_CHARS + 64,
2086 "content truncated"
2087 );
2088 assert!(out[0].snippet.chars().count() <= 200);
2089 }
2090
2091 #[test]
2092 fn searxng_response_parses_results() {
2093 let json = serde_json::json!({
2094 "results": [
2095 {"title": "A", "url": "https://a.com", "content": "alpha"},
2096 {"url": "https://b.com"},
2097 ]
2098 })
2099 .to_string();
2100 let parsed: SearxngResponse = serde_json::from_str(&json).unwrap();
2101 assert_eq!(parsed.results.len(), 2);
2102 assert_eq!(parsed.results[0].url, "https://a.com");
2103 assert_eq!(parsed.results[1].title, "");
2105 assert_eq!(parsed.results[1].content, "");
2106 }
2107
2108 #[test]
2109 fn extract_readable_produces_markdown() {
2110 let html = r#"<html><head><title>My Page</title></head>
2111 <body><article><h1>Heading</h1><p>Hello <a href="https://x.com">link</a>.</p>
2112 <p>More text to satisfy the readability length heuristic so this block is
2113 treated as the main article content rather than boilerplate chrome.</p>
2114 </article></body></html>"#;
2115 let (title, md, mode) = extract_readable(html, "https://example.com/page").unwrap();
2116 assert!(!title.is_empty(), "title extracted, got {title:?}");
2117 assert!(matches!(
2118 mode,
2119 ExtractionMode::Readability | ExtractionMode::HtmlToMarkdown
2120 ));
2121 assert!(
2122 md.contains("Hello"),
2123 "content converted to markdown: {md:?}"
2124 );
2125 assert!(md.contains("](https://x.com)"), "links preserved: {md:?}");
2126 }
2127
2128 #[test]
2129 fn extract_readable_fallback_title_on_unparseable() {
2130 let html = "<title>Bare</title><p>just a snippet</p>";
2132 let (title, md, _) = extract_readable(html, "https://example.com").unwrap();
2133 assert_eq!(title, "Bare");
2134 assert!(md.contains("just a snippet"));
2135 }
2136
2137 #[test]
2138 fn html_extraction_uses_final_url_as_relative_link_base() {
2139 let html = r#"<html><body><article><h1>Page</h1>
2140 <p>This article has enough useful prose for extraction and a
2141 <a href="../source">relative source link</a> that must be resolved
2142 against the final response URL after redirects.</p></article></body></html>"#;
2143 let (_, markdown, _) =
2144 extract_readable(html, "https://example.com/redirected/page").unwrap();
2145 assert!(
2146 markdown.contains("https://example.com/source"),
2147 "relative link was not resolved against final URL: {markdown}"
2148 );
2149 }
2150
2151 fn response_media(raw: Option<&str>) -> FetchResult<ResponseMedia> {
2152 let mut headers = reqwest::header::HeaderMap::new();
2153 if let Some(raw) = raw {
2154 headers.insert(
2155 reqwest::header::CONTENT_TYPE,
2156 reqwest::header::HeaderValue::from_str(raw).unwrap(),
2157 );
2158 }
2159 ResponseMedia::from_headers(&headers)
2160 }
2161
2162 #[test]
2163 fn content_type_dispatch_is_exact_and_rejects_binary() {
2164 assert_eq!(
2165 response_media(Some("text/html; charset=utf-8"))
2166 .unwrap()
2167 .kind,
2168 MediaKind::Html
2169 );
2170 assert_eq!(
2171 response_media(Some("application/problem+json"))
2172 .unwrap()
2173 .kind,
2174 MediaKind::Json
2175 );
2176 assert_eq!(
2177 response_media(Some("application/atom+xml")).unwrap().kind,
2178 MediaKind::Xml
2179 );
2180 assert_eq!(
2181 response_media(Some("text/markdown")).unwrap().kind,
2182 MediaKind::Markdown
2183 );
2184 assert!(response_media(Some("application/octet-stream")).is_err());
2185 assert!(response_media(Some("text/html-ish")).is_ok());
2186 assert_eq!(
2187 response_media(Some("text/html-ish")).unwrap().kind,
2188 MediaKind::PlainText
2189 );
2190 }
2191
2192 #[test]
2193 fn declared_charset_is_decoded_without_lossy_replacement() {
2194 let media = response_media(Some("text/plain; charset=windows-1252")).unwrap();
2195 let (decoded, charset) = decode_body(b"caf\xe9", &media).unwrap();
2196 assert_eq!(decoded, "caf\u{e9}");
2197 assert_eq!(charset, "windows-1252");
2198
2199 let utf8 = response_media(Some("text/plain")).unwrap();
2200 assert!(decode_body(&[0xff], &utf8).is_err());
2201 }
2202
2203 #[test]
2204 fn html_meta_charset_is_honored_when_header_omits_it() {
2205 let media = response_media(Some("text/html")).unwrap();
2206 let body = b"<meta charset=\"windows-1252\"><p>caf\xe9</p>";
2207 let (decoded, charset) = decode_body(body, &media).unwrap();
2208 assert!(decoded.contains("caf\u{e9}"));
2209 assert_eq!(charset, "windows-1252");
2210 }
2211
2212 #[test]
2213 fn xhtml_xml_declaration_charset_is_honored() {
2214 let media = response_media(Some("application/xhtml+xml")).unwrap();
2215 assert_eq!(media.kind, MediaKind::Xhtml);
2216 let body = b"<?xml version=\"1.0\" encoding=\"windows-1252\"?><html><p>caf\xe9</p></html>";
2217 let (decoded, charset) = decode_body(body, &media).unwrap();
2218 assert!(decoded.contains("caf\u{e9}"));
2219 assert_eq!(charset, "windows-1252");
2220 }
2221
2222 #[test]
2223 fn bomless_xml_and_xhtml_utf16_signatures_are_decoded() {
2224 let source = "<?xml version=\"1.0\" encoding=\"UTF-16\"?><root>café</root>";
2225 for media_type in ["application/xml", "application/xhtml+xml"] {
2226 let media = response_media(Some(media_type)).unwrap();
2227 for (big_endian, expected_charset) in [(false, "utf-16le"), (true, "utf-16be")] {
2228 let body = source
2229 .encode_utf16()
2230 .flat_map(|unit| {
2231 if big_endian {
2232 unit.to_be_bytes()
2233 } else {
2234 unit.to_le_bytes()
2235 }
2236 })
2237 .collect::<Vec<_>>();
2238 let (decoded, charset) = decode_body(&body, &media).unwrap();
2239 assert_eq!(decoded, source);
2240 assert_eq!(charset, expected_charset);
2241 }
2242 }
2243 }
2244
2245 #[test]
2246 fn xml_utf32_signatures_return_a_typed_unsupported_charset() {
2247 let media = response_media(Some("application/xml")).unwrap();
2248 for (body, expected) in [
2249 (vec![0x00, 0x00, 0x00, 0x3c], "utf-32be"),
2250 (vec![0x3c, 0x00, 0x00, 0x00], "utf-32le"),
2251 (vec![0x00, 0x00, 0xfe, 0xff], "utf-32be"),
2252 (vec![0xff, 0xfe, 0x00, 0x00], "utf-32le"),
2253 ] {
2254 assert_eq!(
2255 decode_body(&body, &media),
2256 Err(WebFetchError::UnsupportedCharset(expected.to_string()))
2257 );
2258 }
2259 }
2260
2261 #[test]
2262 fn json_and_xml_are_not_interpreted_as_html() {
2263 let json_media = response_media(Some("application/json")).unwrap();
2264 let json_source = " \n{\"markup\":\"<h1>literal</h1>\",\"items\":[1,2]}\t ";
2265 let (_, json, mode, _) = decode_and_extract(
2266 json_source.as_bytes().to_vec(),
2267 "https://example.com/data",
2268 &json_media,
2269 )
2270 .unwrap();
2271 assert_eq!(mode, ExtractionMode::Json);
2272 assert_eq!(json, json_source);
2273
2274 let invalid = decode_and_extract(
2275 br#"{"markup": "unterminated}"#.to_vec(),
2276 "https://example.com/data",
2277 &json_media,
2278 )
2279 .unwrap_err();
2280 assert!(matches!(invalid, WebFetchError::Extraction(_)));
2281
2282 let xml_media = response_media(Some("application/xml")).unwrap();
2283 let xml_source = "<root><script>literal text</script></root>";
2284 let (_, xml, mode, _) = decode_and_extract(
2285 xml_source.as_bytes().to_vec(),
2286 "https://example.com/data",
2287 &xml_media,
2288 )
2289 .unwrap();
2290 assert_eq!(mode, ExtractionMode::Xml);
2291 assert_eq!(xml, xml_source);
2292 }
2293
2294 #[test]
2295 fn empty_extractions_are_errors() {
2296 let plain = response_media(Some("text/plain")).unwrap();
2297 assert!(decode_and_extract(b" \r\n".to_vec(), "https://example.com", &plain).is_err());
2298 let html = response_media(Some("text/html")).unwrap();
2299 assert!(
2300 decode_and_extract(
2301 b"<html><head></head><body></body></html>".to_vec(),
2302 "https://example.com",
2303 &html,
2304 )
2305 .is_err()
2306 );
2307 let unlabeled = response_media(None).unwrap();
2308 assert!(
2309 decode_and_extract(
2310 b"GIF89a\0binary".to_vec(),
2311 "https://example.com/image",
2312 &unlabeled,
2313 )
2314 .is_err()
2315 );
2316 }
2317
2318 #[tokio::test]
2319 async fn vetting_resolver_rejects_a_name_resolving_to_loopback() {
2320 use reqwest::dns::Resolve;
2321 use std::str::FromStr;
2322 let name = reqwest::dns::Name::from_str("localhost").expect("valid host name");
2327 assert!(
2328 VettingResolver.resolve(name).await.is_err(),
2329 "a name resolving to a loopback address must be rejected"
2330 );
2331 }
2332
2333 #[test]
2334 fn resolved_address_vetting_rejects_empty_and_mixed_answers() {
2335 use std::net::{IpAddr, Ipv4Addr, SocketAddr};
2336
2337 let public = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(93, 184, 216, 34)), 443);
2338 let second_public = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(1, 1, 1, 1)), 443);
2339 let loopback = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443);
2340
2341 assert!(vet_resolved_addresses("example.com", &[]).is_err());
2342 assert!(vet_resolved_addresses("example.com", &[public, loopback]).is_err());
2343 assert!(vet_resolved_addresses("example.com", &[loopback, public]).is_err());
2344 assert!(vet_resolved_addresses("example.com", &[public, second_public]).is_ok());
2345 }
2346
2347 #[tokio::test]
2348 async fn download_limiter_enforces_global_origin_and_cancellation_release() {
2349 let limiter = DownloadLimiter::new(
2350 mermaid_model::constants::MAX_WEB_DOWNLOAD_CONCURRENCY,
2351 mermaid_model::constants::MAX_WEB_PER_ORIGIN_CONCURRENCY,
2352 );
2353
2354 let mut global_permits = Vec::new();
2355 for index in 0..mermaid_model::constants::MAX_WEB_DOWNLOAD_CONCURRENCY {
2356 let url = reqwest::Url::parse(&format!("https://origin-{index}.example.test/"))
2357 .expect("test URL");
2358 global_permits.push(limiter.acquire(&url).await.expect("global permit"));
2359 }
2360 let ninth_url = reqwest::Url::parse("https://ninth.example.test/").unwrap();
2361 let mut ninth = Box::pin(limiter.acquire(&ninth_url));
2362 assert!(
2363 tokio::time::timeout(Duration::from_millis(50), &mut ninth)
2364 .await
2365 .is_err(),
2366 "a ninth global download was admitted"
2367 );
2368 drop(global_permits.pop());
2369 let ninth_permit = tokio::time::timeout(Duration::from_secs(1), &mut ninth)
2370 .await
2371 .expect("global waiter was not released")
2372 .expect("global permit");
2373 drop(ninth_permit);
2374 drop(global_permits);
2375
2376 let same_origin = reqwest::Url::parse("https://same.example.test/page").unwrap();
2377 let first = limiter.acquire(&same_origin).await.unwrap();
2378 let second = limiter.acquire(&same_origin).await.unwrap();
2379 let mut third = Box::pin(limiter.acquire(&same_origin));
2380 assert!(
2381 tokio::time::timeout(Duration::from_millis(50), &mut third)
2382 .await
2383 .is_err(),
2384 "a third same-origin download was admitted"
2385 );
2386 drop(first);
2387 let third = tokio::time::timeout(Duration::from_secs(1), &mut third)
2388 .await
2389 .expect("origin waiter was not released")
2390 .expect("origin permit");
2391 drop(third);
2392 drop(second);
2393
2394 let cancellation_limiter = Arc::new(DownloadLimiter::new(1, 1));
2395 let held_url = reqwest::Url::parse("https://cancel.example.test/").unwrap();
2396 let holder_limiter = cancellation_limiter.clone();
2397 let holder_url = held_url.clone();
2398 let (ready_tx, ready_rx) = tokio::sync::oneshot::channel();
2399 let holder = tokio::spawn(async move {
2400 let _permit = holder_limiter.acquire(&holder_url).await.unwrap();
2401 let _ = ready_tx.send(());
2402 std::future::pending::<()>().await;
2403 });
2404 ready_rx.await.expect("holder acquired its permit");
2405 let mut waiter = Box::pin(cancellation_limiter.acquire(&held_url));
2406 assert!(
2407 tokio::time::timeout(Duration::from_millis(50), &mut waiter)
2408 .await
2409 .is_err()
2410 );
2411 holder.abort();
2412 let _ = holder.await;
2413 let released = tokio::time::timeout(Duration::from_secs(1), &mut waiter)
2414 .await
2415 .expect("cancelled holder leaked its permits")
2416 .expect("released permit");
2417 drop(released);
2418 }
2419}