yggdrasil/api/
url_fetch.rs1#![cfg(feature = "server")]
26
27use std::net::{IpAddr, SocketAddr, ToSocketAddrs};
28use std::time::Duration;
29
30use bytes::Bytes;
31use http::Uri;
32
33use crate::api::upload::{process_image_upload, UploadError, UploadOutcome};
34use crate::utils::server::MAX_FILE_SIZE;
35
36pub(crate) enum FetchError {
38 Invalid(String),
40 BadStatus(String),
42 TooLarge,
44 Fetch(&'static str),
46}
47
48impl FetchError {
49 fn invalid<E: std::fmt::Display>(e: E, ctx: &str) -> Self {
50 tracing::warn!("url_fetch {ctx}: {e}");
51 FetchError::Invalid(format!("无效的图片 URL:{ctx}"))
52 }
53}
54
55pub(crate) async fn fetch_and_ingest(url: &str) -> Result<UploadOutcome, FetchError> {
59 let uri: Uri = url
61 .parse()
62 .map_err(|e| FetchError::invalid(e, "URL 格式"))?;
63 let scheme = uri.scheme_str().unwrap_or("");
64 if scheme != "https" {
65 return Err(FetchError::Invalid("仅支持 https:// 图片 URL".into()));
66 }
67 let host = uri
68 .host()
69 .ok_or_else(|| FetchError::Invalid("URL 缺少主机名".into()))?;
70 if host.is_empty() {
71 return Err(FetchError::Invalid("URL 缺少主机名".into()));
72 }
73
74 let host_owned = host.to_string();
77 let port = uri.port_u16().unwrap_or(443);
78 let resolve_target = format!("{host_owned}:{port}");
79 let addrs = tokio::task::spawn_blocking(move || {
80 resolve_target
81 .to_socket_addrs()
82 .map(|i| i.collect::<Vec<SocketAddr>>())
83 })
84 .await
85 .map_err(|e| FetchError::invalid(e, "解析任务失败"))?
86 .map_err(|e| FetchError::invalid(e, "DNS 解析失败"))?;
87
88 if addrs.is_empty() {
89 return Err(FetchError::Invalid("DNS 解析无结果".into()));
90 }
91
92 for addr in &addrs {
94 let ip = addr.ip();
95 if is_forbidden_ip(&ip) {
96 tracing::warn!("url_fetch SSRF 拒绝:{} → {}", host_owned, ip);
97 return Err(FetchError::Invalid("目标地址位于禁止的网络段".into()));
98 }
99 }
100 let locked_ip = addrs[0].ip();
102
103 let client = reqwest::Client::builder()
105 .redirect(reqwest::redirect::Policy::none())
106 .resolve(&host_owned, SocketAddr::from((locked_ip, port)))
107 .timeout(Duration::from_secs(30))
108 .connect_timeout(Duration::from_secs(10))
109 .user_agent("yggdrasil-mcp-media/1.0")
110 .build()
111 .map_err(|e| {
112 tracing::error!("url_fetch client build: {e}");
113 FetchError::Fetch("构建抓取客户端失败")
114 })?;
115
116 let resp = client.get(url).send().await.map_err(|e| {
118 tracing::warn!("url_fetch GET {url} failed: {e}");
119 FetchError::Fetch("抓取图片失败")
120 })?;
121
122 if !resp.status().is_success() {
123 return Err(FetchError::BadStatus(format!("远端返回 {}", resp.status())));
124 }
125
126 let mut total: usize = 0;
129 let mut buf: Vec<u8> = Vec::new();
130 let mut stream = resp.bytes_stream();
131 use futures::StreamExt;
132 while let Some(chunk) = stream.next().await {
133 let chunk = chunk.map_err(|e| {
134 tracing::warn!("url_fetch stream read failed: {e}");
135 FetchError::Fetch("读取图片数据失败")
136 })?;
137 total += chunk.len();
138 if total > MAX_FILE_SIZE {
139 return Err(FetchError::TooLarge);
140 }
141 buf.extend_from_slice(&chunk);
142 }
143
144 let data = Bytes::from(buf);
145
146 let original_filename = uri
148 .path()
149 .rsplit('/')
150 .next()
151 .filter(|s| !s.is_empty())
152 .map(|s| s.to_string());
153
154 process_image_upload(data, original_filename)
156 .await
157 .map_err(|e| match e {
158 UploadError::TooLarge => FetchError::TooLarge,
159 other => {
160 let ctx = match other {
161 UploadError::Empty => "空响应",
162 UploadError::BadType => "非图片或格式不支持",
163 UploadError::Oversized => "图片尺寸超限",
164 UploadError::Corrupt => "图片损坏",
165 UploadError::TooLarge => "文件过大",
166 UploadError::Internal(_) => "入库失败",
167 };
168 FetchError::BadStatus(ctx.into())
169 }
170 })
171}
172
173fn is_forbidden_ip(ip: &IpAddr) -> bool {
175 match ip {
176 IpAddr::V4(v4) => {
177 v4.is_private() || v4.is_loopback() || v4.is_link_local() || v4.is_unspecified() || v4.is_broadcast() || v4.is_documentation() || {
185 let o = v4.octets();
186 o[0] == 100 && o[1] >= 64 && o[1] <= 127
187 }
188 }
189 IpAddr::V6(v6) => {
190 v6.is_loopback() || v6.is_unspecified() || v6.is_multicast() || {
194 let s = v6.segments();
195 (s[0] & 0xfe00) == 0xfc00 || (s[0] & 0xffc0) == 0xfe80 }
198 }
199 }
200}
201
202#[cfg(all(test, feature = "server"))]
203mod tests {
204 use super::is_forbidden_ip;
205 use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
206
207 #[test]
208 fn forbidden_ipv4_ranges() {
209 assert!(is_forbidden_ip(&IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1))));
210 assert!(is_forbidden_ip(&IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1))));
211 assert!(is_forbidden_ip(&IpAddr::V4(Ipv4Addr::new(172, 16, 0, 1))));
212 assert!(is_forbidden_ip(&IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1))));
213 assert!(is_forbidden_ip(&IpAddr::V4(Ipv4Addr::new(169, 254, 1, 1))));
214 assert!(is_forbidden_ip(&IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0))));
215 assert!(is_forbidden_ip(&IpAddr::V4(Ipv4Addr::new(100, 64, 0, 1)))); assert!(is_forbidden_ip(&IpAddr::V4(Ipv4Addr::new(
217 100, 127, 255, 255
218 )))); assert!(is_forbidden_ip(&IpAddr::V4(Ipv4Addr::new(192, 0, 2, 1)))); }
221
222 #[test]
223 fn allowed_ipv4_ranges() {
224 assert!(!is_forbidden_ip(&IpAddr::V4(Ipv4Addr::new(8, 8, 8, 8))));
225 assert!(!is_forbidden_ip(&IpAddr::V4(Ipv4Addr::new(1, 1, 1, 1))));
226 assert!(!is_forbidden_ip(&IpAddr::V4(Ipv4Addr::new(99, 63, 0, 1)))); assert!(!is_forbidden_ip(&IpAddr::V4(Ipv4Addr::new(
228 100, 63, 255, 255
229 )))); assert!(!is_forbidden_ip(&IpAddr::V4(Ipv4Addr::new(100, 128, 0, 1)))); assert!(!is_forbidden_ip(&IpAddr::V4(Ipv4Addr::new(
232 172, 15, 255, 255
233 )))); assert!(!is_forbidden_ip(&IpAddr::V4(Ipv4Addr::new(172, 32, 0, 1)))); }
236
237 #[test]
238 fn forbidden_ipv6_ranges() {
239 assert!(is_forbidden_ip(&IpAddr::V6(Ipv6Addr::LOCALHOST))); assert!(is_forbidden_ip(&IpAddr::V6(Ipv6Addr::UNSPECIFIED))); assert!(is_forbidden_ip(&IpAddr::V6("ff02::1".parse().unwrap()))); assert!(is_forbidden_ip(&IpAddr::V6("fc00::1".parse().unwrap()))); assert!(is_forbidden_ip(&IpAddr::V6("fd12::1".parse().unwrap()))); assert!(is_forbidden_ip(&IpAddr::V6("fe80::1".parse().unwrap()))); }
246
247 #[test]
248 fn allowed_ipv6_ranges() {
249 assert!(!is_forbidden_ip(&IpAddr::V6(
250 "2606:4700::1".parse().unwrap()
251 ))); assert!(!is_forbidden_ip(&IpAddr::V6(
253 "2001:4860:4860::8888".parse().unwrap()
254 ))); }
256}