Skip to main content

lc_tools/extended/
http.rs

1//! HTTP tool with SSRF protection
2
3use std::time::Duration;
4
5use async_trait::async_trait;
6use serde_json::Value;
7
8use crate::ssrf::{guarded_get, guarded_post_json, read_body_bounded, MAX_FETCH_BYTES};
9use lc_core::tools::{ToolError, ToolRiskProfile};
10use lc_core::BaseTool;
11
12/// HTTP request tool (GET/POST) with SSRF protection.
13pub struct HTTPTool {
14    /// Per-request timeout handed to the per-request pinned client.
15    timeout: Duration,
16    allow_private_ips: bool,
17}
18
19impl HTTPTool {
20    /// Creates an HTTP tool with a 30s timeout and SSRF protection enabled.
21    pub fn new() -> Self {
22        Self {
23            timeout: Duration::from_secs(30),
24            allow_private_ips: false,
25        }
26    }
27
28    /// Creates an HTTP tool with a custom timeout (SSRF protection enabled).
29    pub fn with_timeout(timeout: Duration) -> Self {
30        Self {
31            timeout,
32            allow_private_ips: false,
33        }
34    }
35
36    /// Allow requests to private/internal IP addresses (SSRF opt-in).
37    pub fn with_allow_private_ips(mut self, allow: bool) -> Self {
38        self.allow_private_ips = allow;
39        self
40    }
41
42    /// Sends a GET request, following redirects with SSRF checks per hop.
43    pub async fn get(&self, url: &str) -> Result<String, ToolError> {
44        // SSRF: guarded_get resolves once, validates every answer, pins the validated
45        // IPs, then follows redirects manually with the same treatment per hop.
46        let resp = guarded_get(url, !self.allow_private_ips, Some(self.timeout)).await?;
47        let (body, _) = read_body_bounded(resp, MAX_FETCH_BYTES).await?;
48        Ok(body)
49    }
50
51    /// Sends a POST request with a JSON body (single hop, IP-pinned SSRF guard).
52    pub async fn post(&self, url: &str, body: Value) -> Result<String, ToolError> {
53        // POST never follows redirects (the 3xx response is returned as-is); resolve-once
54        // plus IP pinning still closes the DNS-rebinding window on this single hop.
55        let resp =
56            guarded_post_json(url, &body, !self.allow_private_ips, Some(self.timeout)).await?;
57        let (text, _) = read_body_bounded(resp, MAX_FETCH_BYTES).await?;
58        Ok(text)
59    }
60}
61
62impl Default for HTTPTool {
63    fn default() -> Self {
64        Self::new()
65    }
66}
67
68#[async_trait]
69impl BaseTool for HTTPTool {
70    fn name(&self) -> &str {
71        "http_request"
72    }
73
74    fn description(&self) -> &str {
75        "Make HTTP requests. Input JSON: {\"url\": \"...\", \"method\": \"get|post\", \"body\": {...}}. \
76         SSRF protection enabled by default (blocks private IPs)."
77    }
78
79    /// S3: consumes raw, untrusted web content — declare the input risk so a policy
80    /// layer (Rule of Two) can weigh prompt-injection / tool-abuse surface.
81    fn risk(&self) -> ToolRiskProfile {
82        ToolRiskProfile {
83            untrusted_input: true,
84            ..Default::default()
85        }
86    }
87
88    async fn run(&self, input: String) -> Result<String, ToolError> {
89        let v: Value =
90            serde_json::from_str(&input).map_err(|e| ToolError::InvalidInput(e.to_string()))?;
91        let url = v
92            .get("url")
93            .and_then(|x| x.as_str())
94            .ok_or_else(|| ToolError::InvalidInput("Missing 'url' field".to_string()))?;
95        let method = v.get("method").and_then(|x| x.as_str()).unwrap_or("get");
96        match method {
97            "get" => self.get(url).await,
98            "post" => {
99                self.post(url, v.get("body").cloned().unwrap_or(Value::Null))
100                    .await
101            }
102            other => Err(ToolError::InvalidInput(format!(
103                "Unknown method: {}. Supported: get, post",
104                other
105            ))),
106        }
107    }
108}
109
110#[cfg(test)]
111mod tests {
112    use super::*;
113    use crate::ssrf::is_private_ip;
114    use std::net::IpAddr;
115
116    #[test]
117    fn test_name_description() {
118        let t = HTTPTool::new();
119        assert_eq!(t.name(), "http_request");
120        assert!(t.description().contains("HTTP"));
121    }
122
123    #[test]
124    fn test_private_ip_detection() {
125        assert!(is_private_ip(&IpAddr::from([127, 0, 0, 1])));
126        assert!(is_private_ip(&IpAddr::from([10, 0, 0, 1])));
127        assert!(is_private_ip(&IpAddr::from([172, 16, 0, 1])));
128        assert!(is_private_ip(&IpAddr::from([172, 31, 255, 255])));
129        assert!(is_private_ip(&IpAddr::from([192, 168, 1, 1])));
130        assert!(is_private_ip(&IpAddr::from([169, 254, 169, 254])));
131        assert!(is_private_ip(&IpAddr::from([0, 0, 0, 0])));
132
133        assert!(!is_private_ip(&IpAddr::from([8, 8, 8, 8])));
134        assert!(!is_private_ip(&IpAddr::from([1, 1, 1, 1])));
135        assert!(!is_private_ip(&IpAddr::from([172, 15, 0, 1])));
136        assert!(!is_private_ip(&IpAddr::from([172, 32, 0, 1])));
137    }
138
139    #[tokio::test]
140    async fn test_ssrf_blocks_localhost() {
141        // Rejection happens after resolve but before connect, so no listener is
142        // contacted — works offline even though Windows answers every loopback port.
143        let tool = HTTPTool::new();
144        let result = tool.get("http://127.0.0.1:6379/").await;
145        assert!(result.is_err());
146        assert!(result.unwrap_err().to_string().contains("SSRF"));
147    }
148
149    #[tokio::test]
150    async fn test_ssrf_blocks_cloud_metadata() {
151        let tool = HTTPTool::new();
152        let result = tool.get("http://169.254.169.254/latest/meta-data/").await;
153        assert!(result.is_err());
154        assert!(result.unwrap_err().to_string().contains("SSRF"));
155    }
156
157    #[test]
158    fn test_with_timeout_is_recorded() {
159        // The opt-in flag and custom timeout are plumbed into guarded_get as
160        // (!allow, Some(timeout)); assert the configuration side directly.
161        let tool = HTTPTool::with_timeout(Duration::from_secs(7)).with_allow_private_ips(true);
162        assert_eq!(tool.timeout, Duration::from_secs(7));
163        assert!(tool.allow_private_ips);
164    }
165
166    #[tokio::test]
167    async fn test_run_invalid_json() {
168        let t = HTTPTool::new();
169        assert!(t.run("not json".to_string()).await.is_err());
170    }
171
172    #[tokio::test]
173    async fn test_run_missing_url() {
174        let t = HTTPTool::new();
175        assert!(t.run(r#"{"method":"get"}"#.to_string()).await.is_err());
176    }
177
178    #[tokio::test]
179    async fn test_run_unknown_method() {
180        let t = HTTPTool::new();
181        let r = t
182            .run(r#"{"url":"http://x","method":"put"}"#.to_string())
183            .await;
184        assert!(r.is_err());
185    }
186}