返回 CodeWhale
fetch.rs
根目录 / crates / tui / src / tools / web / fetch.rs
1 //! Unified guarded fetch pipeline for `fetch_url` and `web.run`.
2
3 use std::collections::BTreeMap;
4 use std::sync::Arc;
5 use std::time::{Duration, Instant};
6 #[cfg(not(test))]
7 use std::time::{SystemTime, UNIX_EPOCH};
8
9 use futures_util::StreamExt;
10
11 use super::cache::{self, CachedFetch};
12 use super::guard::{
13 DnsPin, guarded_reqwest_client_builder, validate_fetch_target, validate_network_policy,
14 };
15 use crate::tools::spec::{ToolContext, ToolError};
16
17 pub(crate) const DEFAULT_TIMEOUT: Duration = Duration::from_secs(15);
18 pub(crate) const HARD_MAX_TIMEOUT: Duration = Duration::from_secs(60);
19 pub(crate) const DEFAULT_MAX_BYTES: usize = 1_000_000;
20 pub(crate) const HARD_MAX_BYTES: usize = 10 * 1024 * 1024;
21 const MAX_REDIRECTS: usize = 5;
22 const USER_AGENT: &str = concat!(
23 "Mozilla/5.0 (compatible; codewhale/",
24 env!("CARGO_PKG_VERSION"),
25 "; +https://github.com/Hmbown/CodeWhale)"
26 );
27
28 #[derive(Debug, Clone)]
29 pub(crate) struct FetchOptions {
30 pub(crate) timeout: Duration,
31 pub(crate) max_bytes: usize,
32 pub(crate) accept: &'static str,
33 }
34
35 impl FetchOptions {
36 pub(crate) fn new(timeout: Duration, max_bytes: usize, accept: &'static str) -> Self {
37 Self {
38 timeout: timeout.min(HARD_MAX_TIMEOUT),
39 max_bytes: max_bytes.clamp(1, HARD_MAX_BYTES),
40 accept,
41 }
42 }
43 }
44
45 #[derive(Debug, Clone)]
46 pub(crate) struct FetchedPayload {
47 pub(crate) url: String,
48 pub(crate) status: u16,
49 pub(crate) headers: BTreeMap<String, String>,
50 pub(crate) content_type: String,
51 pub(crate) bytes: Arc<Vec<u8>>,
52 pub(crate) truncated: bool,
53 pub(crate) cache_hit: bool,
54 pub(crate) retries: usize,
55 pub(crate) redirects: usize,
56 }
57
58 pub(crate) async fn fetch(
59 url: &str,
60 options: &FetchOptions,
61 context: &ToolContext,
62 tool_label: &str,
63 ) -> Result<FetchedPayload, ToolError> {
64 fetch_inner(url, options, context, tool_label, None).await
65 }
66
67 #[cfg(test)]
68 pub(crate) async fn fetch_with_initial_pin(
69 url: &str,
70 options: &FetchOptions,
71 context: &ToolContext,
72 tool_label: &str,
73 initial_pin: DnsPin,
74 ) -> Result<FetchedPayload, ToolError> {
75 fetch_inner(url, options, context, tool_label, Some(initial_pin)).await
76 }
77
78 async fn fetch_inner(
79 url: &str,
80 options: &FetchOptions,
81 context: &ToolContext,
82 tool_label: &str,
83 test_initial_pin: Option<DnsPin>,
84 ) -> Result<FetchedPayload, ToolError> {
85 let initial_url = reqwest::Url::parse(url)
86 .map_err(|err| ToolError::invalid_input(format!("invalid URL: {err}")))?;
87 if !matches!(initial_url.scheme(), "http" | "https") {
88 return Err(ToolError::invalid_input(
89 "only http:// and https:// URLs are supported",
90 ));
91 }
92
93 // Validation precedes cache lookup so a policy tightened during the
94 // session cannot be bypassed by a previously cached response.
95 let validated_initial_pin = match test_initial_pin {
96 Some(pin) => pin,
97 None => validate_fetch_target(&initial_url, context, tool_label).await?,
98 };
99
100 if let Some(cached) = cache::get(
101 &context.state_namespace,
102 &initial_url,
103 options.accept,
104 options.max_bytes,
105 ) {
106 let cached_url = reqwest::Url::parse(&cached.url).map_err(|err| {
107 ToolError::execution_failed(format!("cached response URL was invalid: {err}"))
108 })?;
109 let cached_host = cached_url.host_str().ok_or_else(|| {
110 ToolError::execution_failed("cached response URL did not include a host")
111 })?;
112 // No network request occurs on a cache hit, so DNS/SSRF validation is
113 // unnecessary. The final redirect destination still needs a policy
114 // check in case the session policy was tightened after insertion.
115 validate_network_policy(cached_host, context, tool_label)?;
116 return Ok(from_cached(cached, true, 0));
117 }
118
119 let deadline = Instant::now() + options.timeout;
120 let mut last_transient = None;
121 for attempt in 0..=1 {
122 let remaining = deadline.saturating_duration_since(Instant::now());
123 if remaining.is_zero() {
124 break;
125 }
126 match fetch_attempt(
127 initial_url.clone(),
128 options,
129 context,
130 tool_label,
131 remaining,
132 validated_initial_pin.clone(),
133 )
134 .await
135 {
136 Ok(payload) if is_transient_status(payload.status) && attempt == 0 => {
137 last_transient = Some(format!("HTTP {}", payload.status));
138 }
139 Ok(payload) => {
140 let fetched = from_cached(payload.clone(), false, attempt);
141 if (200..300).contains(&payload.status) {
142 cache::insert(
143 &context.state_namespace,
144 &initial_url,
145 options.accept,
146 payload,
147 );
148 }
149 return Ok(fetched);
150 }
151 Err(AttemptError::Fatal(error)) => return Err(error),
152 Err(AttemptError::Transient(message)) if attempt == 0 => {
153 last_transient = Some(message);
154 }
155 Err(AttemptError::Transient(message)) => {
156 return Err(ToolError::execution_failed(format!(
157 "request failed after one retry: {message}"
158 )));
159 }
160 }
161
162 let delay = retry_delay();
163 if deadline.saturating_duration_since(Instant::now()) <= delay {
164 break;
165 }
166 tokio::time::sleep(delay).await;
167 }
168
169 Err(ToolError::execution_failed(format!(
170 "request timed out before retry completed{}",
171 last_transient
172 .map(|message| format!(" (last failure: {message})"))
173 .unwrap_or_default()
174 )))
175 }
176
177 #[derive(Debug)]
178 enum AttemptError {
179 Fatal(ToolError),
180 Transient(String),
181 }
182
183 async fn fetch_attempt(
184 initial_url: reqwest::Url,
185 options: &FetchOptions,
186 context: &ToolContext,
187 tool_label: &str,
188 timeout: Duration,
189 initial_pin: DnsPin,
190 ) -> Result<CachedFetch, AttemptError> {
191 let mut current_url = initial_url;
192 let mut redirects = 0usize;
193 let mut initial_pin = initial_pin;
194 let deadline = Instant::now() + timeout;
195
196 let response = loop {
197 let dns_pin = if redirects == 0 {
198 match initial_pin.take() {
199 Some(pin) => Some(pin),
200 None => validate_fetch_target(&current_url, context, tool_label)
201 .await
202 .map_err(AttemptError::Fatal)?,
203 }
204 } else {
205 validate_fetch_target(&current_url, context, tool_label)
206 .await
207 .map_err(AttemptError::Fatal)?
208 };
209
210 let remaining = deadline.saturating_duration_since(Instant::now());
211 if remaining.is_zero() {
212 return Err(AttemptError::Transient(
213 "request timed out while following redirects".to_string(),
214 ));
215 }
216 let mut builder = guarded_reqwest_client_builder()
217 .timeout(remaining)
218 .user_agent(USER_AGENT)
219 .redirect(reqwest::redirect::Policy::none());
220 if let Some((hostname, validated_ip)) = dns_pin {
221 builder = builder.resolve(&hostname, std::net::SocketAddr::new(validated_ip, 0));
222 }
223 let client = builder.build().map_err(|err| {
224 AttemptError::Fatal(ToolError::execution_failed(format!(
225 "failed to build HTTP client: {err}"
226 )))
227 })?;
228 let response = client
229 .get(current_url.clone())
230 .header("Accept", options.accept)
231 .header("Accept-Language", "en-US,en;q=0.5")
232 .send()
233 .await
234 .map_err(|err| AttemptError::Transient(err.to_string()))?;
235
236 if !response.status().is_redirection() {
237 break response;
238 }
239 if redirects >= MAX_REDIRECTS {
240 return Err(AttemptError::Fatal(ToolError::execution_failed(
241 "request exceeded the five-redirect limit",
242 )));
243 }
244 let Some(location) = response
245 .headers()
246 .get(reqwest::header::LOCATION)
247 .and_then(|value| value.to_str().ok())
248 else {
249 break response;
250 };
251 current_url = response.url().join(location).map_err(|err| {
252 AttemptError::Fatal(ToolError::execution_failed(format!(
253 "invalid redirect location: {err}"
254 )))
255 })?;
256 redirects += 1;
257 };
258
259 let final_url = response.url().to_string();
260 let status = response.status().as_u16();
261 let content_type = response
262 .headers()
263 .get(reqwest::header::CONTENT_TYPE)
264 .and_then(|value| value.to_str().ok())
265 .unwrap_or("application/octet-stream")
266 .to_string();
267 let headers = response_headers(response.headers());
268 let mut stream = response.bytes_stream();
269 let mut bytes = Vec::with_capacity(options.max_bytes.min(64 * 1024));
270 let mut truncated = false;
271 while let Some(chunk) = stream.next().await {
272 let chunk = chunk.map_err(|err| AttemptError::Transient(err.to_string()))?;
273 let remaining = options.max_bytes.saturating_sub(bytes.len());
274 if chunk.len() > remaining {
275 bytes.extend_from_slice(&chunk[..remaining]);
276 truncated = true;
277 break;
278 }
279 bytes.extend_from_slice(&chunk);
280 if bytes.len() == options.max_bytes {
281 // A response exactly at the cap may be complete. Ask for one more
282 // chunk to distinguish exact length from actual truncation.
283 if let Some(next) = stream.next().await {
284 let next = next.map_err(|err| AttemptError::Transient(err.to_string()))?;
285 truncated = !next.is_empty();
286 }
287 break;
288 }
289 }
290
291 Ok(CachedFetch {
292 url: final_url,
293 status,
294 headers,
295 content_type,
296 bytes: Arc::new(bytes),
297 truncated,
298 redirects,
299 })
300 }
301
302 fn from_cached(payload: CachedFetch, cache_hit: bool, retries: usize) -> FetchedPayload {
303 FetchedPayload {
304 url: payload.url,
305 status: payload.status,
306 headers: payload.headers,
307 content_type: payload.content_type,
308 bytes: payload.bytes,
309 truncated: payload.truncated,
310 cache_hit,
311 retries,
312 redirects: payload.redirects,
313 }
314 }
315
316 fn response_headers(headers: &reqwest::header::HeaderMap) -> BTreeMap<String, String> {
317 headers
318 .iter()
319 .filter(|(name, _)| {
320 !matches!(
321 name.as_str(),
322 "authorization"
323 | "proxy-authorization"
324 | "cookie"
325 | "set-cookie"
326 | "set-cookie2"
327 | "x-api-key"
328 | "api-key"
329 )
330 })
331 .filter_map(|(name, value)| {
332 value
333 .to_str()
334 .ok()
335 .map(|value| (name.as_str().to_ascii_lowercase(), value.to_string()))
336 })
337 .collect()
338 }
339
340 fn is_transient_status(status: u16) -> bool {
341 (500..600).contains(&status)
342 }
343
344 fn retry_delay() -> Duration {
345 #[cfg(test)]
346 return Duration::ZERO;
347
348 #[cfg(not(test))]
349 {
350 let jitter_ms = SystemTime::now()
351 .duration_since(UNIX_EPOCH)
352 .map(|duration| u64::from(duration.subsec_nanos()) % 41)
353 .unwrap_or(0);
354 Duration::from_millis(30 + jitter_ms)
355 }
356 }
357
358 #[cfg(test)]
359 mod tests {
360 use std::collections::BTreeMap;
361 use std::net::{IpAddr, Ipv4Addr};
362 use std::sync::Arc;
363 use std::sync::atomic::{AtomicUsize, Ordering};
364
365 use serde_json::json;
366 use wiremock::matchers::{method, path};
367 use wiremock::{Mock, MockServer, Request, Respond, ResponseTemplate};
368
369 use super::*;
370
371 fn context(namespace: &str) -> ToolContext {
372 ToolContext::new(".").with_state_namespace(namespace)
373 }
374
375 fn pin() -> DnsPin {
376 Some((
377 "public.example".to_string(),
378 IpAddr::V4(Ipv4Addr::LOCALHOST),
379 ))
380 }
381
382 #[derive(Clone)]
383 struct FailOnce {
384 calls: Arc<AtomicUsize>,
385 }
386
387 impl Respond for FailOnce {
388 fn respond(&self, _request: &Request) -> ResponseTemplate {
389 if self.calls.fetch_add(1, Ordering::SeqCst) == 0 {
390 ResponseTemplate::new(503).set_body_json(json!({"error": "retry"}))
391 } else {
392 ResponseTemplate::new(200)
393 .insert_header("content-type", "text/plain")
394 .set_body_string("recovered response")
395 }
396 }
397 }
398
399 #[tokio::test]
400 async fn transient_server_error_retries_once_then_caches() {
401 let server = MockServer::start().await;
402 let calls = Arc::new(AtomicUsize::new(0));
403 Mock::given(method("GET"))
404 .and(path("/retry"))
405 .respond_with(FailOnce {
406 calls: Arc::clone(&calls),
407 })
408 .mount(&server)
409 .await;
410 let url = format!("http://public.example:{}/retry", server.address().port());
411 let options = FetchOptions::new(Duration::from_secs(5), 1_024, "text/plain");
412 let context = context("fetch-retry-cache");
413
414 let first = fetch_with_initial_pin(&url, &options, &context, "test", pin())
415 .await
416 .expect("retry succeeds");
417 assert_eq!(first.status, 200);
418 assert_eq!(first.retries, 1);
419 assert!(!first.cache_hit);
420 assert_eq!(&*first.bytes, b"recovered response");
421
422 let second = fetch_with_initial_pin(&url, &options, &context, "test", pin())
423 .await
424 .expect("cache hit");
425 assert!(second.cache_hit);
426 assert_eq!(calls.load(Ordering::SeqCst), 2);
427 }
428
429 #[tokio::test]
430 async fn truncated_cache_refetches_when_larger_body_is_requested() {
431 let server = MockServer::start().await;
432 Mock::given(method("GET"))
433 .and(path("/large"))
434 .respond_with(
435 ResponseTemplate::new(200)
436 .insert_header("content-type", "text/plain")
437 .set_body_string("0123456789"),
438 )
439 .mount(&server)
440 .await;
441 let url = format!("http://public.example:{}/large", server.address().port());
442 let context = context("fetch-truncated-refetch");
443
444 let small = fetch_with_initial_pin(
445 &url,
446 &FetchOptions::new(Duration::from_secs(5), 4, "text/plain"),
447 &context,
448 "test",
449 pin(),
450 )
451 .await
452 .expect("small fetch");
453 assert_eq!(&*small.bytes, b"0123");
454 assert!(small.truncated);
455
456 let large = fetch_with_initial_pin(
457 &url,
458 &FetchOptions::new(Duration::from_secs(5), 16, "text/plain"),
459 &context,
460 "test",
461 pin(),
462 )
463 .await
464 .expect("larger refetch");
465 assert_eq!(&*large.bytes, b"0123456789");
466 assert!(!large.truncated);
467 assert!(!large.cache_hit);
468 }
469
470 #[test]
471 fn fetch_user_agent_tracks_the_crate_version() {
472 assert!(
473 USER_AGENT.contains(concat!("codewhale/", env!("CARGO_PKG_VERSION"))),
474 "guarded-fetch UA must never pin a stale release: {USER_AGENT}"
475 );
476 }
477
478 #[test]
479 fn response_headers_drop_set_cookie_values() {
480 let mut headers = reqwest::header::HeaderMap::new();
481 headers.insert("content-type", "text/plain".parse().unwrap());
482 headers.insert("set-cookie", "session=secret".parse().unwrap());
483 headers.insert("x-api-key", "secret".parse().unwrap());
484
485 let filtered = response_headers(&headers);
486
487 assert_eq!(
488 filtered.get("content-type").map(String::as_str),
489 Some("text/plain")
490 );
491 assert!(!filtered.contains_key("set-cookie"));
492 assert!(!filtered.contains_key("x-api-key"));
493 }
494
495 #[tokio::test]
496 async fn tightened_network_policy_blocks_an_existing_cache_entry() {
497 use crate::network_policy::{Decision, NetworkPolicy, NetworkPolicyDecider};
498
499 let url = reqwest::Url::parse("https://example.com/cached").unwrap();
500 cache::insert(
501 "policy-cache",
502 &url,
503 "text/plain",
504 CachedFetch {
505 url: url.to_string(),
506 status: 200,
507 headers: BTreeMap::new(),
508 content_type: "text/plain".to_string(),
509 bytes: Arc::new(b"cached".to_vec()),
510 truncated: false,
511 redirects: 0,
512 },
513 );
514 let policy = NetworkPolicy {
515 default: Decision::Deny.into(),
516 allow: Vec::new(),
517 deny: Vec::new(),
518 proxy: Vec::new(),
519 proxy_fake_ip_cidrs: Vec::new(),
520 audit: false,
521 };
522 let context =
523 context("policy-cache").with_network_policy(NetworkPolicyDecider::new(policy, None));
524
525 let error = fetch(
526 url.as_str(),
527 &FetchOptions::new(Duration::from_secs(1), 100, "text/plain"),
528 &context,
529 "fetch_url",
530 )
531 .await
532 .expect_err("policy must win over cache");
533 assert!(error.to_string().contains("blocked by network policy"));
534 }
535
536 #[tokio::test]
537 async fn tightened_network_policy_checks_cached_redirect_destination() {
538 use crate::network_policy::{Decision, NetworkPolicy, NetworkPolicyDecider};
539
540 let initial_url = reqwest::Url::parse("https://8.8.8.8/cached").unwrap();
541 cache::insert(
542 "redirect-policy-cache",
543 &initial_url,
544 "text/plain",
545 CachedFetch {
546 url: "https://1.1.1.1/redirected".to_string(),
547 status: 200,
548 headers: BTreeMap::new(),
549 content_type: "text/plain".to_string(),
550 bytes: Arc::new(b"cached".to_vec()),
551 truncated: false,
552 redirects: 1,
553 },
554 );
555 let policy = NetworkPolicy {
556 default: Decision::Allow.into(),
557 allow: Vec::new(),
558 deny: vec!["1.1.1.1".to_string()],
559 proxy: Vec::new(),
560 proxy_fake_ip_cidrs: Vec::new(),
561 audit: false,
562 };
563 let context = context("redirect-policy-cache")
564 .with_network_policy(NetworkPolicyDecider::new(policy, None));
565
566 let error = fetch(
567 initial_url.as_str(),
568 &FetchOptions::new(Duration::from_secs(1), 100, "text/plain"),
569 &context,
570 "fetch_url",
571 )
572 .await
573 .expect_err("final redirect policy must win over cache");
574 assert!(error.to_string().contains("1.1.1.1"));
575 assert!(error.to_string().contains("blocked by network policy"));
576 }
577 }
578
578 lines RUST