Skip to main content

stealthscraper_rs/state/
model.rs

1//! Per-domain session state model.
2//!
3//! Pure, serializable value types describing what we have learned about a host
4//! across requests (and across process restarts, when a persistent store is
5//! used): the last outcome, which egress proxy was in play, success/failure
6//! tallies, and an optional rate-limit cooldown.
7//!
8//! All time is expressed as Unix seconds and supplied by the caller, so the
9//! model stays free of `SystemTime::now()` and remains trivially testable.
10
11use std::time::Duration;
12
13use serde::{Deserialize, Serialize};
14
15/// The result of an attempt against a host.
16#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
17pub enum Outcome {
18    /// The page loaded cleanly with no bot-protection challenge.
19    Success,
20    /// A challenge was encountered (and possibly cleared) on the page.
21    Challenged,
22    /// The host hard-blocked the request (e.g. Cloudflare error 1020).
23    Blocked,
24    /// The host rate-limited the request (HTTP 429 / error 1015).
25    RateLimited,
26}
27
28/// Accumulated knowledge about a single host.
29#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
30#[serde(deny_unknown_fields)]
31pub struct DomainState {
32    /// The host this state describes (e.g. `example.com`).
33    pub host: String,
34    /// The most recently recorded outcome, if any.
35    pub last_outcome: Option<Outcome>,
36    /// The egress proxy URL used on the most recent recorded attempt.
37    pub last_proxy: Option<String>,
38    /// Count of successful attempts.
39    pub successes: u32,
40    /// Count of failed attempts (challenged/blocked/rate-limited).
41    pub failures: u32,
42    /// Unix timestamp (secs) before which the host should not be hit again.
43    pub cooldown_until: Option<u64>,
44    /// Unix timestamp (secs) of the last update.
45    pub updated_at: u64,
46}
47
48impl DomainState {
49    /// Create a fresh, empty state for `host`.
50    pub fn new(host: impl Into<String>) -> Self {
51        Self {
52            host: host.into(),
53            last_outcome: None,
54            last_proxy: None,
55            successes: 0,
56            failures: 0,
57            cooldown_until: None,
58            updated_at: 0,
59        }
60    }
61
62    /// Returns a new state reflecting `outcome` recorded at `now` (Unix secs).
63    ///
64    /// `proxy` is the egress in play (kept if `None`). Counting:
65    /// - [`Outcome::Success`] and [`Outcome::Challenged`] increment `successes` —
66    ///   in both cases the page was ultimately obtained (a challenge that cleared
67    ///   is a success, not a failure).
68    /// - [`Outcome::Blocked`] and [`Outcome::RateLimited`] increment `failures`.
69    ///
70    /// `cooldown_until` models the rate-limit back-off window: it is set only by
71    /// [`Outcome::RateLimited`] and cleared by every other outcome, so it always
72    /// reflects the most recent outcome. The receiver is left untouched
73    /// (immutable update).
74    pub fn record(
75        &self,
76        outcome: Outcome,
77        proxy: Option<String>,
78        now: u64,
79        rate_limit_cooldown: Duration,
80    ) -> Self {
81        let mut next = self.clone();
82        next.last_outcome = Some(outcome);
83        if proxy.is_some() {
84            next.last_proxy = proxy;
85        }
86        next.updated_at = now;
87        match outcome {
88            Outcome::Success | Outcome::Challenged => {
89                next.successes = next.successes.saturating_add(1);
90                next.cooldown_until = None;
91            }
92            Outcome::RateLimited => {
93                next.failures = next.failures.saturating_add(1);
94                next.cooldown_until = Some(now.saturating_add(rate_limit_cooldown.as_secs()));
95            }
96            Outcome::Blocked => {
97                next.failures = next.failures.saturating_add(1);
98                next.cooldown_until = None;
99            }
100        }
101        next
102    }
103
104    /// Whether the host is still within its cooldown window at `now`.
105    pub fn in_cooldown(&self, now: u64) -> bool {
106        self.cooldown_until.is_some_and(|until| now < until)
107    }
108
109    /// Remaining cooldown at `now`, if any.
110    pub fn cooldown_remaining(&self, now: u64) -> Option<Duration> {
111        self.cooldown_until
112            .and_then(|until| until.checked_sub(now))
113            .filter(|&secs| secs > 0)
114            .map(Duration::from_secs)
115    }
116}
117
118#[cfg(test)]
119mod tests {
120    use super::*;
121
122    #[test]
123    fn record_success_increments_and_clears_cooldown() {
124        let s = DomainState::new("example.com").record(
125            Outcome::RateLimited,
126            None,
127            100,
128            Duration::from_secs(60),
129        );
130        assert!(s.in_cooldown(120));
131
132        let s = s.record(
133            Outcome::Success,
134            Some("http://p:1".into()),
135            200,
136            Duration::ZERO,
137        );
138        assert_eq!(s.successes, 1);
139        assert_eq!(s.last_outcome, Some(Outcome::Success));
140        assert_eq!(s.last_proxy.as_deref(), Some("http://p:1"));
141        assert!(!s.in_cooldown(200));
142        assert_eq!(s.cooldown_until, None);
143    }
144
145    #[test]
146    fn record_rate_limited_sets_cooldown() {
147        let s = DomainState::new("h").record(
148            Outcome::RateLimited,
149            None,
150            1_000,
151            Duration::from_secs(30),
152        );
153        assert_eq!(s.failures, 1);
154        assert!(s.in_cooldown(1_029));
155        assert!(!s.in_cooldown(1_030));
156        assert_eq!(s.cooldown_remaining(1_010), Some(Duration::from_secs(20)));
157        assert_eq!(s.cooldown_remaining(1_030), None);
158    }
159
160    #[test]
161    fn record_keeps_proxy_when_none_passed() {
162        let s = DomainState::new("h")
163            .record(Outcome::Success, Some("http://a".into()), 1, Duration::ZERO)
164            .record(Outcome::Blocked, None, 2, Duration::ZERO);
165        assert_eq!(s.last_proxy.as_deref(), Some("http://a"));
166        assert_eq!(s.failures, 1);
167        assert_eq!(s.last_outcome, Some(Outcome::Blocked));
168    }
169
170    #[test]
171    fn no_cooldown_by_default() {
172        let s = DomainState::new("h");
173        assert!(!s.in_cooldown(0));
174        assert_eq!(s.cooldown_remaining(0), None);
175    }
176
177    #[test]
178    fn challenged_counts_as_success_and_clears_cooldown() {
179        // A challenge that ultimately cleared is a success, not a failure.
180        let s =
181            DomainState::new("h").record(Outcome::RateLimited, None, 100, Duration::from_secs(60));
182        assert!(s.in_cooldown(120));
183
184        let s = s.record(Outcome::Challenged, None, 200, Duration::ZERO);
185        assert_eq!(s.successes, 1);
186        assert_eq!(s.failures, 1); // the earlier rate-limit
187        assert_eq!(s.last_outcome, Some(Outcome::Challenged));
188        assert!(!s.in_cooldown(200));
189    }
190
191    #[test]
192    fn blocked_counts_as_failure_and_clears_stale_cooldown() {
193        let s =
194            DomainState::new("h").record(Outcome::RateLimited, None, 100, Duration::from_secs(60));
195        assert!(s.in_cooldown(120));
196
197        // A subsequent hard block is a different condition; the stale rate-limit
198        // cooldown must not linger.
199        let s = s.record(Outcome::Blocked, None, 130, Duration::ZERO);
200        assert_eq!(s.failures, 2);
201        assert_eq!(s.successes, 0);
202        assert!(!s.in_cooldown(130));
203        assert_eq!(s.cooldown_until, None);
204    }
205}