Skip to main content

spire_core/frontier/facade/
traps.rs

1use std::collections::{HashMap, HashSet};
2use std::future::Future;
3use std::sync::Mutex;
4use std::sync::atomic::{AtomicUsize, Ordering};
5
6use ::http::Uri;
7
8use crate::frontier::seed::{Fingerprint, Seed};
9use crate::frontier::{Ack, Admission, Frontier, NextSeed, Report};
10use crate::sync::lock;
11use crate::{Result, TRACING_TARGET};
12
13/// URL budgets for [`AvoidTraps`]. These are heuristics, not proof of a trap.
14///
15/// Budgets apply independently to each origin and first nonempty path
16/// segment, e.g. `https://example.com/archive/`. Query variants count as
17/// separate URLs. State is in memory and starts fresh on resume.
18#[derive(Debug, Clone, Copy)]
19pub struct TrapPolicy {
20    urls: usize,
21    fanout: usize,
22    path_depth: usize,
23    repetitions: usize,
24    random_urls: usize,
25}
26
27impl Default for TrapPolicy {
28    fn default() -> Self {
29        Self {
30            urls: 4096,
31            fanout: 256,
32            path_depth: 32,
33            repetitions: 4,
34            random_urls: 8,
35        }
36    }
37}
38
39impl TrapPolicy {
40    /// Sets the number of distinct request fingerprints offered per prefix,
41    /// 4096 by default. Includes candidates refused by inner facades.
42    #[must_use]
43    pub const fn with_max_urls(mut self, max: usize) -> Self {
44        self.urls = max;
45        self
46    }
47
48    /// Sets the number of distinct URLs offered under one directory within
49    /// a prefix, 256 by default. This estimates fan-out from URL structure,
50    /// not from a page's DOM. Query variants share their path's directory.
51    #[must_use]
52    pub const fn with_max_fanout(mut self, max: usize) -> Self {
53        self.fanout = max;
54        self
55    }
56
57    /// Sets the maximum number of nonempty URL path segments, 32 by default.
58    /// This is independent of the seed's link-following depth.
59    #[must_use]
60    pub const fn with_max_path_depth(mut self, max: usize) -> Self {
61        self.path_depth = max;
62        self
63    }
64
65    /// Sets how many times one segment can repeat in a URL, 4 by default.
66    #[must_use]
67    pub const fn with_max_repetitions(mut self, max: usize) -> Self {
68        self.repetitions = max;
69        self
70    }
71
72    /// Sets the number of URLs containing opaque-looking path segments,
73    /// 8 by default. A segment is suspect when it is at least 16 ASCII
74    /// alphanumeric characters, mixes letters and digits, and contains at
75    /// least 8 distinct characters. Legitimate identifiers can match;
76    /// use `usize::MAX` to disable this budget for such sites.
77    #[must_use]
78    pub const fn with_max_random_urls(mut self, max: usize) -> Self {
79        self.random_urls = max;
80        self
81    }
82}
83
84/// A facade that bounds expanding URL branches, see
85/// [`FrontierExt::avoid_traps`](crate::frontier::FrontierExt::avoid_traps).
86///
87/// When a budget is exceeded, the whole origin/path prefix is blocked for
88/// this facade's lifetime. Seeds already queued there are skipped on pop
89/// and reported inwards as [`Report::Skipped`]. In-flight work can finish.
90/// Other prefixes and origins continue. Refusals are logged and counted
91/// by [`refused`](Self::refused), separately from crawler request totals.
92///
93/// This does not inspect page content, execute CSS, or replace robots.txt
94/// and deduplication. Keep this facade outside retry/revisit facades so it
95/// also checks the seeds they hand out. Snapshots contain pending seeds,
96/// not the heuristic state; persistent storage keeps blocked seeds until
97/// they are popped and skipped.
98#[derive(Debug)]
99pub struct AvoidTraps<F> {
100    inner: F,
101    policy: TrapPolicy,
102    prefixes: Mutex<HashMap<String, Prefix>>,
103    refused: AtomicUsize,
104}
105
106#[derive(Debug, Default)]
107struct Prefix {
108    seen: HashSet<Fingerprint>,
109    directories: HashMap<String, usize>,
110    random_urls: usize,
111    blocked: Option<Admission>,
112}
113
114impl<F> AvoidTraps<F> {
115    /// Wraps `inner` with the URL budgets in `policy`.
116    pub fn new(inner: F, policy: TrapPolicy) -> Self {
117        Self {
118            inner,
119            policy,
120            prefixes: Mutex::new(HashMap::new()),
121            refused: AtomicUsize::new(0),
122        }
123    }
124
125    /// Returns the number of refused pushes and skipped queued seeds.
126    #[must_use]
127    pub fn refused(&self) -> usize {
128        self.refused.load(Ordering::Relaxed)
129    }
130
131    fn refusal(&self, seed: &Seed) -> Option<Admission> {
132        let key = prefix(seed.uri());
133        let mut prefixes = lock(&self.prefixes);
134        let state = prefixes.entry(key.clone()).or_default();
135        if state.blocked.is_some() {
136            return state.blocked;
137        }
138        let fingerprint = seed.fingerprint();
139        if state.seen.contains(&fingerprint) {
140            return None;
141        }
142        let path = seed.uri().path();
143        let segments: Vec<_> = path.split('/').filter(|part| !part.is_empty()).collect();
144        let directory = path.rsplit_once('/').map_or("", |(parent, _)| parent);
145        let fanout = state.directories.get(directory).copied().unwrap_or(0);
146        let random = segments.iter().any(|segment| opaque(segment));
147        let reason = if segments.len() > self.policy.path_depth {
148            Some((Admission::PrefixBudgetExceeded, "path depth"))
149        } else if segments.iter().any(|segment| {
150            segments.iter().filter(|other| *other == segment).count() > self.policy.repetitions
151        }) {
152            Some((Admission::SuspectedTrap, "repeated path segments"))
153        } else if state.seen.len() >= self.policy.urls {
154            Some((Admission::PrefixBudgetExceeded, "URLs per prefix"))
155        } else if fanout >= self.policy.fanout {
156            Some((Admission::PrefixBudgetExceeded, "directory fan-out"))
157        } else if random && state.random_urls >= self.policy.random_urls {
158            Some((Admission::SuspectedTrap, "opaque path segments"))
159        } else {
160            None
161        };
162        if let Some((admission, reason)) = reason {
163            state.blocked = Some(admission);
164            // Keep only the blocked marker once this branch is closed.
165            state.seen.clear();
166            state.directories.clear();
167            tracing::warn!(target: TRACING_TARGET, prefix = %key, reason, "blocked URL prefix");
168            return Some(admission);
169        }
170        state.seen.insert(fingerprint);
171        *state.directories.entry(directory.to_owned()).or_default() += 1;
172        state.random_urls += usize::from(random);
173        None
174    }
175
176    fn is_blocked(&self, seed: &Seed) -> bool {
177        lock(&self.prefixes)
178            .get(&prefix(seed.uri()))
179            .is_some_and(|state| state.blocked.is_some())
180    }
181
182    fn count_refusal(&self, seed: &Seed) {
183        self.refused.fetch_add(1, Ordering::Relaxed);
184        tracing::debug!(target: TRACING_TARGET, uri = %seed.uri(), "refused blocked URL prefix");
185    }
186}
187
188impl<F: Frontier> Frontier for AvoidTraps<F> {
189    async fn push(&self, seed: Seed) -> Result<Admission> {
190        if let Some(admission) = self.refusal(&seed) {
191            self.count_refusal(&seed);
192            return Ok(admission);
193        }
194        self.inner.push(seed).await
195    }
196
197    async fn pop(&self) -> Result<NextSeed> {
198        loop {
199            match self.inner.pop().await? {
200                NextSeed::Ready(seed) if self.refusal(&seed).is_some() => {
201                    self.count_refusal(&seed);
202                    self.inner.report(Report::Skipped(seed)).await?;
203                }
204                pop => return Ok(pop),
205            }
206        }
207    }
208
209    fn report(&self, report: Report) -> impl Future<Output = Result<Ack>> + Send {
210        self.inner.report(report)
211    }
212
213    fn size_hint(&self) -> Option<usize> {
214        self.inner.size_hint()
215    }
216
217    fn snapshot(&self) -> Option<Vec<Seed>> {
218        self.inner.snapshot().map(|seeds| {
219            seeds
220                .into_iter()
221                .filter(|seed| !self.is_blocked(seed))
222                .collect()
223        })
224    }
225}
226
227fn prefix(uri: &Uri) -> String {
228    let origin = url::Url::parse(&uri.to_string())
229        .map_or_else(|_| String::new(), |url| url.origin().ascii_serialization());
230    let first = uri
231        .path()
232        .split('/')
233        .find(|part| !part.is_empty())
234        .unwrap_or_default();
235    if first.is_empty() || opaque(first) {
236        format!("{origin}/")
237    } else {
238        format!("{origin}/{first}/")
239    }
240}
241
242fn opaque(segment: &str) -> bool {
243    segment.len() >= 16
244        && segment.bytes().all(|byte| byte.is_ascii_alphanumeric())
245        && segment.bytes().any(|byte| byte.is_ascii_digit())
246        && segment.bytes().any(|byte| byte.is_ascii_alphabetic())
247        && segment.bytes().collect::<HashSet<_>>().len() >= 8
248}
249
250#[cfg(test)]
251mod tests {
252    use std::sync::Arc;
253
254    use super::*;
255    use crate::frontier::{FrontierExt, InMemFrontier};
256
257    fn seed(path: &str) -> Seed {
258        Seed::get(format!("https://site.test{path}")).unwrap()
259    }
260
261    #[tokio::test]
262    async fn blocks_only_the_over_budget_prefix_and_skips_pending_seeds() {
263        let frontier = InMemFrontier::fifo()
264            .dedup()
265            .avoid_traps_with(TrapPolicy::default().with_max_urls(2));
266        for path in ["/trap/a", "/trap/b", "/safe/a"] {
267            assert!(frontier.push(seed(path)).await.unwrap().is_queued());
268        }
269        assert_eq!(
270            frontier.push(seed("/trap/c")).await.unwrap(),
271            Admission::PrefixBudgetExceeded
272        );
273        assert_eq!(
274            frontier.push(seed("/trap/d")).await.unwrap(),
275            Admission::PrefixBudgetExceeded
276        );
277        assert_eq!(frontier.snapshot().unwrap().len(), 1);
278        let NextSeed::Ready(safe) = frontier.pop().await.unwrap() else {
279            panic!("safe seed remains");
280        };
281        assert_eq!(safe.uri().path(), "/safe/a");
282        assert!(matches!(frontier.pop().await.unwrap(), NextSeed::Empty));
283        assert_eq!(frontier.refused(), 4);
284    }
285
286    #[tokio::test]
287    async fn detects_fanout_depth_repetition_and_opaque_segments() {
288        let cases = [
289            (TrapPolicy::default().with_max_fanout(1), "/pit/a", "/pit/b"),
290            (
291                TrapPolicy::default().with_max_path_depth(3),
292                "/pit/a/b",
293                "/pit/a/b/c",
294            ),
295            (TrapPolicy::default(), "/pit/a/a/a/a", "/pit/a/a/a/a/a"),
296            (
297                TrapPolicy::default().with_max_random_urls(1),
298                "/pit/aB12cD34eF56gH78",
299                "/pit/iJ90kL12mN34oP56",
300            ),
301        ];
302        for (index, (policy, allowed, refused)) in cases.into_iter().enumerate() {
303            let frontier = InMemFrontier::fifo().avoid_traps_with(policy);
304            assert!(frontier.push(seed(allowed)).await.unwrap().is_queued());
305            assert_eq!(
306                frontier.push(seed(refused)).await.unwrap(),
307                if index < 2 {
308                    Admission::PrefixBudgetExceeded
309                } else {
310                    Admission::SuspectedTrap
311                },
312                "{refused}"
313            );
314        }
315    }
316
317    #[tokio::test]
318    async fn repeated_candidates_do_not_spend_the_budget_again() {
319        let frontier = InMemFrontier::fifo()
320            .dedup()
321            .avoid_traps_with(TrapPolicy::default().with_max_urls(1));
322        assert!(frontier.push(seed("/pit/a")).await.unwrap().is_queued());
323        assert_eq!(
324            frontier.push(seed("/pit/a")).await.unwrap(),
325            Admission::Duplicate
326        );
327        assert_eq!(frontier.refused(), 0);
328    }
329
330    #[tokio::test]
331    async fn concurrent_pushes_cannot_bypass_the_cap() {
332        let frontier = Arc::new(
333            InMemFrontier::fifo().avoid_traps_with(TrapPolicy::default().with_max_urls(4)),
334        );
335        let mut tasks = Vec::new();
336        for n in 0..20 {
337            let frontier = frontier.clone();
338            tasks.push(tokio::spawn(async move {
339                frontier.push(seed(&format!("/pit/{n}"))).await.unwrap()
340            }));
341        }
342        let mut queued = 0;
343        for task in tasks {
344            queued += usize::from(task.await.unwrap().is_queued());
345        }
346        assert_eq!(queued, 4);
347        assert_eq!(frontier.refused(), 16);
348    }
349
350    #[tokio::test]
351    async fn checks_preexisting_storage_on_pop() {
352        let stored = InMemFrontier::fifo();
353        stored.push(seed("/pit/a")).await.unwrap();
354        stored.push(seed("/pit/b")).await.unwrap();
355        let frontier = stored.avoid_traps_with(TrapPolicy::default().with_max_urls(1));
356        assert!(matches!(frontier.pop().await.unwrap(), NextSeed::Ready(_)));
357        assert!(matches!(frontier.pop().await.unwrap(), NextSeed::Empty));
358        assert_eq!(frontier.refused(), 1);
359    }
360
361    #[tokio::test(start_paused = true)]
362    async fn a_blocked_prefix_does_not_escape_through_a_delayed_retry() {
363        let frontier = InMemFrontier::fifo()
364            .retry(1)
365            .avoid_traps_with(TrapPolicy::default().with_max_urls(1));
366        frontier.push(seed("/pit/a")).await.unwrap();
367        let NextSeed::Ready(retried) = frontier.pop().await.unwrap() else {
368            panic!("seed");
369        };
370        frontier
371            .report(Report::Retry {
372                seed: retried,
373                after: Some(std::time::Duration::from_secs(1)),
374            })
375            .await
376            .unwrap();
377        assert_eq!(
378            frontier.push(seed("/pit/b")).await.unwrap(),
379            Admission::PrefixBudgetExceeded
380        );
381        tokio::time::advance(std::time::Duration::from_secs(2)).await;
382        assert!(matches!(frontier.pop().await.unwrap(), NextSeed::Empty));
383        assert_eq!(frontier.refused(), 2);
384    }
385
386    #[tokio::test]
387    async fn opaque_root_segments_share_a_budget_and_real_ids_can_opt_out() {
388        let ids = ["aB12cD34eF56gH78", "iJ90kL12mN34oP56"];
389        let guarded =
390            InMemFrontier::fifo().avoid_traps_with(TrapPolicy::default().with_max_random_urls(1));
391        assert!(
392            guarded
393                .push(seed(&format!("/{}", ids[0])))
394                .await
395                .unwrap()
396                .is_queued()
397        );
398        assert_eq!(
399            guarded.push(seed(&format!("/{}", ids[1]))).await.unwrap(),
400            Admission::SuspectedTrap
401        );
402        let ids_allowed = InMemFrontier::fifo()
403            .avoid_traps_with(TrapPolicy::default().with_max_random_urls(usize::MAX));
404        for id in ids {
405            assert!(
406                ids_allowed
407                    .push(seed(&format!("/products/{id}")))
408                    .await
409                    .unwrap()
410                    .is_queued()
411            );
412        }
413    }
414
415    #[tokio::test]
416    async fn origins_ports_and_path_boundaries_are_independent() {
417        let frontier =
418            InMemFrontier::fifo().avoid_traps_with(TrapPolicy::default().with_max_urls(1));
419        frontier.push(seed("/pit/a")).await.unwrap();
420        assert_eq!(
421            frontier.push(seed("/pit/b")).await.unwrap(),
422            Admission::PrefixBudgetExceeded
423        );
424        for uri in [
425            "https://site.test/pitfall/a",
426            "https://other.test/pit/a",
427            "http://site.test/pit/a",
428            "https://site.test:8443/pit/a",
429        ] {
430            assert!(
431                frontier
432                    .push(Seed::get(uri).unwrap())
433                    .await
434                    .unwrap()
435                    .is_queued(),
436                "{uri}"
437            );
438        }
439    }
440}