spire_core/frontier/facade/
traps.rs1use 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#[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 #[must_use]
43 pub const fn with_max_urls(mut self, max: usize) -> Self {
44 self.urls = max;
45 self
46 }
47
48 #[must_use]
52 pub const fn with_max_fanout(mut self, max: usize) -> Self {
53 self.fanout = max;
54 self
55 }
56
57 #[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 #[must_use]
67 pub const fn with_max_repetitions(mut self, max: usize) -> Self {
68 self.repetitions = max;
69 self
70 }
71
72 #[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#[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 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 #[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 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}