use std::collections::{BTreeMap, HashMap};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use chrono::{DateTime, Utc};
use futures_util::StreamExt;
use reqwest::header::{HeaderMap, HeaderName, HeaderValue, CONTENT_TYPE, COOKIE, USER_AGENT};
use robots_txt::{matcher::SimpleMatcher, parts::RequestRate, Robots};
use serde::{Deserialize, Serialize};
use tokio::time::sleep;
use url::Url;
use crate::config::Config;
use crate::errors::{Result, SiteforgeError};
const ROBOTS_TXT_MAX_BYTES: u64 = 1024 * 1024;
#[derive(Debug, Clone)]
pub struct HttpFetcher {
client: reqwest::Client,
user_agent: String,
retry_count: usize,
delay: Duration,
robots_cache: Arc<Mutex<HashMap<String, String>>>,
last_request_start: Arc<Mutex<HashMap<String, Instant>>>,
extra_headers: HeaderMap,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FetchResult {
pub requested_url: Url,
pub final_url: Url,
pub status: u16,
pub headers: BTreeMap<String, String>,
pub mime_type: Option<String>,
pub fetched_at: DateTime<Utc>,
pub content_hash: String,
pub bytes: Vec<u8>,
}
impl HttpFetcher {
pub fn new(
config: &Config,
delay_ms: u64,
headers: &[(String, String)],
cookie: Option<&str>,
) -> Result<Self> {
let client = reqwest::Client::builder()
.user_agent(config.user_agent.clone())
.timeout(Duration::from_secs(config.timeout_secs))
.redirect(reqwest::redirect::Policy::limited(10))
.build()?;
Ok(Self {
client,
user_agent: config.user_agent.clone(),
retry_count: config.retry_count,
delay: Duration::from_millis(delay_ms),
robots_cache: Arc::new(Mutex::new(HashMap::new())),
last_request_start: Arc::new(Mutex::new(HashMap::new())),
extra_headers: build_extra_headers(headers, cookie)?,
})
}
pub async fn fetch(&self, url: Url) -> Result<FetchResult> {
self.fetch_with_limit(url, 0).await
}
pub async fn fetch_with_limit(&self, url: Url, max_body_bytes: u64) -> Result<FetchResult> {
let robots_delay = self.ensure_robots_allowed(&url).await?;
self.wait_for_origin(&url, robots_delay).await;
let mut last_error = None;
for attempt in 0..=self.retry_count {
match self.fetch_once(url.clone(), max_body_bytes).await {
Ok(result) if should_retry_status(result.status) && attempt < self.retry_count => {
last_error = Some(format!("HTTP {}", result.status));
sleep(backoff(attempt)).await;
}
Ok(result) => return Ok(result),
Err(err @ SiteforgeError::SizeLimitExceeded { .. }) => return Err(err),
Err(err) if attempt < self.retry_count => {
last_error = Some(err.to_string());
sleep(backoff(attempt)).await;
}
Err(err) => return Err(err),
}
}
Err(SiteforgeError::message(
last_error.unwrap_or_else(|| format!("failed to fetch {url}")),
))
}
pub async fn robots_sitemaps(&self, url: &Url) -> Result<Vec<Url>> {
let robots_txt = self.robots_text(url).await?;
Ok(sitemap_urls_from_robots(url, &robots_txt))
}
async fn fetch_once(&self, url: Url, max_body_bytes: u64) -> Result<FetchResult> {
let response = self.request(url.clone()).send().await?;
let status = response.status().as_u16();
let final_url = response.url().clone();
let headers = headers_to_map(response.headers());
let mime_type = response
.headers()
.get(CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.map(|value| {
value
.split(';')
.next()
.unwrap_or(value)
.trim()
.to_ascii_lowercase()
});
let bytes = response_body_limited(response, max_body_bytes, final_url.as_str()).await?;
let content_hash = blake3::hash(&bytes).to_hex().to_string();
Ok(FetchResult {
requested_url: url,
final_url,
status,
headers,
mime_type,
fetched_at: Utc::now(),
content_hash,
bytes,
})
}
async fn ensure_robots_allowed(&self, url: &Url) -> Result<Option<Duration>> {
let origin = origin_key(url);
let robots_txt = self.robots_text_for_origin(url, &origin).await?;
if robots_txt.trim().is_empty() {
return Ok(None);
}
let robots = Robots::from_str_lossy(&robots_txt);
let section = robots.choose_section(&self.user_agent);
let matcher = SimpleMatcher::new(§ion.rules);
let robots_delay = [
duration_from_crawl_delay(section.crawl_delay),
duration_from_request_rate(section.req_rate),
]
.into_iter()
.flatten()
.max();
let mut path = url.path().to_string();
if let Some(query) = url.query() {
path.push('?');
path.push_str(query);
}
if matcher.check_path(&path) {
Ok(robots_delay)
} else {
Err(SiteforgeError::RobotsDenied(url.to_string()))
}
}
async fn fetch_robots(&self, url: &Url) -> Result<String> {
let robots_url = robots_url(url)?;
let response = self.request(robots_url).send().await?;
if !response.status().is_success() {
return Ok(String::new());
}
let label = response.url().to_string();
let bytes = response_body_limited(response, ROBOTS_TXT_MAX_BYTES, &label).await?;
Ok(bytes_to_string(&bytes))
}
async fn robots_text(&self, url: &Url) -> Result<String> {
let origin = origin_key(url);
self.robots_text_for_origin(url, &origin).await
}
async fn robots_text_for_origin(&self, url: &Url, origin: &str) -> Result<String> {
let cached = self
.robots_cache
.lock()
.expect("robots cache poisoned")
.get(origin)
.cloned();
match cached {
Some(value) => Ok(value),
None => {
let value = self.fetch_robots(url).await.unwrap_or_default();
self.robots_cache
.lock()
.expect("robots cache poisoned")
.insert(origin.to_string(), value.clone());
Ok(value)
}
}
}
async fn wait_for_origin(&self, url: &Url, robots_delay: Option<Duration>) {
let min_delay = robots_delay
.map(|delay| delay.max(self.delay))
.unwrap_or(self.delay);
if min_delay.is_zero() {
return;
}
let origin = origin_key(url);
loop {
let sleep_for = {
let mut starts = self
.last_request_start
.lock()
.expect("request pacing state poisoned");
let now = Instant::now();
match starts.get(&origin).copied() {
Some(last) => {
let elapsed = now.saturating_duration_since(last);
if elapsed >= min_delay {
starts.insert(origin.clone(), now);
None
} else {
Some(min_delay - elapsed)
}
}
None => {
starts.insert(origin.clone(), now);
None
}
}
};
match sleep_for {
Some(duration) => sleep(duration).await,
None => break,
}
}
}
fn request(&self, url: Url) -> reqwest::RequestBuilder {
let mut request = self
.client
.get(url)
.header(USER_AGENT, self.user_agent.clone());
for (name, value) in &self.extra_headers {
request = request.header(name, value);
}
request
}
}
async fn response_body_limited(
response: reqwest::Response,
max_body_bytes: u64,
label: &str,
) -> Result<Vec<u8>> {
if max_body_bytes > 0 {
if let Some(length) = response.content_length() {
if length > max_body_bytes {
return Err(SiteforgeError::SizeLimitExceeded {
path: label.to_string(),
size: length,
limit: max_body_bytes,
});
}
}
}
let mut bytes = Vec::new();
let mut stream = response.bytes_stream();
while let Some(chunk) = stream.next().await {
let chunk = chunk?;
let next_len = (bytes.len() as u64).saturating_add(chunk.len() as u64);
if max_body_bytes > 0 && next_len > max_body_bytes {
return Err(SiteforgeError::SizeLimitExceeded {
path: label.to_string(),
size: next_len,
limit: max_body_bytes,
});
}
bytes.extend_from_slice(&chunk);
}
Ok(bytes)
}
pub fn is_html(result: &FetchResult) -> bool {
result
.mime_type
.as_deref()
.map(|mime| mime.contains("html") || mime == "text/plain")
.unwrap_or_else(|| looks_like_html(&result.bytes))
}
pub fn bytes_to_string(bytes: &[u8]) -> String {
String::from_utf8_lossy(bytes).into_owned()
}
fn looks_like_html(bytes: &[u8]) -> bool {
let sample = String::from_utf8_lossy(&bytes[..bytes.len().min(512)]).to_ascii_lowercase();
sample.contains("<html") || sample.contains("<!doctype html") || sample.contains("<article")
}
fn should_retry_status(status: u16) -> bool {
matches!(status, 408 | 425 | 429 | 500 | 502 | 503 | 504)
}
fn backoff(attempt: usize) -> Duration {
Duration::from_millis(500 * 2u64.saturating_pow(attempt as u32))
}
fn duration_from_crawl_delay(seconds: Option<f64>) -> Option<Duration> {
seconds.and_then(|seconds| {
(seconds.is_finite() && seconds > 0.0).then(|| Duration::from_secs_f64(seconds.min(3600.0)))
})
}
fn duration_from_request_rate(rate: Option<RequestRate>) -> Option<Duration> {
rate.and_then(|rate| {
(rate.requests > 0 && rate.seconds > 0)
.then(|| Duration::from_secs_f64(rate.seconds as f64 / rate.requests as f64))
})
}
fn build_extra_headers(headers: &[(String, String)], cookie: Option<&str>) -> Result<HeaderMap> {
let mut map = HeaderMap::new();
for (name, value) in headers {
let name = HeaderName::from_bytes(name.as_bytes()).map_err(|err| {
SiteforgeError::message(format!("invalid header name {name:?}: {err}"))
})?;
let value = HeaderValue::from_str(value).map_err(|err| {
SiteforgeError::message(format!("invalid value for header {name:?}: {err}"))
})?;
map.insert(name, value);
}
if let Some(cookie) = cookie.filter(|cookie| !cookie.trim().is_empty()) {
let value = HeaderValue::from_str(cookie).map_err(|err| {
SiteforgeError::message(format!("invalid cookie header value: {err}"))
})?;
map.insert(COOKIE, value);
}
Ok(map)
}
fn headers_to_map(headers: &HeaderMap) -> BTreeMap<String, String> {
headers
.iter()
.filter_map(|(key, value)| {
value
.to_str()
.ok()
.map(|value| (key.as_str().to_ascii_lowercase(), value.to_string()))
})
.collect()
}
fn origin_key(url: &Url) -> String {
match url.port() {
Some(port) => format!(
"{}://{}:{port}",
url.scheme(),
url.host_str().unwrap_or_default()
),
None => format!("{}://{}", url.scheme(), url.host_str().unwrap_or_default()),
}
}
fn robots_url(url: &Url) -> Result<Url> {
let mut robots = url.clone();
robots.set_path("/robots.txt");
robots.set_query(None);
robots.set_fragment(None);
Ok(robots)
}
fn sitemap_urls_from_robots(base: &Url, robots_txt: &str) -> Vec<Url> {
let mut urls = robots_txt
.lines()
.filter_map(|line| line.split('#').next())
.filter_map(|line| line.split_once(':'))
.filter(|(key, _)| key.trim().eq_ignore_ascii_case("sitemap"))
.filter_map(|(_, value)| {
Url::parse(value.trim())
.or_else(|_| base.join(value.trim()))
.ok()
})
.collect::<Vec<_>>();
urls.sort();
urls.dedup();
urls
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::Config;
use std::io::{Read, Write};
use std::net::TcpListener;
use std::thread;
use wiremock::matchers::{method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
#[test]
fn parses_positive_crawl_delay() {
assert_eq!(
duration_from_crawl_delay(Some(1.5)),
Some(Duration::from_millis(1500))
);
assert_eq!(duration_from_crawl_delay(Some(0.0)), None);
assert_eq!(duration_from_crawl_delay(Some(f64::NAN)), None);
}
#[test]
fn parses_request_rate_delay() {
assert_eq!(
duration_from_request_rate(Some(RequestRate::new(2, 5))),
Some(Duration::from_millis(2500))
);
assert_eq!(
duration_from_request_rate(Some(RequestRate::new(0, 5))),
None
);
}
#[test]
fn origin_key_preserves_non_default_port() {
let url = Url::parse("https://example.test:8443/path").unwrap();
assert_eq!(origin_key(&url), "https://example.test:8443");
}
#[test]
fn builds_authorized_request_headers() {
let headers = build_extra_headers(
&[("Authorization".to_string(), "Bearer token".to_string())],
Some("session=abc"),
)
.unwrap();
assert_eq!(headers.get("authorization").unwrap(), "Bearer token");
assert_eq!(headers.get(COOKIE).unwrap(), "session=abc");
}
#[test]
fn rejects_invalid_request_headers() {
let err = build_extra_headers(&[("Bad Header".to_string(), "value".to_string())], None)
.unwrap_err()
.to_string();
assert!(err.contains("invalid header name"));
}
#[test]
fn extracts_sitemaps_from_robots() {
let base = Url::parse("https://example.test/docs/").unwrap();
let urls = sitemap_urls_from_robots(
&base,
"User-agent: *\nSitemap: /sitemap.xml\nSitemap: https://cdn.example.test/sitemap.xml\n",
);
assert_eq!(urls.len(), 2);
assert!(urls
.iter()
.any(|url| url.as_str() == "https://example.test/sitemap.xml"));
}
#[tokio::test]
async fn fetch_with_limit_rejects_advertised_large_body() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/robots.txt"))
.respond_with(ResponseTemplate::new(404))
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/large"))
.respond_with(ResponseTemplate::new(200).set_body_bytes(vec![b'a'; 32]))
.mount(&server)
.await;
let config = Config {
retry_count: 0,
..Config::default()
};
let fetcher = HttpFetcher::new(&config, 0, &[], None).unwrap();
let err = fetcher
.fetch_with_limit(
Url::parse(&server.uri()).unwrap().join("/large").unwrap(),
8,
)
.await
.unwrap_err();
assert!(matches!(err, SiteforgeError::SizeLimitExceeded { .. }));
}
#[tokio::test]
async fn fetch_with_limit_rejects_streamed_large_body() {
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
thread::spawn(move || {
for _ in 0..2 {
let Ok((mut stream, _)) = listener.accept() else {
return;
};
let mut buffer = [0u8; 1024];
let read = stream.read(&mut buffer).unwrap_or(0);
let request = String::from_utf8_lossy(&buffer[..read]);
let path = request
.lines()
.next()
.and_then(|line| line.split_whitespace().nth(1))
.unwrap_or("/");
if path == "/robots.txt" {
let _ =
stream.write_all(b"HTTP/1.1 404 Not Found\r\nConnection: close\r\n\r\n");
} else {
let _ = stream.write_all(
b"HTTP/1.1 200 OK\r\nContent-Type: text/html\r\nConnection: close\r\n\r\nabcdefghijklmnop",
);
}
}
});
let config = Config {
retry_count: 0,
..Config::default()
};
let fetcher = HttpFetcher::new(&config, 0, &[], None).unwrap();
let err = fetcher
.fetch_with_limit(Url::parse(&format!("http://{addr}/large")).unwrap(), 8)
.await
.unwrap_err();
assert!(matches!(err, SiteforgeError::SizeLimitExceeded { .. }));
}
}