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