leyline/core/crawl/
frontier.rs

1use std::collections::{HashMap, HashSet, VecDeque};
2
3use url::{Host, Url};
4
5use crate::cookie::parse::registrable_domain;
6use crate::core::config::{Origin, normalize_host};
7use crate::core::robots::normalize;
8use crate::html::WEB_SCHEMES;
9
10const MAX_URL_BYTES: usize = 2048;
11pub(super) const MAX_DEPTH: usize = 64;
12
13#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
14#[non_exhaustive]
15pub enum Scope {
16    #[default]
17    Host,
18    Domain,
19    Prefix,
20}
21
22enum Bound {
23    Host(String),
24    Domain(String),
25    Prefix { host: String, dir: String },
26}
27
28fn host(url: &Url) -> Option<String> {
29    Some(normalize_host(url.host_str()?))
30}
31
32fn is_ip(url: &Url) -> bool {
33    matches!(url.host(), Some(Host::Ipv4(_) | Host::Ipv6(_)))
34}
35
36fn directory(path: &str) -> String {
37    let last = path.rsplit('/').next().unwrap_or_default();
38    let named_file = last
39        .rsplit_once('.')
40        .is_some_and(|(_, ext)| ext.bytes().any(|byte| byte.is_ascii_alphabetic()));
41    if named_file {
42        let end = path.rfind('/').map_or(0, |end| end + 1);
43        return path[..end].to_owned();
44    }
45    format!("{}/", path.trim_end_matches('/'))
46}
47
48impl Bound {
49    fn of(scope: Scope, url: &Url) -> Option<Bound> {
50        let host = host(url)?;
51        Some(match scope {
52            Scope::Domain if !is_ip(url) => {
53                Bound::Domain(registrable_domain(&host).unwrap_or(host))
54            }
55            Scope::Host | Scope::Domain => Bound::Host(host),
56            Scope::Prefix => Bound::Prefix {
57                host,
58                dir: directory(url.path()),
59            },
60        })
61    }
62
63    fn holds(&self, url: &Url) -> bool {
64        let Some(candidate) = host(url) else {
65            return false;
66        };
67        match self {
68            Bound::Host(bound) => candidate == *bound,
69            Bound::Domain(domain) => {
70                !is_ip(url) && registrable_domain(&candidate).unwrap_or(candidate) == *domain
71            }
72            Bound::Prefix { host, dir } => {
73                let path = url.path();
74                candidate == *host
75                    && (path.starts_with(dir.as_str()) || path == dir.trim_end_matches('/'))
76            }
77        }
78    }
79}
80
81#[derive(Debug, Clone)]
82pub(super) struct Queued {
83    pub(super) url: Url,
84    pub(super) depth: usize,
85    pub(super) hops: usize,
86    pub(super) seed: bool,
87    pub(super) widen: bool,
88    pub(super) attempts: u32,
89    pub(super) chain: Vec<String>,
90}
91
92impl Queued {
93    pub(super) fn link(url: Url, depth: usize) -> Queued {
94        Queued {
95            url,
96            depth,
97            hops: 0,
98            seed: false,
99            widen: false,
100            attempts: 0,
101            chain: Vec::new(),
102        }
103    }
104
105    pub(super) fn seed(url: Url) -> Queued {
106        Queued {
107            seed: true,
108            widen: true,
109            ..Queued::link(url, 0)
110        }
111    }
112}
113
114#[derive(Default)]
115pub(super) struct Frontier {
116    queues: HashMap<Origin, VecDeque<Queued>>,
117    order: VecDeque<Origin>,
118    seen: HashSet<String>,
119    redirected: HashSet<String>,
120    bounds: Vec<Bound>,
121    pub(super) scope: Scope,
122}
123
124pub(super) fn same_site(from: &Url, to: &Url) -> bool {
125    match (Bound::of(Scope::Domain, from), Bound::of(Scope::Domain, to)) {
126        (Some(Bound::Domain(from)), Some(Bound::Domain(to))) => from == to,
127        _ => host(from) == host(to),
128    }
129}
130
131pub(super) fn page_key(url: &Url) -> String {
132    let mut key = format!("{}://{}", url.scheme(), host(url).unwrap_or_default());
133    if let Some(port) = url.port() {
134        key.push_str(&format!(":{port}"));
135    }
136    key.push_str(&normalize(url.path()));
137    if let Some(query) = url.query() {
138        let mut pairs: Vec<String> = query
139            .split('&')
140            .filter(|pair| !pair.is_empty())
141            .map(normalize)
142            .collect();
143        pairs.sort_unstable();
144        key.push('?');
145        key.push_str(&pairs.join("&"));
146    }
147    key
148}
149
150impl Frontier {
151    pub(super) fn add(&mut self, queued: Queued) -> bool {
152        self.admit(queued, false)
153    }
154
155    pub(super) fn add_redirect(&mut self, queued: Queued) -> bool {
156        self.admit(queued, true)
157    }
158
159    fn admit(&mut self, mut queued: Queued, redirect: bool) -> bool {
160        queued.url.set_fragment(None);
161        let usable = WEB_SCHEMES.contains(&queued.url.scheme())
162            && queued.url.as_str().len() <= MAX_URL_BYTES
163            && queued.depth <= MAX_DEPTH;
164        if !usable {
165            return false;
166        }
167        if queued.widen
168            && let Some(bound) = Bound::of(self.scope, &queued.url)
169        {
170            self.bounds.push(bound);
171        }
172        if !queued.seed && !self.in_scope(&queued.url) {
173            return false;
174        }
175        let key = page_key(&queued.url);
176        let again = redirect && self.redirected.remove(&key);
177        if !self.seen.insert(key) && !again {
178            return false;
179        }
180        self.push_back(queued);
181        true
182    }
183
184    pub(super) fn rekey(&mut self, mut queued: Queued, to: Url) {
185        self.seen.remove(&page_key(&queued.url));
186        queued.url = to;
187        self.add(queued);
188    }
189
190    pub(super) fn has_seen(&self, url: &Url) -> bool {
191        self.seen.contains(&page_key(url))
192    }
193
194    pub(super) fn redirected(&mut self, url: &Url) {
195        self.redirected.insert(page_key(url));
196    }
197
198    pub(super) fn in_scope(&self, url: &Url) -> bool {
199        self.bounds.iter().any(|bound| bound.holds(url))
200    }
201
202    pub(super) fn known(&self) -> usize {
203        self.seen.len()
204    }
205
206    pub(super) fn push_back(&mut self, queued: Queued) {
207        let Some(origin) = Origin::of(&queued.url) else {
208            return;
209        };
210        let queue = self.queues.entry(origin.clone()).or_default();
211        if queue.is_empty() {
212            self.order.push_back(origin);
213        }
214        queue.push_back(queued);
215    }
216
217    pub(super) fn push_front(&mut self, queued: Queued) {
218        let Some(origin) = Origin::of(&queued.url) else {
219            return;
220        };
221        let queue = self.queues.entry(origin.clone()).or_default();
222        if queue.is_empty() {
223            self.order.push_front(origin);
224        }
225        queue.push_front(queued);
226    }
227
228    pub(super) fn pop_where(&mut self, mut ready: impl FnMut(&Origin) -> bool) -> Option<Queued> {
229        for _ in 0..self.order.len() {
230            let origin = self.order.pop_front()?;
231            if !ready(&origin) {
232                self.order.push_back(origin);
233                continue;
234            }
235            let queue = self.queues.get_mut(&origin)?;
236            let queued = queue.pop_front();
237            if queue.is_empty() {
238                self.queues.remove(&origin);
239            } else {
240                self.order.push_back(origin);
241            }
242            return queued;
243        }
244        None
245    }
246
247    pub(super) fn origins(&self) -> impl Iterator<Item = &Origin> {
248        self.queues.keys()
249    }
250
251    pub(super) fn is_empty(&self) -> bool {
252        self.queues.is_empty()
253    }
254
255    pub(super) fn clear(&mut self) {
256        self.queues.clear();
257        self.order.clear();
258    }
259}
leyline/core/crawl/frontier.rs

leyline/core/crawl/
frontier.rs
