use crate::{
enqueue::SkipReason,
errors::CrawlError,
request::{HeaderMap, Method, UserData},
};
use globset::{GlobBuilder, GlobSet, GlobSetBuilder};
use regex::Regex;
use std::{fmt, sync::Arc};
use url::{Host, Url};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum EnqueueStrategy {
All,
#[default]
SameHostname,
SameDomain,
SameOrigin,
}
pub fn strategy_allows(strategy: EnqueueStrategy, parent: &Url, candidate: &Url) -> bool {
if !matches!(candidate.scheme(), "http" | "https") {
return false;
}
match strategy {
EnqueueStrategy::All => true,
EnqueueStrategy::SameHostname => hosts_equal(parent.host_str(), candidate.host_str()),
EnqueueStrategy::SameDomain => same_domain(parent, candidate),
EnqueueStrategy::SameOrigin => {
parent.scheme() == candidate.scheme()
&& hosts_equal(parent.host_str(), candidate.host_str())
&& parent.port_or_known_default() == candidate.port_or_known_default()
}
}
}
fn hosts_equal(left: Option<&str>, right: Option<&str>) -> bool {
match (left, right) {
(Some(left), Some(right)) => left.eq_ignore_ascii_case(right),
_ => false,
}
}
fn same_domain(parent: &Url, candidate: &Url) -> bool {
let (Some(parent_host), Some(candidate_host)) = (parent.host(), candidate.host()) else {
return false;
};
match (parent_host, candidate_host) {
(Host::Domain(parent_host), Host::Domain(candidate_host)) => {
match (
psl::domain_str(parent_host),
psl::domain_str(candidate_host),
) {
(Some(parent_domain), Some(candidate_domain)) => {
parent_domain.eq_ignore_ascii_case(candidate_domain)
}
_ => parent_host.eq_ignore_ascii_case(candidate_host),
}
}
(parent_host, candidate_host) => parent_host == candidate_host,
}
}
#[derive(Debug, Clone)]
#[non_exhaustive]
pub enum UrlPattern {
Glob(String),
Regex(Regex),
}
impl From<&str> for UrlPattern {
fn from(pattern: &str) -> Self {
Self::Glob(pattern.to_owned())
}
}
impl From<String> for UrlPattern {
fn from(pattern: String) -> Self {
Self::Glob(pattern)
}
}
impl From<Regex> for UrlPattern {
fn from(pattern: Regex) -> Self {
Self::Regex(pattern)
}
}
#[allow(dead_code)] pub(crate) fn compile_globs(patterns: &[String]) -> Result<GlobSet, LinkPatternError> {
let mut builder = GlobSetBuilder::new();
for pattern in patterns {
let glob = GlobBuilder::new(pattern)
.literal_separator(false)
.build()
.map_err(|source| LinkPatternError::InvalidGlob {
pattern: pattern.clone(),
source,
})?;
builder.add(glob);
}
builder
.build()
.map_err(|source| LinkPatternError::InvalidGlob {
pattern: patterns.last().cloned().unwrap_or_default(),
source,
})
}
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum LinkPatternError {
#[error("invalid glob pattern {pattern:?}: {source}")]
InvalidGlob {
pattern: String,
#[source]
source: globset::Error,
},
}
#[derive(Debug, Clone)]
pub struct UrlMatch {
pub pattern: UrlPattern,
pub label: Option<String>,
pub user_data: Option<UserData>,
pub method: Option<Method>,
pub headers: Option<HeaderMap>,
}
impl UrlMatch {
pub fn new(pattern: impl Into<UrlPattern>) -> Self {
Self {
pattern: pattern.into(),
label: None,
user_data: None,
method: None,
headers: None,
}
}
pub fn label(mut self, label: impl Into<String>) -> Self {
self.label = Some(label.into());
self
}
pub fn user_data(mut self, user_data: UserData) -> Self {
self.user_data = Some(user_data);
self
}
pub fn method(mut self, method: Method) -> Self {
self.method = Some(method);
self
}
pub fn headers(mut self, headers: HeaderMap) -> Self {
self.headers = Some(headers);
self
}
}
#[derive(Debug, Clone)]
pub struct GlobPattern {
matched: UrlMatch,
}
impl From<&str> for GlobPattern {
fn from(pattern: &str) -> Self {
Self {
matched: UrlMatch::new(pattern),
}
}
}
impl From<String> for GlobPattern {
fn from(pattern: String) -> Self {
Self {
matched: UrlMatch::new(pattern),
}
}
}
impl From<Regex> for GlobPattern {
fn from(pattern: Regex) -> Self {
Self {
matched: UrlMatch::new(pattern),
}
}
}
impl From<UrlMatch> for GlobPattern {
fn from(matched: UrlMatch) -> Self {
Self { matched }
}
}
#[allow(dead_code)]
impl GlobPattern {
pub(crate) fn pattern(&self) -> &UrlPattern {
&self.matched.pattern
}
pub(crate) fn label(&self) -> Option<&str> {
self.matched.label.as_deref()
}
pub(crate) fn user_data(&self) -> Option<&UserData> {
self.matched.user_data.as_ref()
}
pub(crate) fn method(&self) -> Option<&Method> {
self.matched.method.as_ref()
}
pub(crate) fn headers(&self) -> Option<&HeaderMap> {
self.matched.headers.as_ref()
}
}
#[derive(Debug, Clone)]
pub struct ExtractedLink {
pub url: String,
pub base: Option<Url>,
}
#[async_trait::async_trait]
pub trait LinkExtractor: Send + Sync {
async fn extract(&self, selector: Option<&str>) -> Result<Vec<ExtractedLink>, CrawlError>;
}
#[derive(Debug)]
#[non_exhaustive]
pub enum TransformResult {
Enqueue,
Skip {
reason: String,
},
}
pub trait SkippedHandler: Send + Sync + 'static {
fn on_skip(&self, url: &str, reason: &SkipReason);
}
impl<F> SkippedHandler for F
where
F: Fn(&str, &SkipReason) + Send + Sync + 'static,
{
fn on_skip(&self, url: &str, reason: &SkipReason) {
self(url, reason);
}
}
#[non_exhaustive]
#[derive(Default)]
pub struct CrawlPolicy {
pub strategy: EnqueueStrategy,
pub max_crawl_depth: Option<u32>,
pub max_requests_per_crawl: Option<u64>,
pub on_skipped: Option<Arc<dyn SkippedHandler>>,
}
impl CrawlPolicy {
pub fn new() -> Self {
Self::default()
}
pub fn strategy(mut self, strategy: EnqueueStrategy) -> Self {
self.strategy = strategy;
self
}
pub fn max_crawl_depth(mut self, max_crawl_depth: u32) -> Self {
self.max_crawl_depth = Some(max_crawl_depth);
self
}
pub fn max_requests_per_crawl(mut self, max_requests_per_crawl: u64) -> Self {
self.max_requests_per_crawl = Some(max_requests_per_crawl);
self
}
pub fn on_skipped<H: SkippedHandler>(mut self, handler: H) -> Self {
self.on_skipped = Some(Arc::new(handler));
self
}
}
impl fmt::Debug for CrawlPolicy {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("CrawlPolicy")
.field("strategy", &self.strategy)
.field("max_crawl_depth", &self.max_crawl_depth)
.field("max_requests_per_crawl", &self.max_requests_per_crawl)
.field(
"on_skipped",
&self.on_skipped.as_ref().map(|_| "<dyn SkippedHandler>"),
)
.finish()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn compiled_globs_match_complete_url_strings() {
let patterns = vec![
"**/products/*".to_owned(),
"https://example.com/**".to_owned(),
];
let compiled = compile_globs(&patterns).expect("valid globs");
assert!(compiled.is_match("https://shop.test/products/p1"));
assert!(compiled.is_match("https://example.com/anything/here"));
assert!(!compiled.is_match("https://shop.test/categories/c1"));
}
#[test]
fn invalid_glob_retains_source_pattern() {
let error = compile_globs(&["[".to_owned()]).expect_err("glob should be invalid");
assert!(matches!(
error,
LinkPatternError::InvalidGlob { ref pattern, .. } if pattern == "["
));
}
#[test]
fn glob_pattern_accessors_expose_overrides() {
let mut user_data = UserData::default();
user_data
.set_typed("kind", &"product")
.expect("serializable data");
let mut headers = HeaderMap::new();
headers.insert("x-test", "yes".parse().expect("valid header value"));
let pattern = GlobPattern::from(
UrlMatch::new(Regex::new("products").expect("valid regex"))
.label("product")
.user_data(user_data)
.method(Method::POST)
.headers(headers),
);
assert!(matches!(pattern.pattern(), UrlPattern::Regex(_)));
assert_eq!(pattern.label(), Some("product"));
assert_eq!(
pattern.user_data().and_then(|data| data.get("kind")),
Some(&serde_json::json!("product"))
);
assert_eq!(pattern.method(), Some(&Method::POST));
assert_eq!(
pattern.headers().and_then(|map| map.get("x-test")),
Some(&"yes".parse().expect("valid header value"))
);
}
}