lc_tools/
url_fetch.rs

1// lc-tools/src/url_fetch.rs
2//! Web page fetching tool
3//!
4//! Provides web content fetching and parsing.
5
6use async_trait::async_trait;
7use regex::Regex;
8use schemars::JsonSchema;
9use serde::{Deserialize, Serialize};
10
11use crate::ssrf::{guarded_get, read_body_bounded, MAX_FETCH_BYTES};
12use lc_core::tools::{BaseTool, Tool, ToolError, ToolRiskProfile};
13
14static SCRIPT_REGEX: std::sync::LazyLock<Regex> = std::sync::LazyLock::new(|| {
15    Regex::new(r"<script[^>]*>.*?</script>").expect("static regex literal must compile")
16});
17
18static STYLE_REGEX: std::sync::LazyLock<Regex> = std::sync::LazyLock::new(|| {
19    Regex::new(r"<style[^>]*>.*?</style>").expect("static regex literal must compile")
20});
21
22static TAG_REGEX: std::sync::LazyLock<Regex> =
23    std::sync::LazyLock::new(|| Regex::new(r"<[^>]+>").expect("static regex literal must compile"));
24
25static WHITESPACE_REGEX: std::sync::LazyLock<Regex> =
26    std::sync::LazyLock::new(|| Regex::new(r"\s+").expect("static regex literal must compile"));
27
28static LINK_REGEX: std::sync::LazyLock<Regex> = std::sync::LazyLock::new(|| {
29    Regex::new(r#"<a[^>]+href\s*=\s*['"]([^'"]+)['"][^>]*>"#)
30        .expect("static regex literal must compile")
31});
32
33static IMG_REGEX: std::sync::LazyLock<Regex> = std::sync::LazyLock::new(|| {
34    Regex::new(r#"<img[^>]+src\s*=\s*['"]([^'"]+)['"][^>]*>"#)
35        .expect("static regex literal must compile")
36});
37
38static TITLE_REGEX: std::sync::LazyLock<Regex> = std::sync::LazyLock::new(|| {
39    Regex::new(r"<title[^>]*>(.*?)</title>").expect("static regex literal must compile")
40});
41
42static DESC_REGEX: std::sync::LazyLock<Regex> = std::sync::LazyLock::new(|| {
43    Regex::new(r#"<meta[^>]+name\s*=\s*['"]description['"][^>]+content\s*=\s*['"]([^'"]+)['"]"#)
44        .expect("static regex literal must compile")
45});
46
47static KW_REGEX: std::sync::LazyLock<Regex> = std::sync::LazyLock::new(|| {
48    Regex::new(r#"<meta[^>]+name\s*=\s*['"]keywords['"][^>]+content\s*=\s*['"]([^'"]+)['"]"#)
49        .expect("static regex literal must compile")
50});
51
52/// Extracts links from HTML and dedups them (preserving first-occurrence order).
53///
54/// Returns `(deduped links, raw count)`. The raw count is used by details to show
55/// "before/after dedup". Q4: the old implementation only renamed the Vec to
56/// `unique_links` without actually deduplicating.
57fn extract_unique_links(html: &str) -> (Vec<String>, usize) {
58    let raw: Vec<String> = LINK_REGEX
59        .captures_iter(html)
60        .map(|cap| cap[1].to_string())
61        .collect();
62    let raw_count = raw.len();
63    let mut seen = std::collections::HashSet::new();
64    let unique: Vec<String> = raw
65        .into_iter()
66        .filter(|link| seen.insert(link.clone()))
67        .collect();
68    (unique, raw_count)
69}
70
71/// URLFetch tool input
72#[derive(Debug, Deserialize, JsonSchema)]
73pub struct URLFetchInput {
74    /// Operation type: "fetch", "extract_text", "extract_links", "extract_images", "metadata"
75    pub operation: String,
76
77    /// The URL
78    pub url: String,
79
80    /// Whether to include header info (for the fetch operation)
81    pub include_headers: Option<bool>,
82
83    /// Maximum content length (bytes)
84    pub max_length: Option<usize>,
85}
86
87/// URLFetch tool output
88#[derive(Debug, Serialize)]
89pub struct URLFetchOutput {
90    /// Operation result
91    pub result: String,
92
93    /// Operation type
94    pub operation: String,
95
96    /// URL
97    pub url: String,
98
99    /// Content length
100    pub content_length: usize,
101
102    /// Extra details
103    pub details: Option<String>,
104}
105
106/// Web page fetching tool
107pub struct URLFetchTool {
108    /// Whether access to private/internal IPs is allowed (default false)
109    allow_private_ips: bool,
110}
111
112impl URLFetchTool {
113    /// Creates a web fetching tool (SSRF protection enabled by default).
114    ///
115    /// The HTTP client is built per request inside `guarded_get` so the
116    /// validated DNS answers can be pinned onto it (A3 DNS-rebinding closure).
117    pub fn new() -> Self {
118        Self {
119            allow_private_ips: false,
120        }
121    }
122
123    /// Allow requests to private/internal IP addresses (SSRF opt-in).
124    pub fn with_allow_private_ips(mut self, allow: bool) -> Self {
125        self.allow_private_ips = allow;
126        self
127    }
128
129    /// Fetches web content
130    async fn fetch_url(
131        &self,
132        url: &str,
133        max_length: Option<usize>,
134        include_headers: Option<bool>,
135    ) -> Result<URLFetchOutput, ToolError> {
136        if !url.starts_with("http://") && !url.starts_with("https://") {
137            return Err(ToolError::InvalidInput(
138                "URL must start with http:// or https://".to_string(),
139            ));
140        }
141
142        // SSRF: guarded_get checks each hop, pins the validated IPs, and follows
143        // redirects manually (both the first hop and every redirect target are
144        // re-resolved and re-checked against intranet addresses).
145        let response = guarded_get(
146            url,
147            !self.allow_private_ips,
148            Some(std::time::Duration::from_secs(30)),
149        )
150        .await?;
151
152        let status = response.status();
153        if !status.is_success() {
154            return Err(ToolError::ExecutionFailed(format!(
155                "HTTP error: {} - {}",
156                status.as_u16(),
157                status.canonical_reason().unwrap_or("unknown")
158            )));
159        }
160
161        // Q3: when include_headers = Some(true), merge the response headers into the output
162        // (details) instead of silently ignoring them. Headers must be read before consuming
163        // the response (reqwest's headers() borrows, text() consumes).
164        let header_block: String = if include_headers.unwrap_or(false) {
165            let mut lines: Vec<String> = response
166                .headers()
167                .iter()
168                .map(|(name, value)| format!("{}: {}", name, value.to_str().unwrap_or("<非UTF-8>")))
169                .collect();
170            lines.sort();
171            let mut block = String::from("响应头:\n");
172            for line in lines {
173                block.push_str(&line);
174                block.push('\n');
175            }
176            block
177        } else {
178            String::new()
179        };
180
181        // M2: read the body with a hard cap, streaming — never buffer unbounded.
182        let max_len = max_length.unwrap_or(50_000).min(MAX_FETCH_BYTES);
183        let (body, truncated) = read_body_bounded(response, max_len).await?;
184        let content_len = body.len();
185        let result = if truncated {
186            format!("{body}\n... [内容已截断]")
187        } else {
188            body
189        };
190
191        let mut details = format!(
192            "状态码: {}, 内容长度: {} 字节{}",
193            status.as_u16(),
194            content_len,
195            if truncated { " (已截断)" } else { "" }
196        );
197        if !header_block.is_empty() {
198            details.push('\n');
199            details.push_str(&header_block);
200        }
201
202        Ok(URLFetchOutput {
203            result,
204            operation: "fetch".to_string(),
205            url: url.to_string(),
206            content_length: content_len,
207            details: Some(details),
208        })
209    }
210
211    /// Extracts plain text content
212    async fn extract_text(&self, url: &str) -> Result<URLFetchOutput, ToolError> {
213        let fetch_result = self.fetch_url(url, Some(100000), None).await?;
214        let html = &fetch_result.result;
215
216        let html = SCRIPT_REGEX.replace_all(html, "");
217        let html = STYLE_REGEX.replace_all(&html, "");
218
219        let text = TAG_REGEX.replace_all(&html, "");
220
221        let clean_text = WHITESPACE_REGEX.replace_all(&text, " ").trim().to_string();
222
223        let max_len = 5000;
224        let clean_len = clean_text.len();
225        let result = if clean_len > max_len {
226            clean_text.chars().take(max_len).collect::<String>() + "..."
227        } else {
228            clean_text
229        };
230
231        Ok(URLFetchOutput {
232            result,
233            operation: "extract_text".to_string(),
234            url: url.to_string(),
235            content_length: clean_len,
236            details: Some(format!("提取了 {} 字符的纯文本", clean_len)),
237        })
238    }
239
240    /// Extracts links (deduped, first-occurrence order preserved)
241    async fn extract_links(&self, url: &str) -> Result<URLFetchOutput, ToolError> {
242        let fetch_result = self.fetch_url(url, Some(100000), None).await?;
243        let html = &fetch_result.result;
244
245        let (unique_links, raw_count) = extract_unique_links(html);
246        let result = unique_links.join("\n");
247
248        Ok(URLFetchOutput {
249            result,
250            operation: "extract_links".to_string(),
251            url: url.to_string(),
252            content_length: html.len(), // Q4: the real body length, not the link count
253            details: Some(format!(
254                "找到 {} 个唯一链接(原始 {} 个)",
255                unique_links.len(),
256                raw_count
257            )),
258        })
259    }
260
261    /// Extracts image links
262    async fn extract_images(&self, url: &str) -> Result<URLFetchOutput, ToolError> {
263        let fetch_result = self.fetch_url(url, Some(100000), None).await?;
264        let html = &fetch_result.result;
265
266        let images: Vec<String> = IMG_REGEX
267            .captures_iter(html)
268            .map(|cap| cap[1].to_string())
269            .collect();
270
271        let result = images.join("\n");
272
273        Ok(URLFetchOutput {
274            result,
275            operation: "extract_images".to_string(),
276            url: url.to_string(),
277            content_length: images.len(),
278            details: Some(format!("找到 {} 张图片", images.len())),
279        })
280    }
281
282    /// Extracts metadata
283    async fn extract_metadata(&self, url: &str) -> Result<URLFetchOutput, ToolError> {
284        let fetch_result = self.fetch_url(url, Some(50000), None).await?;
285        let html = &fetch_result.result;
286
287        let title = TITLE_REGEX
288            .captures(html)
289            .map(|cap| cap[1].trim().to_string())
290            .unwrap_or_default();
291
292        let description = DESC_REGEX
293            .captures(html)
294            .map(|cap| cap[1].to_string())
295            .unwrap_or_default();
296
297        let keywords = KW_REGEX
298            .captures(html)
299            .map(|cap| cap[1].to_string())
300            .unwrap_or_default();
301
302        let result = format!(
303            "标题: {}\n描述: {}\n关键词: {}",
304            title, description, keywords
305        );
306
307        Ok(URLFetchOutput {
308            result,
309            operation: "metadata".to_string(),
310            url: url.to_string(),
311            content_length: title.len() + description.len() + keywords.len(),
312            details: Some("提取了网页元数据".to_string()),
313        })
314    }
315}
316
317impl Default for URLFetchTool {
318    fn default() -> Self {
319        Self::new()
320    }
321}
322
323#[async_trait]
324impl Tool for URLFetchTool {
325    type Input = URLFetchInput;
326    type Output = URLFetchOutput;
327
328    async fn invoke(&self, input: Self::Input) -> Result<Self::Output, ToolError> {
329        match input.operation.as_str() {
330            "fetch" => {
331                self.fetch_url(&input.url, input.max_length, input.include_headers)
332                    .await
333            }
334            "extract_text" => self.extract_text(&input.url).await,
335            "extract_links" => self.extract_links(&input.url).await,
336            "extract_images" => self.extract_images(&input.url).await,
337            "metadata" => self.extract_metadata(&input.url).await,
338            _ => Err(ToolError::InvalidInput(
339                format!("unsupported operation: {}, use: fetch, extract_text, extract_links, extract_images, metadata", input.operation)
340            )),
341        }
342    }
343}
344
345#[async_trait]
346impl BaseTool for URLFetchTool {
347    fn name(&self) -> &str {
348        "url_fetch"
349    }
350
351    fn description(&self) -> &str {
352        "网页抓取工具。支持多种操作：
353
354操作类型:
355- fetch: 抓取完整网页内容
356- extract_text: 提取纯文本内容（去除HTML标签）
357- extract_links: 提取所有链接
358- extract_images: 提取所有图片链接
359- metadata: 提取网页元数据（标题、描述、关键词）
360
361参数:
362- url: 网页地址（必须以 http:// 或 https:// 开头）
363- max_length: 最大内容长度（可选，默认50KB）
364- include_headers: 是否包含头部信息（可选）
365
366示例:
367- 抓取网页: {\"operation\": \"fetch\", \"url\": \"https://example.com\"}
368- 提取文本: {\"operation\": \"extract_text\", \"url\": \"https://example.com\"}
369- 提取链接: {\"operation\": \"extract_links\", \"url\": \"https://example.com\"}"
370    }
371
372    /// S3: this tool consumes raw, untrusted web content — declare the input risk so a
373    /// policy layer (Rule of Two) can weigh prompt-injection / tool-abuse surface.
374    fn risk(&self) -> ToolRiskProfile {
375        ToolRiskProfile {
376            untrusted_input: true,
377            ..Default::default()
378        }
379    }
380
381    async fn run(&self, input: String) -> Result<String, ToolError> {
382        let parsed: URLFetchInput = serde_json::from_str(&input)
383            .map_err(|e| ToolError::InvalidInput(format!("JSON parse failed: {}", e)))?;
384
385        let output = self.invoke(parsed).await?;
386
387        Ok(format!(
388            "URL: {}\n操作: {}\n内容长度: {} 字节\n\n{}\n详细信息: {}",
389            output.url,
390            output.operation,
391            output.content_length,
392            output.result,
393            output.details.unwrap_or_default()
394        ))
395    }
396
397    fn args_schema(&self) -> Option<serde_json::Value> {
398        use schemars::schema_for;
399        serde_json::to_value(schema_for!(URLFetchInput)).ok()
400    }
401}
402
403#[cfg(test)]
404mod tests {
405    use super::*;
406
407    #[test]
408    fn test_url_validation() {
409        let valid_url = "https://example.com";
410        assert!(valid_url.starts_with("http://") || valid_url.starts_with("https://"));
411
412        let valid_url2 = "http://example.org";
413        assert!(valid_url2.starts_with("http://") || valid_url2.starts_with("https://"));
414    }
415
416    #[tokio::test]
417    async fn test_url_fetch_invalid_url() {
418        let tool = URLFetchTool::new();
419
420        let input = URLFetchInput {
421            operation: "fetch".to_string(),
422            url: "invalid-url".to_string(),
423            include_headers: None,
424            max_length: None,
425        };
426
427        let result = tool.invoke(input).await;
428        assert!(result.is_err());
429        assert!(result.unwrap_err().to_string().contains("http://"));
430    }
431
432    /// Q4: extract_links truly dedups and preserves first-occurrence order; content_length is the body length.
433    #[test]
434    fn test_extract_unique_links_dedups() {
435        let html = r#"
436            <a href="https://a.com/1">first</a>
437            <a href="https://a.com/1">dup</a>
438            <a href="https://a.com/2">second</a>
439            <a href="https://a.com/1">dup2</a>
440        "#;
441        let (unique, raw) = extract_unique_links(html);
442        assert_eq!(raw, 4, "原始链接数应为 4");
443        assert_eq!(
444            unique,
445            vec!["https://a.com/1".to_string(), "https://a.com/2".to_string()],
446            "应去重且保持首次出现顺序"
447        );
448    }
449
450    /// Q1: after SSRF was extracted into a shared module, URLFetch still blocks intranet addresses by default.
451    #[tokio::test]
452    async fn test_url_fetch_blocks_localhost_by_default() {
453        let tool = URLFetchTool::new();
454        let input = URLFetchInput {
455            operation: "fetch".to_string(),
456            url: "http://127.0.0.1:6379/".to_string(),
457            include_headers: None,
458            max_length: None,
459        };
460        let result = tool.invoke(input).await;
461        assert!(result.is_err());
462        let err = result.unwrap_err().to_string();
463        assert!(err.contains("SSRF"), "expected SSRF error, got: {}", err);
464    }
465
466    /// Q3: when include_headers = Some(true), response headers are merged into the output (details).
467    #[tokio::test]
468    async fn test_fetch_include_headers() {
469        use tokio::io::{AsyncReadExt, AsyncWriteExt};
470        use tokio::net::TcpListener;
471
472        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
473        let addr = listener.local_addr().unwrap();
474
475        let server = tokio::spawn(async move {
476            let (mut socket, _) = listener.accept().await.unwrap();
477            let mut buf = [0u8; 4096];
478            let _ = socket.read(&mut buf).await;
479            let body = "hello world";
480            let resp = format!(
481                "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nContent-Type: text/plain\r\nX-Test-Header: hello\r\n\r\n{}",
482                body.len(),
483                body
484            );
485            socket.write_all(resp.as_bytes()).await.unwrap();
486            drop(socket);
487        });
488
489        let tool = URLFetchTool::new().with_allow_private_ips(true);
490        let input = URLFetchInput {
491            operation: "fetch".to_string(),
492            url: format!("http://{}/", addr),
493            include_headers: Some(true),
494            max_length: None,
495        };
496
497        let output = tool.invoke(input).await.unwrap();
498        let details = output.details.unwrap();
499        assert!(details.contains("响应头:"), "details: {}", details);
500        // reqwest normalizes response header names to lowercase (HTTP header names are case-insensitive)
501        assert!(
502            details.contains("x-test-header: hello"),
503            "details: {}",
504            details
505        );
506        assert!(output.result.contains("hello world"));
507        server.await.unwrap();
508    }
509
510    /// Q3: when include_headers = false/None, no response headers are returned.
511    #[tokio::test]
512    async fn test_fetch_without_include_headers() {
513        use tokio::io::{AsyncReadExt, AsyncWriteExt};
514        use tokio::net::TcpListener;
515
516        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
517        let addr = listener.local_addr().unwrap();
518
519        let server = tokio::spawn(async move {
520            let (mut socket, _) = listener.accept().await.unwrap();
521            let mut buf = [0u8; 4096];
522            let _ = socket.read(&mut buf).await;
523            let body = "hello world";
524            let resp = format!(
525                "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nX-Test-Header: hello\r\n\r\n{}",
526                body.len(),
527                body
528            );
529            socket.write_all(resp.as_bytes()).await.unwrap();
530            drop(socket);
531        });
532
533        let tool = URLFetchTool::new().with_allow_private_ips(true);
534        let input = URLFetchInput {
535            operation: "fetch".to_string(),
536            url: format!("http://{}/", addr),
537            include_headers: Some(false),
538            max_length: None,
539        };
540
541        let output = tool.invoke(input).await.unwrap();
542        let details = output.details.unwrap();
543        assert!(
544            !details.contains("X-Test-Header"),
545            "不应包含响应头, got: {}",
546            details
547        );
548        server.await.unwrap();
549    }
550
551    #[test]
552    fn test_tool_properties() {
553        let tool = URLFetchTool::new();
554
555        assert_eq!(tool.name(), "url_fetch");
556        assert!(tool.description().contains("fetch"));
557        assert!(BaseTool::args_schema(&tool).is_some());
558    }
559}
lc_tools/url_fetch.rs

lc_tools/
url_fetch.rs
