Skip to main content

antibot_rs/
session_cache.rs

1//! Per-domain session cache so previously-solved cookies can be reused
2//! without round-tripping through the (slow) solver.
3
4use crate::cookie::Cookie;
5use dashmap::DashMap;
6use std::sync::Arc;
7use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
8
9#[derive(Debug, Clone)]
10pub struct CachedSession {
11    pub cookies: Vec<Cookie>,
12    pub user_agent: String,
13    pub solved_at: Instant,
14    pub solved_at_system: SystemTime,
15    pub expires_at: Option<Instant>,
16}
17
18impl CachedSession {
19    pub fn age(&self) -> Duration {
20        self.solved_at.elapsed()
21    }
22}
23
24#[derive(Debug, Clone)]
25pub struct SessionCacheConfig {
26    /// TTL applied when the underlying cookies have no usable expiry.
27    pub default_ttl: Duration,
28    /// Hard cap on cached domains. Soft eviction via random sampling once hit.
29    pub max_entries: usize,
30    /// If `true`, derive `expires_at` from the soonest cookie expiry.
31    pub respect_cookie_expiry: bool,
32}
33
34impl Default for SessionCacheConfig {
35    fn default() -> Self {
36        Self {
37            default_ttl: Duration::from_secs(30 * 60),
38            max_entries: 1000,
39            respect_cookie_expiry: true,
40        }
41    }
42}
43
44#[derive(Clone)]
45pub(crate) struct SessionCache {
46    entries: Arc<DashMap<String, CachedSession>>,
47    config: SessionCacheConfig,
48}
49
50impl SessionCache {
51    pub fn new(config: SessionCacheConfig) -> Self {
52        Self {
53            entries: Arc::new(DashMap::new()),
54            config,
55        }
56    }
57
58    pub fn get(&self, domain: &str) -> Option<CachedSession> {
59        let entry = self.entries.get(domain)?;
60        if let Some(expires_at) = entry.expires_at {
61            if Instant::now() >= expires_at {
62                let key = entry.key().clone();
63                drop(entry);
64                self.entries.remove(&key);
65                return None;
66            }
67        }
68        Some(entry.clone())
69    }
70
71    pub fn insert(&self, domain: String, cookies: Vec<Cookie>, user_agent: String) {
72        let expires_at = self.compute_expiry(&cookies);
73        let now = Instant::now();
74        let session = CachedSession {
75            cookies,
76            user_agent,
77            solved_at: now,
78            solved_at_system: SystemTime::now(),
79            expires_at,
80        };
81        self.entries.insert(domain, session);
82        self.evict_if_needed();
83    }
84
85    pub fn invalidate(&self, domain: &str) {
86        self.entries.remove(domain);
87    }
88
89    pub fn clear(&self) {
90        self.entries.clear();
91    }
92
93    pub fn len(&self) -> usize {
94        self.entries.len()
95    }
96
97    fn compute_expiry(&self, cookies: &[Cookie]) -> Option<Instant> {
98        let default_deadline = Instant::now() + self.config.default_ttl;
99
100        if !self.config.respect_cookie_expiry {
101            return Some(default_deadline);
102        }
103
104        let now_unix = SystemTime::now()
105            .duration_since(UNIX_EPOCH)
106            .map(|d| d.as_secs_f64())
107            .unwrap_or(0.0);
108
109        let earliest_cookie = cookies
110            .iter()
111            .filter_map(|c| c.expires)
112            .filter(|&exp| exp > now_unix)
113            .fold(None, |acc: Option<f64>, exp| {
114                Some(acc.map_or(exp, |a| a.min(exp)))
115            });
116
117        match earliest_cookie {
118            Some(exp) => {
119                let secs_remaining = (exp - now_unix).max(0.0);
120                let from_cookie =
121                    Instant::now() + Duration::from_secs_f64(secs_remaining);
122                Some(from_cookie.min(default_deadline))
123            }
124            None => Some(default_deadline),
125        }
126    }
127
128    fn evict_if_needed(&self) {
129        if self.entries.len() <= self.config.max_entries {
130            return;
131        }
132
133        // Cheap eviction: take the first key the iterator yields and drop it.
134        // DashMap iteration order is shard-dependent, which is good enough for
135        // bounding memory without doing a real LRU.
136        if let Some(victim) = self.entries.iter().next().map(|e| e.key().clone()) {
137            self.entries.remove(&victim);
138        }
139    }
140}
141
142/// Extract the registrable domain from a URL. Returns `None` if `url` can't be
143/// parsed or has no host.
144pub(crate) fn extract_domain(url: &str) -> Option<String> {
145    let parsed = url::Url::parse(url).ok()?;
146    parsed.host_str().map(|s| s.to_ascii_lowercase())
147}