parse_book_source/
cookie.rs

1//! cookie 库:按**注册域(eTLD+1)**归并,session / persistent 分离。
2//!
3//! - `registrable_domain`:用 `psl` 公共后缀表取注册域(`example.co.uk` / `site.com.cn` 正确)。
4//! - [`CookieJar`]:内存态 cookie 库,`Set-Cookie` 回灌(`enabledCookieJar`)、请求前合并进 `Cookie` 头;
5//!   session cookie(无 `Expires`/`Max-Age`)仅内存、重启失效,persistent 可落盘([`CookieJar::persistent`])。
6//!   `cf_clearance`、headful 登录 cookie、`Set-Cookie` 三路可汇入同一库。
7
8use std::collections::{BTreeMap, HashMap};
9
10/// 由 URL 或裸 host 取**注册域(eTLD+1)**作为 cookie 归并键。
11/// `psl` 正确处理 `example.com` / `example.co.uk` / `site.com.cn`;IP / 单标签 / 未知后缀回退「末两段」。
12/// 大小写归一(host 不区分大小写)。
13pub fn registrable_domain(url: &str) -> String {
14    let host = url
15        .trim_start_matches("https://")
16        .trim_start_matches("http://")
17        .split(['/', '?', '#'])
18        .next()
19        .unwrap_or("")
20        .split(':') // 去端口
21        .next()
22        .unwrap_or("")
23        .to_ascii_lowercase();
24    // IPv4 优先判定:psl 会把纯数字 host 误套默认规则(192.168.1.1 → 1.1),故先短路。
25    let labels: Vec<&str> = host.split('.').filter(|s| !s.is_empty()).collect();
26    let is_ip = !labels.is_empty() && labels.iter().all(|l| l.chars().all(|c| c.is_ascii_digit()));
27    if is_ip {
28        return host;
29    }
30    if let Some(d) = psl::domain_str(&host) {
31        return d.to_string();
32    }
33    if labels.len() >= 2 {
34        labels[labels.len() - 2..].join(".")
35    } else {
36        host
37    }
38}
39
40/// 解析一段 cookie 串(`k=v; k2=v2`)为有序 map:key/value 各自 trim、同名 last-wins、
41/// 无 `=` 的段忽略。merge / 反序列化 / 按 key 查值统一走此函数,避免多处手写解析漂移。
42pub(crate) fn parse_cookie_str(s: &str) -> BTreeMap<String, String> {
43    let mut map = BTreeMap::new();
44    for kv in s.split(';').map(str::trim).filter(|s| !s.is_empty()) {
45        if let Some((k, v)) = kv.split_once('=') {
46            map.insert(k.trim().to_string(), v.trim().to_string());
47        }
48    }
49    map
50}
51
52/// 合并两段 cookie 串(`k=v; k2=v2`)按 key 去重(`second` 同名覆盖 `first`),按字典序输出。
53pub fn merge_cookie_str(first: &str, second: &str) -> String {
54    let mut map = parse_cookie_str(first);
55    map.extend(parse_cookie_str(second));
56    pairs_to_str(&map)
57}
58
59/// `name -> value` map 序列化回 `k=v; k2=v2` 串(字典序)。
60pub(crate) fn pairs_to_str(map: &BTreeMap<String, String>) -> String {
61    map.iter()
62        .map(|(k, v)| format!("{k}={v}"))
63        .collect::<Vec<_>>()
64        .join("; ")
65}
66
67/// 剥除 header 值中的 CR/LF——防 header 注入,并避免「多个 Set-Cookie 以 `\n` 连接的值」
68/// 被回写为出站请求头时让 reqwest 构建失败(审查确认的真问题)。
69pub(crate) fn sanitize_header_value(v: &str) -> String {
70    v.replace(['\r', '\n'], "")
71}
72
73/// 请求的注册域:绝对 URL 取其注册域,相对 URL 归一到书源注册域(`source_domain` 应已是注册域)。
74pub(crate) fn request_registrable_domain(url: &str, source_domain: &str) -> String {
75    if url.starts_with("http://") || url.starts_with("https://") {
76        registrable_domain(url)
77    } else {
78        source_domain.to_string()
79    }
80}
81
82/// 把登录态并入出站请求头——host(`net.*`)与引擎(`apply_auth`)共用的单一真相源:
83///
84/// - **同注册域门控**:仅当请求注册域与书源注册域一致时注入 loginHeader(含其 `Cookie` 字段),
85///   防止「页面内容(toc/next_page/bookUrl 等)诱导的第三方绝对 URL」把 `Authorization`/Cookie
86///   外泄(此防护针对可信书源 + 被挂马页面内容的威胁模型;恶意书源本可经脚本自行外泄,不在此列)。
87///   相对 URL 已归一到书源域,天然放行;登录域与 API 域分属不同注册域的书源会被静默跳过,
88///   如需支持留待书源 schema 显式声明授权域名集合。
89/// - Cookie 合并优先级:已有头 Cookie ← loginHeader Cookie ← `jar_cookie`(调用方按请求注册域取,
90///   本就按域隔离,不受门控),按 key 去重。
91/// - 全部值剥 CR/LF(防 header 注入;亦兜底已落盘的脏数据)。
92pub(crate) fn merge_login_into_headers(
93    login_header: &BTreeMap<String, String>,
94    source_domain: &str,
95    request_domain: &str,
96    jar_cookie: Option<&str>,
97    headers: &mut HashMap<String, String>,
98) {
99    let mut cookie = headers
100        .remove("Cookie")
101        .or_else(|| headers.remove("cookie"));
102    if request_domain == source_domain {
103        for (k, v) in login_header {
104            if k.eq_ignore_ascii_case("cookie") {
105                let v = sanitize_header_value(v);
106                cookie = Some(match cookie {
107                    Some(c) => merge_cookie_str(&c, &v),
108                    None => v,
109                });
110            } else {
111                headers.insert(k.clone(), sanitize_header_value(v));
112            }
113        }
114    }
115    if let Some(jar) = jar_cookie {
116        cookie = Some(match cookie {
117            Some(c) => merge_cookie_str(&c, jar),
118            None => jar.to_string(),
119        });
120    }
121    // 最终再 sanitize 一次:cookie 值本身可能含 `\n`(如脏落盘数据),merge 不会剥除。
122    if let Some(c) = cookie.map(|c| sanitize_header_value(&c))
123        && !c.is_empty()
124    {
125        headers.insert("Cookie".into(), c);
126    }
127}
128
129/// 一条 cookie 值 + 是否持久(有 `Expires`/`Max-Age`)。
130#[derive(Debug, Clone, PartialEq, Eq)]
131struct CookieVal {
132    value: String,
133    persistent: bool,
134}
135
136/// 内存态 cookie 库:`注册域 -> (name -> CookieVal)`。
137#[derive(Debug, Clone, Default)]
138pub struct CookieJar {
139    jar: BTreeMap<String, BTreeMap<String, CookieVal>>,
140}
141
142impl CookieJar {
143    /// 从持久化的 `注册域 -> "k=v; k2=v2"` 映射重建(全部标记为 persistent)。
144    pub fn from_persistent(saved: &BTreeMap<String, String>) -> Self {
145        let mut jar = BTreeMap::new();
146        for (domain, cookie) in saved {
147            let m: BTreeMap<String, CookieVal> = parse_cookie_str(cookie)
148                .into_iter()
149                .map(|(k, v)| {
150                    (
151                        k,
152                        CookieVal {
153                            value: v,
154                            persistent: true,
155                        },
156                    )
157                })
158                .collect();
159            if !m.is_empty() {
160                jar.insert(registrable_domain(domain), m);
161            }
162        }
163        Self { jar }
164    }
165
166    /// 取某域名(自动归一到注册域)的 `Cookie` 头串;空则 `None`。
167    pub fn cookie_header(&self, domain: &str) -> Option<String> {
168        let key = registrable_domain(domain);
169        let m = self.jar.get(&key)?;
170        if m.is_empty() {
171            return None;
172        }
173        let flat: BTreeMap<String, String> = m
174            .iter()
175            .map(|(k, v)| (k.clone(), v.value.clone()))
176            .collect();
177        Some(pairs_to_str(&flat))
178    }
179
180    /// 回灌一条响应的 `Set-Cookie`(可能多条以 `\n` 连接)到某请求域名(归一到注册域)。
181    /// `Max-Age<=0` 视为删除;有 `Expires`/`Max-Age(>0)` 为 persistent,否则 session。
182    pub fn absorb_set_cookie(&mut self, request_domain: &str, set_cookie: &str) {
183        let key = registrable_domain(request_domain);
184        let entry = self.jar.entry(key).or_default();
185        for line in set_cookie
186            .split('\n')
187            .map(str::trim)
188            .filter(|s| !s.is_empty())
189        {
190            let mut parts = line.split(';').map(str::trim);
191            let Some(nv) = parts.next() else { continue };
192            let Some((name, value)) = nv.split_once('=') else {
193                continue;
194            };
195            let (name, value) = (name.trim().to_string(), value.trim().to_string());
196            if name.is_empty() {
197                continue;
198            }
199            let mut persistent = false;
200            let mut deleted = false;
201            for attr in parts {
202                let lower = attr.to_ascii_lowercase();
203                if let Some(ma) = lower.strip_prefix("max-age=") {
204                    match ma.trim().parse::<i64>() {
205                        Ok(n) if n <= 0 => deleted = true,
206                        Ok(_) => persistent = true,
207                        Err(_) => {}
208                    }
209                } else if lower.starts_with("expires=") {
210                    persistent = true;
211                }
212            }
213            if deleted {
214                entry.remove(&name);
215            } else {
216                entry.insert(name, CookieVal { value, persistent });
217            }
218        }
219        if entry.is_empty() {
220            self.jar.remove(&registrable_domain(request_domain));
221        }
222    }
223
224    /// 仅取 persistent cookie 的 `注册域 -> "k=v; k2=v2"` 映射,供落盘(session cookie 不保存)。
225    pub fn persistent(&self) -> BTreeMap<String, String> {
226        let mut out = BTreeMap::new();
227        for (domain, m) in &self.jar {
228            let flat: BTreeMap<String, String> = m
229                .iter()
230                .filter(|(_, v)| v.persistent)
231                .map(|(k, v)| (k.clone(), v.value.clone()))
232                .collect();
233            if !flat.is_empty() {
234                out.insert(domain.clone(), pairs_to_str(&flat));
235            }
236        }
237        out
238    }
239
240    /// 库是否为空(无任何 cookie)。
241    pub fn is_empty(&self) -> bool {
242        self.jar.values().all(BTreeMap::is_empty)
243    }
244}
245
246#[cfg(test)]
247mod tests {
248    use super::*;
249
250    #[test]
251    fn registrable_domain_publicsuffix_and_fallbacks() {
252        assert_eq!(
253            registrable_domain("https://www.fanqienovel.com/x"),
254            "fanqienovel.com"
255        );
256        assert_eq!(registrable_domain("http://api.site.com:8080/p"), "site.com");
257        assert_eq!(registrable_domain("WWW.Site.COM"), "site.com");
258        assert_eq!(
259            registrable_domain("https://www.example.com.cn/p"),
260            "example.com.cn"
261        );
262        assert_eq!(
263            registrable_domain("http://a.b.example.co.uk"),
264            "example.co.uk"
265        );
266        assert_eq!(registrable_domain("http://192.168.1.1:80"), "192.168.1.1");
267        assert_eq!(registrable_domain("localhost"), "localhost");
268        assert_eq!(registrable_domain("http:///path"), "");
269    }
270
271    #[test]
272    fn merge_cookie_str_dedups_second_wins() {
273        assert_eq!(
274            merge_cookie_str("sid=old; theme=dark", "sid=new; lang=zh"),
275            "lang=zh; sid=new; theme=dark"
276        );
277    }
278
279    #[test]
280    fn sanitize_header_value_strips_crlf() {
281        assert_eq!(
282            sanitize_header_value("a=1; Path=/\nb=2; HttpOnly"),
283            "a=1; Path=/b=2; HttpOnly"
284        );
285        assert_eq!(sanitize_header_value("Bearer\r\n token"), "Bearer token");
286    }
287
288    // ── 审查/security:loginHeader 仅注入同注册域请求;Cookie 三方合并;jar cookie 不受门控 ──
289    #[test]
290    fn merge_login_gates_cross_domain_and_merges_cookie() {
291        let mut lh = BTreeMap::new();
292        lh.insert("Authorization".into(), "Bearer T".into());
293        lh.insert("Cookie".into(), "lang=zh".into());
294        // 同注册域:loginHeader 注入,Cookie = 已有头 ← loginHeader ← jar 三方按 key 合并。
295        let mut h = HashMap::new();
296        h.insert("Cookie".into(), "a=1".into());
297        merge_login_into_headers(&lh, "site.com", "site.com", Some("sid=9"), &mut h);
298        assert_eq!(h.get("Authorization").map(String::as_str), Some("Bearer T"));
299        assert_eq!(
300            h.get("Cookie").map(String::as_str),
301            Some("a=1; lang=zh; sid=9")
302        );
303        // 跨注册域:loginHeader(含其 Cookie)整体跳过,jar cookie(按请求域取)照常注入。
304        let mut h2 = HashMap::new();
305        merge_login_into_headers(&lh, "site.com", "evil.com", Some("sid=9"), &mut h2);
306        assert!(!h2.contains_key("Authorization"), "跨域不应注入登录头");
307        assert_eq!(h2.get("Cookie").map(String::as_str), Some("sid=9"));
308    }
309
310    #[test]
311    fn absorb_splits_session_and_persistent() {
312        let mut jar = CookieJar::default();
313        // 子域请求 → 归并到注册域 site.com。
314        jar.absorb_set_cookie(
315            "www.site.com",
316            "sid=abc; Path=/\nremember=1; Max-Age=3600; HttpOnly\ntmp=x; Path=/",
317        );
318        // 请求头含全部(session + persistent)。
319        let header = jar.cookie_header("api.site.com").unwrap();
320        assert!(header.contains("sid=abc"));
321        assert!(header.contains("remember=1"));
322        assert!(header.contains("tmp=x"));
323        // 落盘只留 persistent(remember 有 Max-Age),session 的 sid/tmp 不保存。
324        let persisted = jar.persistent();
325        assert_eq!(
326            persisted.get("site.com").map(String::as_str),
327            Some("remember=1")
328        );
329    }
330
331    #[test]
332    fn absorb_max_age_zero_deletes() {
333        let mut jar = CookieJar::default();
334        jar.absorb_set_cookie("site.com", "sid=abc; Max-Age=3600");
335        assert!(jar.cookie_header("site.com").unwrap().contains("sid=abc"));
336        jar.absorb_set_cookie("site.com", "sid=; Max-Age=0");
337        assert!(jar.cookie_header("site.com").is_none(), "Max-Age=0 应删除");
338    }
339
340    #[test]
341    fn from_persistent_round_trip() {
342        let mut saved = BTreeMap::new();
343        saved.insert("site.com".to_string(), "a=1; b=2".to_string());
344        let jar = CookieJar::from_persistent(&saved);
345        assert_eq!(
346            jar.cookie_header("www.site.com"),
347            Some("a=1; b=2".to_string())
348        );
349        assert_eq!(
350            jar.persistent().get("site.com").map(String::as_str),
351            Some("a=1; b=2")
352        );
353    }
354
355    #[test]
356    fn expires_attribute_marks_persistent() {
357        let mut jar = CookieJar::default();
358        jar.absorb_set_cookie("site.com", "t=1; Expires=Wed, 09 Jun 2027 10:18:14 GMT");
359        assert_eq!(
360            jar.persistent().get("site.com").map(String::as_str),
361            Some("t=1")
362        );
363    }
364}
parse_book_source/cookie.rs

parse_book_source/
cookie.rs
