use quick_xml::Reader;
use quick_xml::XmlVersion;
use quick_xml::events::Event;
use url::Url;
use crate::http::http_fetch;
use crate::map::MapFilter;
use crate::normalize::{resolve_redirect, rewrite_url_host};
use crate::types::{CrawlConfig, SitemapUrl};
pub fn parse_sitemap_xml(body: &str) -> Vec<SitemapUrl> {
let mut urls = Vec::new();
let mut reader = Reader::from_str(body);
let mut buf = Vec::new();
let mut in_url = false;
let mut in_loc = false;
let mut in_lastmod = false;
let mut in_changefreq = false;
let mut in_priority = false;
let mut current_loc = String::new();
let mut current_lastmod: Option<String> = None;
let mut current_changefreq: Option<String> = None;
let mut current_priority: Option<String> = None;
loop {
match reader.read_event_into(&mut buf) {
Ok(Event::Start(ref e)) | Ok(Event::Empty(ref e)) => match e.name().as_ref() {
b"url" => {
in_url = true;
current_loc.clear();
current_lastmod = None;
current_changefreq = None;
current_priority = None;
}
b"loc" if in_url => in_loc = true,
b"lastmod" if in_url => in_lastmod = true,
b"changefreq" if in_url => in_changefreq = true,
b"priority" if in_url => in_priority = true,
_ => {}
},
Ok(Event::End(ref e)) => match e.name().as_ref() {
b"url" => {
if in_url && !current_loc.is_empty() {
urls.push(SitemapUrl {
url: current_loc.clone(),
lastmod: current_lastmod.clone(),
changefreq: current_changefreq.clone(),
priority: current_priority.clone(),
});
}
in_url = false;
}
b"loc" => in_loc = false,
b"lastmod" => in_lastmod = false,
b"changefreq" => in_changefreq = false,
b"priority" => in_priority = false,
_ => {}
},
Ok(Event::Text(ref e)) => {
if let Ok(text) = e.xml_content(XmlVersion::default()) {
let text = text.trim().to_owned();
if in_loc {
current_loc = text;
} else if in_lastmod {
current_lastmod = Some(text);
} else if in_changefreq {
current_changefreq = Some(text);
} else if in_priority {
current_priority = Some(text);
}
}
}
Ok(Event::Eof) => break,
Err(_) => break,
_ => {}
}
buf.clear();
}
urls
}
pub fn parse_sitemap_index(body: &str) -> Vec<String> {
let mut child_urls = Vec::new();
let mut reader = Reader::from_str(body);
let mut buf = Vec::new();
let mut in_sitemap = false;
let mut in_loc = false;
let mut current_loc = String::new();
loop {
match reader.read_event_into(&mut buf) {
Ok(Event::Start(ref e)) => match e.name().as_ref() {
b"sitemap" => {
in_sitemap = true;
current_loc.clear();
}
b"loc" if in_sitemap => in_loc = true,
_ => {}
},
Ok(Event::End(ref e)) => match e.name().as_ref() {
b"sitemap" => {
if in_sitemap && !current_loc.is_empty() {
child_urls.push(current_loc.clone());
}
in_sitemap = false;
}
b"loc" => in_loc = false,
_ => {}
},
Ok(Event::Text(ref e)) => {
if in_loc && let Ok(text) = e.xml_content(XmlVersion::default()) {
current_loc = text.trim().to_owned();
}
}
Ok(Event::Eof) => break,
Err(_) => break,
_ => {}
}
buf.clear();
}
child_urls
}
pub fn is_sitemap_index(body: &str) -> bool {
body.contains("<sitemapindex") || body.contains("<sitemapindex>")
}
pub(crate) async fn fetch_sitemap_tree(
sitemap_url: &str,
config: &CrawlConfig,
client: &reqwest::Client,
filter: &MapFilter,
limit: Option<usize>,
) -> Vec<SitemapUrl> {
let resp = match http_fetch(sitemap_url, config, &std::collections::HashMap::new(), client).await {
Ok(r) => r,
Err(_) => return Vec::new(),
};
process_sitemap_response(
sitemap_url,
&resp.body,
&resp.body_bytes,
&resp.content_type,
config,
client,
filter,
limit,
)
.await
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn process_sitemap_response(
sitemap_url: &str,
body: &str,
body_bytes: &[u8],
content_type: &str,
config: &CrawlConfig,
client: &reqwest::Client,
filter: &MapFilter,
limit: Option<usize>,
) -> Vec<SitemapUrl> {
let decompressed;
let xml_body = if content_type.contains("gzip") || content_type.contains("x-gzip") {
match decompress_gzip(body_bytes) {
Ok(d) => {
decompressed = d;
&decompressed
}
Err(_) => body,
}
} else {
body
};
let reached_limit = |len: usize| limit.is_some_and(|limit| len >= limit);
if is_sitemap_index(xml_body) {
let child_urls = parse_sitemap_index(xml_body);
let base = Url::parse(sitemap_url).ok();
let mut all_urls = Vec::new();
let max_children = 100;
for child_url in child_urls.iter().take(max_children) {
if reached_limit(all_urls.len()) {
break;
}
let resolved = if let Some(ref base_parsed) = base {
if Url::parse(child_url).is_ok() {
rewrite_url_host(child_url, base_parsed)
} else {
resolve_redirect(sitemap_url, child_url)
}
} else {
child_url.clone()
};
let child_resp = match http_fetch(&resolved, config, &std::collections::HashMap::new(), client).await {
Ok(r) => r,
Err(_) => continue,
};
let child_body = &child_resp.body;
let child_decompressed;
let child_xml = if child_resp.content_type.contains("gzip") {
match decompress_gzip(&child_resp.body_bytes) {
Ok(d) => {
child_decompressed = d;
&child_decompressed
}
Err(_) => child_body,
}
} else {
child_body
};
for entry in parse_sitemap_xml(child_xml) {
if !filter.matches(&entry.url) {
continue;
}
all_urls.push(entry);
if reached_limit(all_urls.len()) {
break;
}
}
}
all_urls
} else {
let mut urls = Vec::new();
for entry in parse_sitemap_xml(xml_body) {
if !filter.matches(&entry.url) {
continue;
}
urls.push(entry);
if reached_limit(urls.len()) {
break;
}
}
urls
}
}
pub(crate) fn decompress_gzip(data: &[u8]) -> Result<String, std::io::Error> {
use flate2::read::GzDecoder;
use std::io::Read;
const MAX_DECOMPRESSED_SIZE: u64 = 50 * 1024 * 1024;
let decoder = GzDecoder::new(data);
let mut limited = decoder.take(MAX_DECOMPRESSED_SIZE);
let mut result = String::new();
limited.read_to_string(&mut result)?;
Ok(result)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::map::MapFilter;
use crate::types::CrawlConfig;
fn urlset(count: usize) -> String {
let mut body =
String::from(r#"<?xml version="1.0"?><urlset xmlns="http://www.sitemaps.org/schemas/sitemap/0.9">"#);
for i in 0..count {
body.push_str(&format!("<url><loc>https://example.com/page-{i}</loc></url>"));
}
body.push_str("</urlset>");
body
}
#[tokio::test]
async fn process_sitemap_response_stops_at_limit_on_single_sitemap() {
let config = CrawlConfig::default();
let filter = MapFilter::from_config(&config).unwrap();
let client = reqwest::Client::new();
let body = urlset(1000);
let urls = process_sitemap_response(
"https://example.com/sitemap.xml",
&body,
body.as_bytes(),
"application/xml",
&config,
&client,
&filter,
Some(10),
)
.await;
assert_eq!(
urls.len(),
10,
"map_limit must cap the parsed URLs, not just the returned slice"
);
}
#[tokio::test]
async fn process_sitemap_response_returns_all_when_unlimited() {
let config = CrawlConfig::default();
let filter = MapFilter::from_config(&config).unwrap();
let client = reqwest::Client::new();
let body = urlset(25);
let urls = process_sitemap_response(
"https://example.com/sitemap.xml",
&body,
body.as_bytes(),
"application/xml",
&config,
&client,
&filter,
None,
)
.await;
assert_eq!(urls.len(), 25);
}
#[tokio::test]
async fn process_sitemap_response_applies_filter_before_limit() {
let config = CrawlConfig {
map_search: Some("keep".to_string()),
..CrawlConfig::default()
};
let filter = MapFilter::from_config(&config).unwrap();
let client = reqwest::Client::new();
let body = concat!(
r#"<?xml version="1.0"?><urlset xmlns="http://www.sitemaps.org/schemas/sitemap/0.9">"#,
"<url><loc>https://example.com/keep-1</loc></url>",
"<url><loc>https://example.com/drop-1</loc></url>",
"<url><loc>https://example.com/keep-2</loc></url>",
"<url><loc>https://example.com/drop-2</loc></url>",
"</urlset>",
);
let urls = process_sitemap_response(
"https://example.com/sitemap.xml",
body,
body.as_bytes(),
"application/xml",
&config,
&client,
&filter,
None,
)
.await;
assert_eq!(urls.len(), 2);
assert!(urls.iter().all(|entry| entry.url.contains("keep")));
}
}