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

lc_tools/
url_fetch.rs
