1use std::collections::HashMap;
2use std::sync::{Arc, Mutex};
3use std::time::SystemTime;
4
5use serde::{Deserialize, Deserializer, Serialize, Serializer};
6use url::Url;
7
8use crate::cookie::parse;
9use crate::cookie::record::Cookie;
10use crate::util::lock;
11
12const MAX_COOKIES_PER_DOMAIN: usize = 180;
13const EVICT_PER_DOMAIN: usize = 30;
14const MAX_COOKIES_GLOBAL: usize = 3300;
15const EVICT_GLOBAL: usize = 300;
16
17#[derive(Clone)]
18pub struct Jar {
19 inner: Arc<Mutex<JarInner>>,
20}
21
22struct JarInner {
23 cookies: HashMap<String, Vec<Cookie>>,
24 total: usize,
25}
26
27impl Jar {
28 pub fn new() -> Self {
29 Self::from_buckets(HashMap::new(), 0)
30 }
31
32 fn from_buckets(cookies: HashMap<String, Vec<Cookie>>, total: usize) -> Self {
33 Self {
34 inner: Arc::new(Mutex::new(JarInner { cookies, total })),
35 }
36 }
37
38 pub fn snapshot(&self) -> Self {
39 let jar = lock(&self.inner);
40 Self::from_buckets(jar.cookies.clone(), jar.total)
41 }
42
43 pub fn extend_from(&self, other: &Jar) {
44 if Arc::ptr_eq(&self.inner, &other.inner) {
45 return;
46 }
47 let incoming: Vec<Cookie> = lock(&other.inner)
48 .cookies
49 .values()
50 .flatten()
51 .cloned()
52 .collect();
53 let mut jar = lock(&self.inner);
54 for cookie in incoming {
55 let domain = cookie.domain.to_lowercase();
56 let entries = jar.cookies.entry(domain.clone()).or_default();
57 let added = match entries.iter().position(|c| c.same_slot(&cookie)) {
58 Some(pos) => {
59 entries[pos] = cookie;
60 false
61 }
62 None => {
63 entries.push(cookie);
64 true
65 }
66 };
67 settle(&mut jar, &domain, added);
68 }
69 }
70
71 pub fn store_set_cookie(&self, set_cookie: &str, url: &Url) {
72 self.store(set_cookie, url);
73 }
74
75 fn store(&self, set_cookie: &str, url: &Url) {
76 let mut cookie = match parse::parse_set_cookie(set_cookie, url) {
77 Some(c) => c,
78 None => return,
79 };
80 if !cookie.secure
81 && url.scheme() != "https"
82 && lock(&self.inner)
83 .cookies
84 .values()
85 .flatten()
86 .any(|existing| cookie.shadows_secure(existing))
87 {
88 return;
89 }
90
91 if cookie.is_expired() {
92 let mut jar = lock(&self.inner);
93 let domain = cookie.domain.to_lowercase();
94 if let Some(entries) = jar.cookies.get_mut(&domain)
95 && let Some(pos) = entries.iter().position(|c| c.same_slot(&cookie))
96 {
97 if entries[pos].secure && url.scheme() != "https" {
98 return;
99 }
100 jar.remove_in(&domain, |c| c.same_slot(&cookie));
101 }
102 return;
103 }
104
105 let mut jar = lock(&self.inner);
106 let domain = cookie.domain.to_lowercase();
107
108 let entries = jar.cookies.entry(domain.clone()).or_default();
109
110 let mut added = false;
111 if let Some(pos) = entries.iter().position(|c| c.same_slot(&cookie)) {
112 if entries[pos].secure && !cookie.secure {
113 return;
114 }
115 cookie.creation_time = entries[pos].creation_time;
116 entries[pos] = cookie;
117 } else {
118 entries.push(cookie);
119 added = true;
120 }
121 settle(&mut jar, &domain, added);
122 }
123
124 pub(crate) fn store_response_cookies(&self, headers: &[&str], url: &Url) {
125 for header in headers {
126 self.store(header, url);
127 }
128 }
129
130 pub fn cookie_header(&self, url: &Url) -> Option<String> {
131 self.cookie_header_for(url, false, true)
132 }
133
134 pub(crate) fn cookie_header_for(
135 &self,
136 url: &Url,
137 cross_site: bool,
138 lax_allowed: bool,
139 ) -> Option<String> {
140 use crate::cookie::record::SameSite;
141 let domain = url.host_str().unwrap_or("");
142 let path = url.path();
143 let is_secure = url.scheme() == "https";
144
145 let mut jar = lock(&self.inner);
146 let now = SystemTime::now();
147
148 let mut matching: Vec<&mut Cookie> = Vec::with_capacity(8);
149 for entries in jar.cookies.values_mut() {
150 for cookie in entries.iter_mut() {
151 if cookie.is_expired() {
152 continue;
153 }
154 if !cookie.matches(domain, path, is_secure) {
155 continue;
156 }
157 if cross_site {
158 match cookie.same_site {
159 SameSite::Strict => continue,
160 SameSite::Lax if !lax_allowed => continue,
161 _ => {}
162 }
163 }
164 cookie.last_access = now;
165 matching.push(cookie);
166 }
167 }
168
169 if matching.is_empty() {
170 return None;
171 }
172
173 matching.sort_by(|a, b| {
174 b.path
175 .len()
176 .cmp(&a.path.len())
177 .then(a.creation_time.cmp(&b.creation_time))
178 });
179
180 let header = matching
181 .iter()
182 .map(|c| format!("{}={}", c.name, c.value))
183 .collect::<Vec<_>>()
184 .join("; ");
185
186 Some(header)
187 }
188
189 pub fn remove_named(&self, name: &str) -> usize {
190 lock(&self.inner).remove_where(|c| c.name == name)
191 }
192
193 pub fn remove(&self, url: &Url, name: &str) -> usize {
194 let host = url.host_str().unwrap_or("");
195 lock(&self.inner).remove_where(|c| c.name == name && c.matches(host, &c.path, true))
196 }
197
198 pub fn all_cookies(&self) -> Vec<Cookie> {
199 let jar = lock(&self.inner);
200 let mut out: Vec<Cookie> = jar.cookies.values().flatten().cloned().collect();
201 out.sort_by(|a, b| a.domain.cmp(&b.domain).then_with(|| a.name.cmp(&b.name)));
202 out
203 }
204
205 pub fn get_cookie(&self, url: &Url, name: &str) -> Option<String> {
206 let domain = url.host_str().unwrap_or("");
207 let path = url.path();
208 let is_secure = url.scheme() == "https";
209
210 let jar = lock(&self.inner);
211 jar.cookies
212 .values()
213 .flatten()
214 .find(|c| c.name == name && c.matches(domain, path, is_secure) && !c.is_expired())
215 .map(|c| c.value.clone())
216 }
217
218 pub fn set_cookie(&self, url: &Url, name: &str, value: &str) {
219 if plain_token(name) && !name.contains('=') && plain_token(value) {
220 self.store(&format!("{name}={value}; Path=/"), url);
221 }
222 }
223
224 pub fn load_cookies(&self, cookie_str: &str, url: &Url) {
225 for pair in cookie_str.split(';') {
226 if let Some((name, value)) = pair.trim().split_once('=') {
227 self.set_cookie(url, name, value);
228 }
229 }
230 }
231
232 pub fn clear(&self) {
233 let mut jar = lock(&self.inner);
234 jar.cookies.clear();
235 jar.total = 0;
236 }
237}
238
239impl Default for Jar {
240 fn default() -> Self {
241 Self::new()
242 }
243}
244
245impl std::fmt::Debug for Jar {
246 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
247 let jar = lock(&self.inner);
248 f.debug_struct("Jar")
249 .field("domains", &jar.cookies.len())
250 .field("total", &jar.total)
251 .finish()
252 }
253}
254
255impl Serialize for Jar {
256 fn serialize<S: Serializer>(&self, ser: S) -> std::result::Result<S::Ok, S::Error> {
257 let jar = lock(&self.inner);
258 let mut flat: Vec<&Cookie> = jar.cookies.values().flatten().collect();
259 flat.sort_by(|a, b| {
260 a.domain
261 .cmp(&b.domain)
262 .then_with(|| a.path.cmp(&b.path))
263 .then_with(|| a.name.cmp(&b.name))
264 .then_with(|| a.creation_time.cmp(&b.creation_time))
265 });
266 flat.serialize(ser)
267 }
268}
269
270impl<'de> Deserialize<'de> for Jar {
271 fn deserialize<D: Deserializer<'de>>(de: D) -> std::result::Result<Self, D::Error> {
272 let cookies: Vec<Cookie> = Vec::deserialize(de)?;
273 let mut buckets: HashMap<String, Vec<Cookie>> = HashMap::new();
274 let mut total = 0;
275 for c in cookies {
276 let key = c.domain.to_lowercase();
277 buckets.entry(key).or_default().push(c);
278 total += 1;
279 }
280 Ok(Self::from_buckets(buckets, total))
281 }
282}
283
284fn plain_token(text: &str) -> bool {
285 !text.chars().any(|c| c == ';' || c.is_control())
286}
287
288impl JarInner {
289 fn remove_where(&mut self, mut doomed: impl FnMut(&Cookie) -> bool) -> usize {
290 let mut removed = 0;
291 for entries in self.cookies.values_mut() {
292 let before = entries.len();
293 entries.retain(|c| !doomed(c));
294 removed += before - entries.len();
295 }
296 self.prune();
297 self.total -= removed;
298 removed
299 }
300
301 fn remove_in(&mut self, domain: &str, mut doomed: impl FnMut(&Cookie) -> bool) -> usize {
302 let Some(entries) = self.cookies.get_mut(domain) else {
303 return 0;
304 };
305 let before = entries.len();
306 entries.retain(|c| !doomed(c));
307 let removed = before - entries.len();
308 if entries.is_empty() {
309 self.cookies.remove(domain);
310 }
311 self.total -= removed;
312 removed
313 }
314
315 fn prune(&mut self) {
316 self.cookies.retain(|_, entries| !entries.is_empty());
317 }
318}
319
320fn settle(jar: &mut JarInner, domain: &str, added: bool) {
321 let mut evicted = 0;
322 if let Some(entries) = jar.cookies.get_mut(domain)
323 && entries.len() > MAX_COOKIES_PER_DOMAIN
324 {
325 evict_lru(entries, EVICT_PER_DOMAIN);
326 evicted = EVICT_PER_DOMAIN;
327 }
328 if added {
329 jar.total += 1;
330 }
331 jar.total -= evicted;
332 if jar.total > MAX_COOKIES_GLOBAL {
333 evict_global(&mut jar.cookies, EVICT_GLOBAL);
334 jar.prune();
335 jar.total = jar.cookies.values().map(|v| v.len()).sum();
336 }
337}
338
339fn evict_lru(cookies: &mut Vec<Cookie>, count: usize) {
340 cookies.sort_by_key(|a| a.last_access);
341 cookies.drain(..count.min(cookies.len()));
342}
343
344fn evict_global(all: &mut HashMap<String, Vec<Cookie>>, count: usize) {
345 let mut all_cookies: Vec<(String, usize, SystemTime)> = Vec::new();
346 for (domain, entries) in all.iter() {
347 for (i, cookie) in entries.iter().enumerate() {
348 all_cookies.push((domain.clone(), i, cookie.last_access));
349 }
350 }
351 all_cookies.sort_by_key(|a| a.2);
352
353 let to_remove = count.min(all_cookies.len());
354 let mut removals: HashMap<String, Vec<usize>> = HashMap::new();
355 for (domain, idx, _) in &all_cookies[..to_remove] {
356 removals.entry(domain.clone()).or_default().push(*idx);
357 }
358 for (domain, mut indices) in removals {
359 indices.sort_unstable_by(|a, b| b.cmp(a));
360 if let Some(entries) = all.get_mut(&domain) {
361 for idx in indices {
362 if idx < entries.len() {
363 entries.remove(idx);
364 }
365 }
366 }
367 }
368}
369
370#[cfg(test)]
371mod tests;