brainos_backends/
fetch.rs

1//! URL-fetch backend — bounded HTTP GET + HTML→text reduction.
2//!
3//! Used by `cortex::actions::ActionDispatcher` to inline the body of
4//! user-pasted URLs alongside web search results. The fetcher imposes
5//! its own timeout and body cap so a runaway page can't blow the
6//! LLM context window or stall the answer.
7
8use std::net::IpAddr;
9use std::sync::Arc;
10use std::time::Duration;
11
12use cortex::actions::{ActionError, FetchedPage, UrlFetchBackend};
13use metrics::SubsystemMetrics;
14
15use crate::resilience::{http_breaker, resilient_send, CircuitBreaker};
16
17/// Hard cap on bytes read from any single page. Anything bigger is
18/// truncated *after* HTML stripping so most real pages still produce
19/// useful text.
20const MAX_BODY_BYTES: usize = 256 * 1024;
21/// Hard cap on cleaned-text bytes returned to the dispatcher.
22const MAX_TEXT_BYTES: usize = 4 * 1024;
23/// Connect/read timeout for a single fetch. Tight on purpose — the
24/// answering LLM is waiting and a slow site shouldn't hold up the
25/// whole response.
26const DEFAULT_FETCH_TIMEOUT: Duration = Duration::from_secs(8);
27
28/// Default HTTP fetcher: GET, follow redirects, strip HTML, return
29/// truncated text. No JavaScript execution, no auth, no cookies.
30pub struct BasicUrlFetcher {
31    client: reqwest::Client,
32    circuit_breaker: Arc<CircuitBreaker>,
33    max_retries: u32,
34    retry_base_ms: u64,
35    /// When false (production default), `assert_url_safe` rejects URLs whose
36    /// resolved address is in any reserved range — loopback, private,
37    /// link-local, etc. Tests and a deployment that legitimately points the
38    /// LLM at a localhost service flip this on.
39    allow_internal: bool,
40}
41
42impl BasicUrlFetcher {
43    pub fn new(
44        resilience: &brain::config::ResilienceConfig,
45    ) -> Result<Self, crate::error::BackendInitError> {
46        Self::new_with_metrics(resilience, None)
47    }
48
49    pub fn new_with_metrics(
50        resilience: &brain::config::ResilienceConfig,
51        metrics: Option<Arc<SubsystemMetrics>>,
52    ) -> Result<Self, crate::error::BackendInitError> {
53        let client = reqwest::Client::builder()
54            .timeout(DEFAULT_FETCH_TIMEOUT)
55            .user_agent("brainos/url-fetch (+https://github.com/keshavashiya/brain)")
56            .redirect(reqwest::redirect::Policy::limited(5))
57            .build()
58            .map_err(|e| crate::error::BackendInitError::HttpClient("URL fetch client", e))?;
59        let cb = http_breaker(
60            "url-fetch",
61            resilience.circuit_breaker_threshold,
62            resilience.circuit_breaker_cooldown_secs,
63            metrics,
64        );
65        Ok(Self {
66            client,
67            circuit_breaker: Arc::new(cb),
68            max_retries: resilience.max_retries,
69            retry_base_ms: resilience.retry_base_ms,
70            allow_internal: false,
71        })
72    }
73
74    /// Opt out of the SSRF guard for reserved-range targets. Used by tests
75    /// (mockito binds 127.0.0.1) and by deployments pointing the LLM at a
76    /// trusted internal service.
77    pub fn with_allow_internal(mut self, allow: bool) -> Self {
78        self.allow_internal = allow;
79        self
80    }
81}
82
83#[async_trait::async_trait]
84impl UrlFetchBackend for BasicUrlFetcher {
85    async fn fetch(&self, url: &str) -> Result<FetchedPage, ActionError> {
86        assert_url_safe(url, self.allow_internal).await?;
87        let client = self.client.clone();
88        let url_owned = url.to_string();
89        let response = resilient_send(
90            || client.get(&url_owned),
91            &self.circuit_breaker,
92            self.max_retries,
93            self.retry_base_ms,
94        )
95        .await
96        .map_err(|e| ActionError::ExecutionFailed(e.to_string()))?;
97
98        if !response.status().is_success() {
99            return Err(ActionError::ExecutionFailed(format!(
100                "{url} returned HTTP {}",
101                response.status()
102            )));
103        }
104
105        let final_url = response.url().to_string();
106        let bytes = response
107            .bytes()
108            .await
109            .map_err(|e| ActionError::ExecutionFailed(format!("body read failed: {e}")))?;
110
111        let head = &bytes[..bytes.len().min(MAX_BODY_BYTES)];
112        let raw = String::from_utf8_lossy(head);
113        let title = extract_title(&raw).unwrap_or_else(|| final_url.clone());
114        let text = html_to_text(&raw);
115        let text = if text.len() > MAX_TEXT_BYTES {
116            let mut t: String = text.chars().take(MAX_TEXT_BYTES).collect();
117            t.push_str("\n…[truncated]");
118            t
119        } else {
120            text
121        };
122
123        Ok(FetchedPage {
124            url: final_url,
125            title,
126            text,
127        })
128    }
129}
130
131/// Issue 122: SSRF guard. The fetcher is reachable from LLM-controlled
132/// inputs (web-search hits, user-pasted URLs in chat) so we have to assume
133/// the URL is adversarial. Reject:
134///   * non-`http(s)` schemes (file:, gopher:, ftp:, …)
135///   * embedded credentials (`user:pass@`) which most reqwest builds honor
136///   * IP-literal hosts in any reserved range
137///   * DNS hostnames whose resolved addresses include any reserved range
138///
139/// Remaining gap: a TOCTOU between this resolve and reqwest's own resolve
140/// could theoretically be exploited by DNS rebinding. Closing it would
141/// require pinning the connect target to a vetted IP via a custom
142/// reqwest::dns::Resolve — deferred. The current resolver still raises
143/// the floor materially.
144pub(crate) async fn assert_url_safe(url: &str, allow_internal: bool) -> Result<(), ActionError> {
145    let parsed = url::Url::parse(url)
146        .map_err(|e| ActionError::InvalidArguments(format!("invalid url '{url}': {e}")))?;
147
148    match parsed.scheme() {
149        "http" | "https" => {}
150        other => {
151            return Err(ActionError::InvalidArguments(format!(
152                "url-fetch requires http(s):// scheme, got {other}://"
153            )));
154        }
155    }
156
157    if !parsed.username().is_empty() || parsed.password().is_some() {
158        return Err(ActionError::InvalidArguments(
159            "url-fetch refuses URLs with embedded credentials".into(),
160        ));
161    }
162
163    let host = parsed
164        .host()
165        .ok_or_else(|| ActionError::InvalidArguments("url-fetch requires a host".into()))?;
166
167    match host {
168        url::Host::Ipv4(v4) => {
169            reject_reserved_ip(IpAddr::V4(v4), &host.to_string(), allow_internal)
170        }
171        url::Host::Ipv6(v6) => {
172            reject_reserved_ip(IpAddr::V6(v6), &host.to_string(), allow_internal)
173        }
174        url::Host::Domain(name) => {
175            // Hostname — resolve and reject if any answer is reserved.
176            // `lookup_host` needs a `host:port` string; synthesise a port
177            // if the URL omits one.
178            let port = parsed.port_or_known_default().unwrap_or(0);
179            let addrs = tokio::net::lookup_host((name, port))
180                .await
181                .map_err(|e| {
182                    ActionError::ExecutionFailed(format!("dns lookup for {name} failed: {e}"))
183                })?
184                .collect::<Vec<_>>();
185
186            if addrs.is_empty() {
187                return Err(ActionError::ExecutionFailed(format!(
188                    "dns lookup for {name} returned no addresses"
189                )));
190            }
191            for sa in addrs {
192                reject_reserved_ip(sa.ip(), name, allow_internal)?;
193            }
194            Ok(())
195        }
196    }
197}
198
199fn reject_reserved_ip(ip: IpAddr, host: &str, allow_internal: bool) -> Result<(), ActionError> {
200    if allow_internal {
201        return Ok(());
202    }
203    if is_reserved(ip) {
204        return Err(ActionError::InvalidArguments(format!(
205            "url-fetch refuses host '{host}' resolving to reserved address {ip}"
206        )));
207    }
208    Ok(())
209}
210
211fn is_reserved(ip: IpAddr) -> bool {
212    if ip.is_loopback() || ip.is_unspecified() || ip.is_multicast() {
213        return true;
214    }
215    match ip {
216        IpAddr::V4(v4) => {
217            if v4.is_private() || v4.is_link_local() || v4.is_broadcast() || v4.is_documentation() {
218                return true;
219            }
220            // Cloud metadata endpoints (AWS / GCP / Azure all live at 169.254.169.254).
221            // 169.254.0.0/16 is link-local and already caught above; explicit for clarity.
222            if v4.octets()[0] == 169 && v4.octets()[1] == 254 {
223                return true;
224            }
225            // Carrier-grade NAT (RFC 6598).
226            if v4.octets()[0] == 100 && (64..=127).contains(&v4.octets()[1]) {
227                return true;
228            }
229            // Reserved for future use (240.0.0.0/4) — `is_reserved` is unstable; check manually.
230            if v4.octets()[0] >= 240 {
231                return true;
232            }
233            false
234        }
235        IpAddr::V6(v6) => {
236            // Unique local (fc00::/7).
237            if (v6.segments()[0] & 0xfe00) == 0xfc00 {
238                return true;
239            }
240            // Link-local (fe80::/10).
241            if (v6.segments()[0] & 0xffc0) == 0xfe80 {
242                return true;
243            }
244            // IPv4-mapped — recurse on the embedded IPv4 so we reject
245            // ::ffff:127.0.0.1 etc.
246            if let Some(v4) = v6.to_ipv4_mapped() {
247                return is_reserved(IpAddr::V4(v4));
248            }
249            false
250        }
251    }
252}
253
254/// Pull the `<title>` content from raw HTML. Case-insensitive on the
255/// tag name so `<TITLE>` and `<Title>` both match.
256fn extract_title(html: &str) -> Option<String> {
257    let lower = html.to_ascii_lowercase();
258    let start = lower.find("<title")?;
259    // Skip past the closing `>` of the opening tag (which may include attrs).
260    let after_open = lower[start..].find('>').map(|i| start + i + 1)?;
261    let end_rel = lower[after_open..].find("</title>")?;
262    let raw = &html[after_open..after_open + end_rel];
263    let cleaned = decode_entities(raw).trim().to_string();
264    if cleaned.is_empty() {
265        None
266    } else {
267        Some(cleaned)
268    }
269}
270
271/// Strip HTML tags + scripts/styles + decode common entities. Not a
272/// full parser — good enough to feed an LLM a readable excerpt.
273fn html_to_text(html: &str) -> String {
274    let mut s = String::with_capacity(html.len());
275    let stripped = strip_block(html, "script");
276    let stripped = strip_block(&stripped, "style");
277    let stripped = strip_block(&stripped, "noscript");
278    let mut in_tag = false;
279    for ch in stripped.chars() {
280        match ch {
281            '<' => in_tag = true,
282            '>' => {
283                in_tag = false;
284                s.push(' ');
285            }
286            _ if in_tag => {}
287            _ => s.push(ch),
288        }
289    }
290    let decoded = decode_entities(&s);
291    collapse_whitespace(&decoded)
292}
293
294/// Remove every `<tag …>…</tag>` block (case-insensitive on the tag).
295fn strip_block(input: &str, tag: &str) -> String {
296    let lower = input.to_ascii_lowercase();
297    let open_pat = format!("<{tag}");
298    let close_pat = format!("</{tag}>");
299    let mut out = String::with_capacity(input.len());
300    let mut cursor = 0usize;
301    while cursor < input.len() {
302        match lower[cursor..].find(&open_pat) {
303            None => {
304                out.push_str(&input[cursor..]);
305                break;
306            }
307            Some(rel) => {
308                let abs_open = cursor + rel;
309                out.push_str(&input[cursor..abs_open]);
310                let after_open = match lower[abs_open..].find('>') {
311                    Some(i) => abs_open + i + 1,
312                    None => break,
313                };
314                match lower[after_open..].find(&close_pat) {
315                    Some(end_rel) => {
316                        cursor = after_open + end_rel + close_pat.len();
317                    }
318                    None => {
319                        cursor = after_open;
320                    }
321                }
322            }
323        }
324    }
325    out
326}
327
328/// Decode the handful of HTML entities common enough to matter for
329/// readable LLM context. Anything else is passed through verbatim.
330fn decode_entities(s: &str) -> String {
331    s.replace("&nbsp;", " ")
332        .replace("&amp;", "&")
333        .replace("&lt;", "<")
334        .replace("&gt;", ">")
335        .replace("&quot;", "\"")
336        .replace("&#39;", "'")
337        .replace("&apos;", "'")
338}
339
340/// Collapse runs of whitespace (including newlines) into single spaces,
341/// preserving paragraph breaks (a sequence with two or more newlines
342/// becomes a single `\n`).
343fn collapse_whitespace(s: &str) -> String {
344    let mut out = String::with_capacity(s.len());
345    let mut last_was_space = false;
346    let mut newline_run = 0usize;
347    for ch in s.chars() {
348        if ch == '\n' {
349            newline_run += 1;
350            continue;
351        }
352        if ch.is_whitespace() {
353            if !last_was_space && newline_run == 0 {
354                out.push(' ');
355                last_was_space = true;
356            }
357            continue;
358        }
359        if newline_run >= 2 {
360            out.push('\n');
361        } else if newline_run == 1 && !last_was_space {
362            out.push(' ');
363        }
364        newline_run = 0;
365        last_was_space = false;
366        out.push(ch);
367    }
368    out.trim().to_string()
369}
370
371#[cfg(test)]
372mod tests {
373    use super::*;
374
375    #[test]
376    fn extract_title_handles_attrs_and_case() {
377        let html = "<html><head><TITLE id=\"x\">Hello World</TITLE></head><body/></html>";
378        assert_eq!(extract_title(html).as_deref(), Some("Hello World"));
379    }
380
381    #[test]
382    fn extract_title_returns_none_when_missing() {
383        assert!(extract_title("<html><body>no title</body></html>").is_none());
384    }
385
386    #[test]
387    fn html_to_text_strips_scripts_and_styles() {
388        let html = "<html><head><style>body{color:red}</style></head>\
389                    <body><script>alert(1)</script><p>Hello <b>world</b>!</p></body></html>";
390        let txt = html_to_text(html);
391        assert!(!txt.contains("alert"));
392        assert!(!txt.contains("color:red"));
393        assert!(txt.contains("Hello"));
394        assert!(txt.contains("world"));
395    }
396
397    #[test]
398    fn html_to_text_decodes_common_entities() {
399        let html = "<p>tom &amp; jerry &lt;3 &nbsp; coffee</p>";
400        let txt = html_to_text(html);
401        assert!(txt.contains("tom & jerry"));
402        assert!(txt.contains("<3"));
403    }
404
405    #[test]
406    fn collapse_whitespace_keeps_paragraph_breaks() {
407        let s = "alpha\n\nbeta   gamma\n\n\ndelta";
408        let out = collapse_whitespace(s);
409        // Multi-newline sequences collapse to a single \n; runs of spaces collapse to one.
410        assert_eq!(out, "alpha\nbeta gamma\ndelta");
411    }
412
413    fn fast_resilience() -> brain::config::ResilienceConfig {
414        brain::config::ResilienceConfig {
415            max_retries: 0,
416            retry_base_ms: 1,
417            circuit_breaker_threshold: 5,
418            circuit_breaker_cooldown_secs: 60,
419        }
420    }
421
422    #[tokio::test]
423    async fn fetch_returns_text_and_title() {
424        let mut server = mockito::Server::new_async().await;
425        let _mock = server
426            .mock("GET", "/page")
427            .with_status(200)
428            .with_header("content-type", "text/html; charset=utf-8")
429            .with_body(
430                "<html><head><title>Demo Page</title></head>\
431                 <body><script>x=1</script><h1>Hi</h1><p>Some &amp; text.</p></body></html>",
432            )
433            .create_async()
434            .await;
435
436        let fetcher = BasicUrlFetcher::new(&fast_resilience())
437            .unwrap()
438            .with_allow_internal(true);
439        let url = format!("{}/page", server.url());
440        let page = fetcher.fetch(&url).await.unwrap();
441        assert_eq!(page.title, "Demo Page");
442        assert!(page.text.contains("Hi"));
443        assert!(page.text.contains("Some & text"));
444        assert!(!page.text.contains("x=1"));
445    }
446
447    #[tokio::test]
448    async fn fetch_rejects_non_http_schemes() {
449        let fetcher = BasicUrlFetcher::new(&fast_resilience()).unwrap();
450        let err = fetcher.fetch("file:///etc/passwd").await.unwrap_err();
451        assert!(matches!(err, ActionError::InvalidArguments(_)));
452    }
453
454    #[tokio::test]
455    async fn fetch_surfaces_non_2xx_as_error() {
456        let mut server = mockito::Server::new_async().await;
457        let _mock = server
458            .mock("GET", "/missing")
459            .with_status(404)
460            .create_async()
461            .await;
462        let fetcher = BasicUrlFetcher::new(&fast_resilience())
463            .unwrap()
464            .with_allow_internal(true);
465        let url = format!("{}/missing", server.url());
466        assert!(fetcher.fetch(&url).await.is_err());
467    }
468
469    #[tokio::test]
470    async fn fetch_rejects_loopback_by_default() {
471        let fetcher = BasicUrlFetcher::new(&fast_resilience()).unwrap();
472        let err = fetcher.fetch("http://127.0.0.1:9").await.unwrap_err();
473        assert!(matches!(err, ActionError::InvalidArguments(_)), "{err:?}");
474    }
475
476    #[tokio::test]
477    async fn fetch_rejects_private_ranges() {
478        let fetcher = BasicUrlFetcher::new(&fast_resilience()).unwrap();
479        for url in [
480            "http://10.0.0.1/",
481            "http://192.168.1.1/",
482            "http://172.16.0.1/",
483            "http://169.254.169.254/latest/meta-data/",
484            "http://[::1]/",
485            "http://[::ffff:127.0.0.1]/",
486        ] {
487            let err = fetcher.fetch(url).await.unwrap_err();
488            assert!(
489                matches!(err, ActionError::InvalidArguments(_)),
490                "{url} should have been rejected (got {err:?})"
491            );
492        }
493    }
494
495    #[tokio::test]
496    async fn fetch_rejects_embedded_credentials() {
497        let fetcher = BasicUrlFetcher::new(&fast_resilience())
498            .unwrap()
499            .with_allow_internal(true);
500        let err = fetcher
501            .fetch("http://user:pass@127.0.0.1:9/")
502            .await
503            .unwrap_err();
504        assert!(matches!(err, ActionError::InvalidArguments(_)), "{err:?}");
505    }
506
507    #[tokio::test]
508    async fn fetch_rejects_unspecified_address() {
509        let fetcher = BasicUrlFetcher::new(&fast_resilience()).unwrap();
510        let err = fetcher.fetch("http://0.0.0.0/").await.unwrap_err();
511        assert!(matches!(err, ActionError::InvalidArguments(_)), "{err:?}");
512    }
513}
brainos_backends/fetch.rs

brainos_backends/
fetch.rs
