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(
59 url: &str,
60 alt: Option<String>,
61) -> Result<UploadOutcome, FetchError> {
62 let uri: Uri = url
64 .parse()
65 .map_err(|e| FetchError::invalid(e, "URL 格式"))?;
66 let scheme = uri.scheme_str().unwrap_or("");
67 if scheme != "https" {
68 return Err(FetchError::Invalid("仅支持 https:// 图片 URL".into()));
69 }
70 let host = uri
71 .host()
72 .ok_or_else(|| FetchError::Invalid("URL 缺少主机名".into()))?;
73 if host.is_empty() {
74 return Err(FetchError::Invalid("URL 缺少主机名".into()));
75 }
76
77 let host_owned = host.to_string();
80 let port = uri.port_u16().unwrap_or(443);
81 let resolve_target = format!("{host_owned}:{port}");
82 let addrs = tokio::task::spawn_blocking(move || {
83 resolve_target
84 .to_socket_addrs()
85 .map(|i| i.collect::<Vec<SocketAddr>>())
86 })
87 .await
88 .map_err(|e| FetchError::invalid(e, "解析任务失败"))?
89 .map_err(|e| FetchError::invalid(e, "DNS 解析失败"))?;
90
91 if addrs.is_empty() {
92 return Err(FetchError::Invalid("DNS 解析无结果".into()));
93 }
94
95 for addr in &addrs {
97 let ip = addr.ip();
98 if is_forbidden_ip(&ip) {
99 tracing::warn!("url_fetch SSRF 拒绝:{} → {}", host_owned, ip);
100 return Err(FetchError::Invalid("目标地址位于禁止的网络段".into()));
101 }
102 }
103 let locked_ip = addrs[0].ip();
105
106 let client = reqwest::Client::builder()
108 .redirect(reqwest::redirect::Policy::none())
109 .resolve(&host_owned, SocketAddr::from((locked_ip, port)))
110 .timeout(Duration::from_secs(30))
111 .connect_timeout(Duration::from_secs(10))
112 .user_agent("yggdrasil-mcp-media/1.0")
113 .build()
114 .map_err(|e| {
115 tracing::error!("url_fetch client build: {e}");
116 FetchError::Fetch("构建抓取客户端失败")
117 })?;
118
119 let resp = client.get(url).send().await.map_err(|e| {
121 tracing::warn!("url_fetch GET {url} failed: {e}");
122 FetchError::Fetch("抓取图片失败")
123 })?;
124
125 if !resp.status().is_success() {
126 return Err(FetchError::BadStatus(format!("远端返回 {}", resp.status())));
127 }
128
129 let mut total: usize = 0;
132 let mut buf: Vec<u8> = Vec::new();
133 let mut stream = resp.bytes_stream();
134 use futures::StreamExt;
135 while let Some(chunk) = stream.next().await {
136 let chunk = chunk.map_err(|e| {
137 tracing::warn!("url_fetch stream read failed: {e}");
138 FetchError::Fetch("读取图片数据失败")
139 })?;
140 total += chunk.len();
141 if total > MAX_FILE_SIZE {
142 return Err(FetchError::TooLarge);
143 }
144 buf.extend_from_slice(&chunk);
145 }
146
147 let data = Bytes::from(buf);
148
149 let original_filename = uri
151 .path()
152 .rsplit('/')
153 .next()
154 .filter(|s| !s.is_empty())
155 .map(|s| s.to_string());
156
157 process_image_upload(data, original_filename, alt)
159 .await
160 .map_err(|e| match e {
161 UploadError::TooLarge => FetchError::TooLarge,
162 other => {
163 let ctx = match other {
164 UploadError::Empty => "空响应",
165 UploadError::BadType => "非图片或格式不支持",
166 UploadError::Oversized => "图片尺寸超限",
167 UploadError::Corrupt => "图片损坏",
168 UploadError::TooLarge => "文件过大",
169 UploadError::Internal(_) => "入库失败",
170 };
171 FetchError::BadStatus(ctx.into())
172 }
173 })
174}
175
176fn is_forbidden_ip(ip: &IpAddr) -> bool {
178 match ip {
179 IpAddr::V4(v4) => {
180 v4.is_private() || v4.is_loopback() || v4.is_link_local() || v4.is_unspecified() || v4.is_broadcast() || v4.is_documentation() || {
188 let o = v4.octets();
189 o[0] == 100 && o[1] >= 64 && o[1] <= 127
190 }
191 }
192 IpAddr::V6(v6) => {
193 if v6.is_loopback() || v6.is_unspecified() {
194 return true;
195 }
196 if let Some(v4) = v6.to_ipv4() {
199 return is_forbidden_ip(&IpAddr::V4(v4));
200 }
201 v6.is_multicast() || {
203 let s = v6.segments();
204 (s[0] & 0xfe00) == 0xfc00 || (s[0] & 0xffc0) == 0xfe80 }
207 }
208 }
209}
210
211#[cfg(all(test, feature = "server"))]
212mod tests {
213 use super::is_forbidden_ip;
214 use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
215
216 #[test]
217 fn forbidden_ipv4_ranges() {
218 assert!(is_forbidden_ip(&IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1))));
219 assert!(is_forbidden_ip(&IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1))));
220 assert!(is_forbidden_ip(&IpAddr::V4(Ipv4Addr::new(172, 16, 0, 1))));
221 assert!(is_forbidden_ip(&IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1))));
222 assert!(is_forbidden_ip(&IpAddr::V4(Ipv4Addr::new(169, 254, 1, 1))));
223 assert!(is_forbidden_ip(&IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0))));
224 assert!(is_forbidden_ip(&IpAddr::V4(Ipv4Addr::new(100, 64, 0, 1)))); assert!(is_forbidden_ip(&IpAddr::V4(Ipv4Addr::new(
226 100, 127, 255, 255
227 )))); assert!(is_forbidden_ip(&IpAddr::V4(Ipv4Addr::new(192, 0, 2, 1)))); }
230
231 #[test]
232 fn allowed_ipv4_ranges() {
233 assert!(!is_forbidden_ip(&IpAddr::V4(Ipv4Addr::new(8, 8, 8, 8))));
234 assert!(!is_forbidden_ip(&IpAddr::V4(Ipv4Addr::new(1, 1, 1, 1))));
235 assert!(!is_forbidden_ip(&IpAddr::V4(Ipv4Addr::new(99, 63, 0, 1)))); assert!(!is_forbidden_ip(&IpAddr::V4(Ipv4Addr::new(
237 100, 63, 255, 255
238 )))); assert!(!is_forbidden_ip(&IpAddr::V4(Ipv4Addr::new(100, 128, 0, 1)))); assert!(!is_forbidden_ip(&IpAddr::V4(Ipv4Addr::new(
241 172, 15, 255, 255
242 )))); assert!(!is_forbidden_ip(&IpAddr::V4(Ipv4Addr::new(172, 32, 0, 1)))); }
245
246 #[test]
247 fn forbidden_ipv6_ranges() {
248 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()))); }
255
256 #[test]
257 fn forbidden_ipv4_mapped_ipv6_ranges() {
258 assert!(is_forbidden_ip(&IpAddr::V6(
259 "::ffff:127.0.0.1".parse().expect("valid mapped loopback")
260 )));
261 assert!(is_forbidden_ip(&IpAddr::V6(
262 "::ffff:169.254.169.254"
263 .parse()
264 .expect("valid mapped link-local")
265 )));
266 assert!(is_forbidden_ip(&IpAddr::V6(
267 "::ffff:10.0.0.1".parse().expect("valid mapped private")
268 )));
269 }
270
271 #[test]
272 fn allowed_ipv6_ranges() {
273 assert!(!is_forbidden_ip(&IpAddr::V6(
274 "2606:4700::1".parse().unwrap()
275 ))); assert!(!is_forbidden_ip(&IpAddr::V6(
277 "2001:4860:4860::8888".parse().unwrap()
278 ))); }
280}