Skip to main content

mocra_core/common/model/
cookies.rs

1use crate::cacheable::CacheAble;
2use crate::errors::CookieError;
3use async_trait::async_trait;
4use chrono::{TimeZone, Utc};
5use reqwest::cookie::Jar;
6use reqwest::header::HeaderValue;
7use reqwest_cookie_store::CookieStore;
8use serde::{Deserialize, Deserializer, Serialize};
9use std::fmt;
10use url::Url;
11
12/// Cookie validity is managed externally.
13/// This type only performs serialization/deserialization and does not validate cookie validity.
14#[derive(Serialize, Clone)]
15pub struct CookieItem {
16    pub name: String,
17    pub value: String,
18    pub domain: String,
19    pub path: String,
20    // Unix timestamp seconds; accept i64/f64 on input, store as u64 seconds (ms inputs will be converted to seconds)
21    pub expires: Option<u64>,
22    // Seconds; accept numeric forms, store as u64 seconds
23    pub max_age: Option<u64>,
24    pub secure: bool,
25    #[serde(rename = "httpOnly")]
26    pub http_only: Option<bool>,
27}
28
29impl fmt::Debug for CookieItem {
30    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
31        f.debug_struct("CookieItem")
32            .field("name", &self.name)
33            .field("value", &"***REDACTED***")
34            .field("domain", &self.domain)
35            .field("path", &self.path)
36            .field("expires", &self.expires)
37            .field("max_age", &self.max_age)
38            .field("secure", &self.secure)
39            .field("http_only", &self.http_only)
40            .finish()
41    }
42}
43
44impl<'de> Deserialize<'de> for CookieItem {
45    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
46    where
47        D: Deserializer<'de>,
48    {
49        fn parse_expiration_value(v: &serde_json::Value) -> Option<f64> {
50            match v {
51                serde_json::Value::Null => None,
52                serde_json::Value::Number(n) => n.as_f64().map(|u| {
53                    // Handle millisecond timestamps (>= 1e12 is treated as milliseconds).
54                    if u.abs() >= 1_000_000_000_000.0 {
55                        u / 1000.0
56                    } else {
57                        u
58                    }
59                }),
60                serde_json::Value::String(s) => parse_expiration_date_str(s),
61                _ => None,
62            }
63        }
64
65        // Parse "YYYY/M/D HH:MM:SS" or "YYYY-MM-DD HH:MM:SS" into seconds.
66        fn parse_expiration_date_str(s: &str) -> Option<f64> {
67            let (date_part, time_part) = s.split_once(' ')?;
68
69            let mut hms = time_part.split(':');
70            let hh = hms.next()?.parse::<u32>().ok()?;
71            let mm = hms.next()?.parse::<u32>().ok()?;
72            let ss = hms.next()?.parse::<u32>().ok()?;
73
74            let (y, mo, d) = if date_part.contains('/') {
75                let mut it = date_part.split('/');
76                (
77                    it.next()?.parse::<i32>().ok()?,
78                    it.next()?.parse::<u32>().ok()?,
79                    it.next()?.parse::<u32>().ok()?,
80                )
81            } else if date_part.contains('-') {
82                let mut it = date_part.split('-');
83                (
84                    it.next()?.parse::<i32>().ok()?,
85                    it.next()?.parse::<u32>().ok()?,
86                    it.next()?.parse::<u32>().ok()?,
87                )
88            } else {
89                return None;
90            };
91
92            Utc.with_ymd_and_hms(y, mo, d, hh, mm, ss)
93                .single()
94                .map(|dt| dt.timestamp().max(0) as f64)
95        }
96
97        fn parse_seconds_value(v: &serde_json::Value) -> Option<f64> {
98            match v {
99                serde_json::Value::Null => None,
100                serde_json::Value::Number(n) => n.as_f64(),
101                serde_json::Value::String(s) => s.parse::<f64>().ok(),
102                _ => None,
103            }
104        }
105
106        struct CookieItemVisitor;
107
108        impl<'de> serde::de::Visitor<'de> for CookieItemVisitor {
109            type Value = CookieItem;
110
111            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
112                formatter.write_str("CookieItem as map or seq")
113            }
114
115            fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
116            where
117                A: serde::de::SeqAccess<'de>,
118            {
119                let name: String = seq
120                    .next_element()?
121                    .ok_or_else(|| serde::de::Error::invalid_length(0, &self))?;
122                let value: String = seq
123                    .next_element()?
124                    .ok_or_else(|| serde::de::Error::invalid_length(1, &self))?;
125                let domain: String = seq.next_element()?.unwrap_or_default();
126                let path: String = seq.next_element()?.unwrap_or_else(|| "/".to_string());
127                let expires: Option<u64> = seq.next_element()?.unwrap_or(None);
128                let max_age: Option<u64> = seq.next_element()?.unwrap_or(None);
129                let secure: bool = seq.next_element()?.unwrap_or(false);
130                let http_only: Option<bool> = seq.next_element()?.unwrap_or(None);
131
132                Ok(CookieItem {
133                    name,
134                    value,
135                    domain,
136                    path,
137                    expires,
138                    max_age,
139                    secure,
140                    http_only,
141                })
142            }
143
144            fn visit_map<A>(self, mut map: A) -> Result<Self::Value, A::Error>
145            where
146                A: serde::de::MapAccess<'de>,
147            {
148                let mut name: Option<String> = None;
149                let mut value: Option<String> = None;
150                let mut domain: Option<String> = None;
151                let mut path: Option<String> = None;
152                let mut secure: Option<bool> = None;
153                let mut http_only: Option<bool> = None;
154                let mut expires: Option<serde_json::Value> = None;
155                let mut expiration_date: Option<serde_json::Value> = None;
156                let mut max_age: Option<serde_json::Value> = None;
157                let mut max_age_dash: Option<serde_json::Value> = None;
158                let mut max_age_camel: Option<serde_json::Value> = None;
159
160                while let Some(key) = map.next_key::<String>()? {
161                    match key.as_str() {
162                        "name" => name = Some(map.next_value()?),
163                        "value" => value = Some(map.next_value()?),
164                        "domain" => domain = Some(map.next_value()?),
165                        "path" => path = Some(map.next_value()?),
166                        "secure" => secure = Some(map.next_value()?),
167                        "httpOnly" => http_only = Some(map.next_value()?),
168                        "expires" => expires = Some(map.next_value()?),
169                        "expirationDate" => expiration_date = Some(map.next_value()?),
170                        "max_age" => max_age = Some(map.next_value()?),
171                        "max-age" => max_age_dash = Some(map.next_value()?),
172                        "maxAge" => max_age_camel = Some(map.next_value()?),
173                        _ => {
174                            let _: serde::de::IgnoredAny = map.next_value()?;
175                        }
176                    }
177                }
178
179                let name = name.ok_or_else(|| serde::de::Error::missing_field("name"))?;
180                let value = value.ok_or_else(|| serde::de::Error::missing_field("value"))?;
181                let domain = domain.unwrap_or_default();
182                let path = path.unwrap_or_else(|| "/".to_string());
183                let secure = secure.unwrap_or(false);
184                let http_only = http_only;
185
186                let expires = match (expires.as_ref(), expiration_date.as_ref()) {
187                    (Some(v), _) => parse_expiration_value(v),
188                    (None, Some(v)) => parse_expiration_value(v),
189                    _ => None,
190                }
191                .map(|secs| secs.floor().max(0.0) as u64);
192
193                let max_age = max_age
194                    .as_ref()
195                    .and_then(parse_seconds_value)
196                    .or_else(|| max_age_dash.as_ref().and_then(parse_seconds_value))
197                    .or_else(|| max_age_camel.as_ref().and_then(parse_seconds_value))
198                    .or_else(|| {
199                        expires.map(|e| {
200                            let now = Utc::now().timestamp();
201                            let diff = (e as i64) - now;
202                            diff.max(0) as f64
203                        })
204                    })
205                    .map(|secs| secs.floor().max(0.0) as u64);
206
207                Ok(CookieItem {
208                    name,
209                    value,
210                    domain,
211                    path,
212                    expires,
213                    max_age,
214                    secure,
215                    http_only,
216                })
217            }
218        }
219
220        deserializer.deserialize_any(CookieItemVisitor)
221    }
222}
223
224#[derive(Default, Serialize, Deserialize, Debug, Clone)]
225pub struct Cookies {
226    pub cookies: Vec<CookieItem>,
227}
228
229impl Cookies {
230    pub fn add(&mut self, name: impl AsRef<str>, value: impl AsRef<str>, domain: impl AsRef<str>) {
231        self.cookies.push(CookieItem {
232            name: name.as_ref().into(),
233            value: value.as_ref().into(),
234            domain: domain.as_ref().into(),
235            path: "".to_string(),
236            expires: None,
237            max_age: None,
238            secure: false,
239            http_only: None,
240        })
241    }
242    pub fn merge(&mut self, other: &Cookies) {
243        for cookie in &other.cookies {
244            if !self
245                .cookies
246                .iter()
247                .any(|c| c.name == cookie.name && c.domain == cookie.domain)
248            {
249                self.cookies.push(cookie.clone());
250            }
251        }
252    }
253    pub fn merge_if_absent(&mut self, other: &Cookies) {
254        self.merge(other);
255    }
256    pub fn merge_cookie_store(&mut self, other: &CookieStore) {
257        let other_cookies: Cookies = other.clone().into();
258        self.merge(&other_cookies)
259    }
260    pub fn is_empty(&self) -> bool {
261        self.cookies.is_empty()
262    }
263    pub fn contains(&self, name: impl AsRef<str>, domain: impl AsRef<str>) -> bool {
264        self.cookies
265            .iter()
266            .any(|c| c.name == name.as_ref() && c.domain == domain.as_ref())
267    }
268    pub fn string(&self) -> String {
269        self.cookies
270            .iter()
271            .map(|c| format!("{}={}", c.name, c.value))
272            .collect::<Vec<String>>()
273            .join(";")
274    }
275    pub fn str_by_domain(&self, domain: &[impl AsRef<str>]) -> String {
276        self.cookies
277            .iter()
278            .filter(|c| domain.iter().any(|d| d.as_ref() == c.domain))
279            .map(|c| format!("{}={}", c.name, c.value))
280            .collect::<Vec<String>>()
281            .join(";")
282    }
283
284    pub fn cookie_header_for_url(&self, url: &Url) -> Option<String> {
285        if self.cookies.is_empty() {
286            return None;
287        }
288
289        if let Ok(store) = CookieStore::try_from(self) {
290            let pairs = store
291                .get_request_values(url)
292                .map(|(name, value)| format!("{}={}", name, value))
293                .collect::<Vec<_>>();
294            if !pairs.is_empty() {
295                return Some(pairs.join("; "));
296            }
297        }
298
299        let host = url.host_str()?.to_ascii_lowercase();
300        let pairs = self
301            .cookies
302            .iter()
303            .filter_map(|cookie| {
304                let domain = cookie
305                    .domain
306                    .trim()
307                    .trim_start_matches('.')
308                    .to_ascii_lowercase();
309                if domain.is_empty() {
310                    return None;
311                }
312                if host == domain || host.ends_with(&format!(".{domain}")) {
313                    Some(format!("{}={}", cookie.name, cookie.value))
314                } else {
315                    None
316                }
317            })
318            .collect::<Vec<_>>();
319
320        if pairs.is_empty() {
321            None
322        } else {
323            Some(pairs.join("; "))
324        }
325    }
326}
327impl From<Cookies> for Jar {
328    fn from(value: Cookies) -> Self {
329        let jar = Jar::default();
330
331        for cookie_item in value.cookies {
332            // Build cookie string.
333            let mut cookie_str = format!("{}={}", cookie_item.name, cookie_item.value);
334
335            // Add domain.
336            if !cookie_item.domain.is_empty() {
337                cookie_str.push_str(&format!("; Domain={}", cookie_item.domain));
338            }
339
340            // Add path.
341            if !cookie_item.path.is_empty() {
342                cookie_str.push_str(&format!("; Path={}", cookie_item.path));
343            }
344
345            // Add expiration (GMT, per cookie spec).
346            if let Some(expires) = cookie_item.expires
347                && let Some(dt) = Utc.timestamp_opt(expires as i64, 0).single()
348            {
349                cookie_str.push_str(&format!(
350                    "; Expires={}",
351                    dt.format("%a, %d %b %Y %H:%M:%S GMT")
352                ));
353            }
354
355            // Add Max-Age (seconds).
356            if let Some(max_age) = cookie_item.max_age {
357                cookie_str.push_str(&format!("; Max-Age={max_age}"));
358            }
359
360            // Add Secure flag.
361            if cookie_item.secure {
362                cookie_str.push_str("; Secure");
363            }
364
365            // Add HttpOnly flag.
366            if cookie_item.http_only == Some(true) {
367                cookie_str.push_str("; HttpOnly");
368            }
369
370            // Build URL using domain.
371            let url_str = if cookie_item.secure {
372                format!("https://{}", cookie_item.domain)
373            } else {
374                format!("http://{}", cookie_item.domain)
375            };
376
377            if let Ok(url) = url_str.parse::<Url>() {
378                jar.add_cookie_str(&cookie_str, &url);
379            }
380        }
381
382        jar
383    }
384}
385impl From<CookieItem> for HeaderValue {
386    fn from(value: CookieItem) -> Self {
387        let mut cookie_str = format!("{}={}", value.name, value.value);
388        if !value.domain.is_empty() {
389            cookie_str.push_str(&format!("; Domain={}", value.domain));
390        }
391        if !value.path.is_empty() {
392            cookie_str.push_str(&format!("; Path={}", value.path));
393        }
394        if let Some(expires) = value.expires
395            && let Some(dt) = Utc.timestamp_opt(expires as i64, 0).single()
396        {
397            cookie_str.push_str(&format!(
398                "; Expires={}",
399                dt.format("%a, %d %b %Y %H:%M:%S GMT")
400            ));
401        }
402        if let Some(max_age) = value.max_age {
403            cookie_str.push_str(&format!("; Max-Age={max_age}"));
404        }
405        if value.secure {
406            cookie_str.push_str("; Secure");
407        }
408        if value.http_only == Some(true) {
409            cookie_str.push_str("; HttpOnly");
410        }
411        HeaderValue::from_str(&cookie_str).unwrap_or_else(|_| HeaderValue::from_static(""))
412    }
413}
414impl From<CookieStore> for Cookies {
415    fn from(value: CookieStore) -> Self {
416        let mut cookies = Vec::new();
417
418        // Iterate through all cookies in `CookieStore`.
419        for cookie in value.iter_any() {
420            let cookie_item = CookieItem {
421                name: cookie.name().to_string(),
422                value: cookie.value().to_string(),
423                domain: cookie.domain().map(|d| d.to_string()).unwrap_or_default(),
424                path: cookie.path().map(|p| p.to_string()).unwrap_or_default(),
425                // Normalize expiration to seconds (Unix timestamp, `u64`).
426                expires: cookie.expires().and_then(|exp| {
427                    exp.datetime().map(|x| {
428                        let ts = x.unix_timestamp();
429                        if ts < 0 { 0 } else { ts as u64 }
430                    })
431                }),
432                // Max-Age (seconds), or `None` when unavailable.
433                max_age: cookie.max_age().map(|duration| {
434                    // `cookie_store` duration usually comes from `time` crate.
435                    (duration.whole_seconds()).max(0) as u64
436                }),
437                secure: cookie.secure().unwrap_or(false),
438                http_only: cookie.http_only(),
439            };
440            cookies.push(cookie_item);
441        }
442
443        Cookies { cookies }
444    }
445}
446impl TryFrom<&Cookies> for CookieStore {
447    type Error = CookieError;
448
449    fn try_from(value: &Cookies) -> Result<Self, Self::Error> {
450        let mut store = CookieStore::default();
451
452        for cookie_item in &value.cookies {
453            // Build cookie string.
454            let mut cookie_str = format!("{}={}", cookie_item.name, cookie_item.value);
455
456            // If domain is provided, create Domain Cookie; leading dot is stripped per normalization.
457            let normalized_domain = cookie_item.domain.trim().trim_start_matches('.');
458            if !normalized_domain.is_empty() {
459                cookie_str.push_str(&format!("; Domain={}", normalized_domain));
460            }
461
462            // Path (default `/` for broader matching).
463            if !cookie_item.path.is_empty() {
464                cookie_str.push_str(&format!("; Path={}", cookie_item.path));
465            } else {
466                cookie_str.push_str("; Path=/");
467            }
468
469            // Expires (GMT format).
470            // if let Some(expires) = &cookie_item.expires {
471            //     if let Some(dt) = Utc.timestamp_opt(*expires as i64, 0).single() {
472            //         cookie_str.push_str(&format!(
473            //             "; Expires={}",
474            //             dt.format("%a, %d %b %Y %H:%M:%S GMT")
475            //         ));
476            //     }
477            // }
478            // Max-Age (seconds).
479            // if let Some(max_age) = &cookie_item.max_age {
480            //     cookie_str.push_str(&format!("; Max-Age={}", max_age));
481            // }
482
483            // Flags.
484            if cookie_item.secure {
485                cookie_str.push_str("; Secure");
486            }
487            if cookie_item.http_only == Some(true) {
488                cookie_str.push_str("; HttpOnly");
489            }
490
491            // Build URL (trim leading dot to avoid invalid hosts like `.example.com`).
492            let host = normalized_domain;
493            if host.is_empty() {
494                // Without domain we cannot infer host-only owner; skip here.
495                // Use `into_store_for_base_url` when binding to a specific host is required.
496                continue;
497            }
498            let url_str = if cookie_item.secure {
499                format!("https://{}", host)
500            } else {
501                format!("http://{}", host)
502            };
503
504            if let Ok(url) = url_str.parse::<Url>() {
505                store
506                    .parse(&cookie_str, &url)
507                    .map_err(|e| CookieError::ParseError(e.into()))?;
508            }
509        }
510
511        Ok(store)
512    }
513}
514
515impl From<Vec<CookieItem>> for Cookies {
516    fn from(value: Vec<CookieItem>) -> Self {
517        Cookies { cookies: value }
518    }
519}
520#[async_trait]
521impl CacheAble for Cookies {
522    fn field() -> impl AsRef<str> {
523        "cookies".to_string()
524    }
525}