1#[cfg(feature = "browser")]
6pub mod browser;
7pub mod cookie;
8
9use crate::error::FetchError;
10use crate::fetch::cookie::{merge_cookie_str, sanitize_header_value};
11use crate::source::{BookSource, Charset, Method, RateLimit, Retry};
12use async_trait::async_trait;
13use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
14use std::collections::HashMap;
15use std::sync::Mutex;
16use std::time::{Duration, Instant};
17
18struct RateLimiter {
21 interval: Duration,
22 last: Mutex<Option<Instant>>,
23}
24
25impl RateLimiter {
26 fn from_config(rl: &RateLimit) -> Option<Self> {
27 if rl.max_count == 0 || rl.per_ms == 0 {
28 return None;
29 }
30 Some(Self {
31 interval: Duration::from_millis(rl.per_ms / rl.max_count),
32 last: Mutex::new(None),
33 })
34 }
35
36 fn from_rate_str(s: &str) -> Option<Self> {
38 let s = s.trim();
39 if s.is_empty() {
40 return None;
41 }
42 let (max_count, per_ms) = match s.split_once('/') {
43 Some((n, ms)) => (n.trim().parse().ok()?, ms.trim().parse().ok()?),
44 None => (1, s.parse().ok()?),
45 };
46 Self::from_config(&RateLimit { max_count, per_ms })
47 }
48
49 async fn acquire(&self) {
50 let wait = {
51 let mut last = self.last.lock().expect("rate limiter mutex poisoned");
52 let now = Instant::now();
53 let wait = match *last {
54 Some(prev) => self
55 .interval
56 .checked_sub(now.duration_since(prev))
57 .unwrap_or(Duration::ZERO),
58 None => Duration::ZERO,
59 };
60 *last = Some(now + wait);
62 wait
63 }; if !wait.is_zero() {
65 tokio::time::sleep(wait).await;
66 }
67 }
68}
69
70pub fn is_challenge(status: u16, cf_mitigated: Option<&str>, body: &str) -> bool {
77 if cf_mitigated == Some("challenge") {
78 return true;
79 }
80 matches!(status, 403 | 503)
81 && (body.contains("_cf_chl_opt")
82 || body.contains("/cdn-cgi/challenge-platform/")
83 || body.contains("<title>Just a moment"))
84}
85
86#[derive(Debug, Clone, Default)]
88pub struct FetchRequest {
89 pub url: String,
90 pub method: Method,
91 pub body: Option<String>,
92 pub headers: HashMap<String, String>,
93 pub render: bool,
97 pub ready_for: Option<String>,
102 pub intercept_api: Option<String>,
104 pub page: u32,
107 pub page_by: Option<String>,
111}
112
113impl FetchRequest {
114 pub fn get(url: impl Into<String>) -> Self {
116 Self {
117 url: url.into(),
118 ..Default::default()
119 }
120 }
121}
122
123#[derive(Debug, Clone, Default)]
128pub struct FetchResponse {
129 pub body: String,
130 pub status: u16,
131 pub headers: HashMap<String, String>,
132 pub dom_html: Option<String>,
136}
137
138#[async_trait]
140pub trait Fetcher: Send + Sync {
141 async fn fetch(&self, req: FetchRequest) -> Result<String, FetchError>;
143
144 async fn fetch_full(&self, req: FetchRequest) -> Result<FetchResponse, FetchError> {
147 let body = self.fetch(req).await?;
148 Ok(FetchResponse {
149 body,
150 status: 200,
151 headers: HashMap::new(),
152 dom_html: None,
153 })
154 }
155}
156
157pub struct ReqwestFetcher {
159 client: reqwest::Client,
160 base: String,
161 charset: Charset,
162 retry: Option<Retry>,
163 limiter: Option<RateLimiter>,
164 static_cookie: Option<String>,
168}
169
170impl ReqwestFetcher {
171 pub fn new(source: &BookSource) -> Result<Self, FetchError> {
173 let http = &source.http;
174 let mut headers = HeaderMap::new();
175 for (k, v) in &http.headers {
176 let name = HeaderName::from_bytes(k.as_bytes())
177 .map_err(|e| FetchError::Header(e.to_string()))?;
178 let val = HeaderValue::from_str(v).map_err(|e| FetchError::Header(e.to_string()))?;
179 headers.insert(name, val);
180 }
181 let static_cookie = if http.cookies.is_empty() {
184 None
185 } else {
186 let cookie = http
187 .cookies
188 .iter()
189 .map(|(k, v)| format!("{k}={v}"))
190 .collect::<Vec<_>>()
191 .join("; ");
192 let val =
193 HeaderValue::from_str(&cookie).map_err(|e| FetchError::Header(e.to_string()))?;
194 headers.insert(reqwest::header::COOKIE, val);
195 Some(cookie)
196 };
197
198 let mut builder = reqwest::Client::builder()
199 .cookie_store(true)
200 .default_headers(headers);
201 if let Some(ms) = http.timeout {
202 builder = builder.timeout(Duration::from_millis(ms));
203 }
204 let client = builder.build()?;
205
206 Ok(Self {
207 client,
208 base: source.url.trim_end_matches('/').to_string(),
209 charset: http.charset,
210 retry: http.retry.clone(),
211 limiter: http
213 .rate_limit
214 .as_ref()
215 .and_then(RateLimiter::from_config)
216 .or_else(|| RateLimiter::from_rate_str(&source.concurrent_rate)),
217 static_cookie,
218 })
219 }
220
221 async fn send_once(&self, url: &str, req: &FetchRequest) -> Result<FetchResponse, FetchError> {
223 let mut builder = match req.method {
224 Method::Get => self.client.get(url),
225 Method::Post => self.client.post(url),
226 };
227 for (k, v) in &req.headers {
228 builder = builder.header(k, final_header_value(self.static_cookie.as_deref(), k, v));
229 }
230 if let Some(body) = &req.body {
231 builder = builder.body(body.clone());
232 }
233 let resp = builder.send().await?;
234 let status = resp.status();
235 let mut headers = HashMap::new();
237 for (name, value) in resp.headers() {
238 if let Ok(v) = value.to_str() {
239 headers
240 .entry(name.as_str().to_string())
241 .and_modify(|e: &mut String| {
242 e.push('\n');
243 e.push_str(v);
244 })
245 .or_insert_with(|| v.to_string());
246 }
247 }
248 let cf_mitigated = headers.get("cf-mitigated").cloned();
250 let status_err = resp.error_for_status_ref().err();
253 let bytes = resp.bytes().await?;
254 let text = self.decode(&bytes);
255 if is_challenge(status.as_u16(), cf_mitigated.as_deref(), &text) {
256 return Err(FetchError::Challenged(format!(
257 "Cloudflare/反爬挑战 @ {url}"
258 )));
259 }
260 if let Some(e) = status_err {
261 return Err(FetchError::Http(e));
262 }
263 Ok(FetchResponse {
264 body: text,
265 status: status.as_u16(),
266 headers,
267 dom_html: None,
268 })
269 }
270
271 async fn fetch_full_inner(&self, req: FetchRequest) -> Result<FetchResponse, FetchError> {
273 if req.render {
277 return Err(FetchError::Challenged(format!(
278 "此请求需渲染型取页(浏览器辅助),reqwest 不支持 @ {}",
279 req.url
280 )));
281 }
282 if let Some(limiter) = &self.limiter {
284 limiter.acquire().await;
285 }
286 let url = self.resolve(&req.url);
287
288 let max = self.retry.as_ref().map(|r| r.max).unwrap_or(0);
290 let backoff = self.retry.as_ref().map(|r| r.backoff_ms).unwrap_or(0);
291 let mut attempt = 0u32;
292 loop {
293 match self.send_once(&url, &req).await {
294 Ok(resp) => return Ok(resp),
295 Err(e) => {
296 if matches!(e, FetchError::Challenged(_)) || attempt >= max {
298 return Err(e);
299 }
300 attempt += 1;
301 if backoff > 0 {
302 tokio::time::sleep(Duration::from_millis(backoff)).await;
303 }
304 }
305 }
306 }
307 }
308
309 pub(crate) fn resolve(&self, url: &str) -> String {
311 if url.starts_with("http://") || url.starts_with("https://") {
312 url.to_string()
313 } else if let Some(rest) = url.strip_prefix('/') {
314 format!("{}/{}", self.base, rest)
315 } else {
316 format!("{}/{}", self.base, url)
317 }
318 }
319
320 fn decode(&self, bytes: &[u8]) -> String {
322 use encoding_rs::{BIG5, GB18030, GBK, UTF_8};
323 match self.charset {
324 Charset::Utf8 => UTF_8.decode(bytes).0.into_owned(),
325 Charset::Gbk => GBK.decode(bytes).0.into_owned(),
326 Charset::Gb18030 => GB18030.decode(bytes).0.into_owned(),
327 Charset::Big5 => BIG5.decode(bytes).0.into_owned(),
328 Charset::Auto => {
329 let (text, _, had_err) = UTF_8.decode(bytes);
330 if had_err {
331 GBK.decode(bytes).0.into_owned()
332 } else {
333 text.into_owned()
334 }
335 }
336 }
337 }
338}
339
340fn final_header_value(static_cookie: Option<&str>, key: &str, value: &str) -> String {
350 let value = sanitize_header_value(value);
351 match static_cookie {
352 Some(s) if key.eq_ignore_ascii_case("cookie") => merge_cookie_str(s, &value),
353 _ => value,
354 }
355}
356
357#[async_trait]
358impl Fetcher for ReqwestFetcher {
359 async fn fetch(&self, req: FetchRequest) -> Result<String, FetchError> {
360 self.fetch_full_inner(req).await.map(|r| r.body)
361 }
362
363 async fn fetch_full(&self, req: FetchRequest) -> Result<FetchResponse, FetchError> {
364 self.fetch_full_inner(req).await
365 }
366}
367
368#[cfg(test)]
369mod tests {
370 use super::{final_header_value, is_challenge};
371
372 const CHALLENGE_HTML: &str = r#"<html><head><title>Just a moment...</title></head>
374 <body><script>window._cf_chl_opt={cType:'managed'};
375 a.src='/cdn-cgi/challenge-platform/h/g/orchestrate/chl_page/v1';</script></body></html>"#;
376
377 const NORMAL_HTML: &str =
378 r#"<html><head><title>蛊真人 搜索结果</title></head><body>正文</body></html>"#;
379
380 #[test]
381 fn cf_mitigated_header_is_challenge() {
382 assert!(is_challenge(200, Some("challenge"), NORMAL_HTML));
384 }
385
386 #[test]
387 fn challenge_body_with_403_is_challenge() {
388 assert!(is_challenge(403, None, CHALLENGE_HTML));
389 assert!(is_challenge(503, None, CHALLENGE_HTML));
390 }
391
392 #[test]
393 fn normal_200_page_is_not_challenge() {
394 assert!(!is_challenge(200, None, NORMAL_HTML));
395 assert!(!is_challenge(200, None, CHALLENGE_HTML));
397 }
398
399 #[test]
400 fn challenge_markers_without_bad_status_not_challenge() {
401 assert!(!is_challenge(403, None, NORMAL_HTML));
403 }
404
405 #[test]
407 fn final_header_value_merges_static_cookie_and_strips_crlf() {
408 assert_eq!(
410 final_header_value(Some("device=d1; sid=old"), "Cookie", "sid=new"),
411 "device=d1; sid=new"
412 );
413 assert_eq!(
414 final_header_value(Some("device=d1"), "cookie", "sid=1"),
415 "device=d1; sid=1"
416 );
417 assert_eq!(
419 final_header_value(Some("device=d1"), "Authorization", "Bearer\r\nT"),
420 "BearerT"
421 );
422 assert_eq!(final_header_value(None, "Cookie", "a=1\nb=2"), "a=1b=2");
424 assert_eq!(final_header_value(None, "X-Test", "42"), "42");
425 }
426}