crawlkit_engine/
sitemap.rs1use std::sync::Arc;
8
9use dashmap::DashMap;
10use quick_xml::events::Event;
11use quick_xml::Reader;
12
13use crate::analyzers::SitemapEntry;
14use crate::http::HttpClient;
15
16pub struct SitemapCache {
20 cache: DashMap<String, Vec<SitemapEntry>>,
22 client: Arc<HttpClient>,
24}
25
26impl SitemapCache {
27 pub fn new(client: Arc<HttpClient>) -> Self {
29 Self {
30 cache: DashMap::new(),
31 client,
32 }
33 }
34
35 pub async fn fetch(&self, sitemap_url: &str) -> Vec<SitemapEntry> {
41 self.fetch_inner(sitemap_url, 0).await
42 }
43
44 async fn fetch_inner(&self, sitemap_url: &str, depth: usize) -> Vec<SitemapEntry> {
45 if depth > 2 {
46 return Vec::new();
47 }
48
49 if let Some(entries) = self.cache.get(sitemap_url) {
51 return entries.clone();
52 }
53
54 let url = match url::Url::parse(sitemap_url) {
56 Ok(u) => u,
57 Err(_) => return Vec::new(),
58 };
59
60 let body = match self.client.fetch(&url).await {
61 Ok(result) if result.status_code == 200 => result.body,
62 _ => return Vec::new(),
63 };
64
65 let entries = match parse_sitemap(&body) {
67 Ok(SitemapContent::UrlSet(entries)) => entries,
68 Ok(SitemapContent::SitemapIndex(child_urls)) => {
69 let mut all_entries = Vec::new();
71 for child_url in &child_urls {
72 let child_entries = Box::pin(self.fetch_inner(child_url, depth + 1)).await;
73 all_entries.extend(child_entries);
74 }
75 all_entries
76 }
77 Err(_) => return Vec::new(),
78 };
79
80 self.cache.insert(sitemap_url.to_string(), entries.clone());
81 entries
82 }
83
84 pub async fn fetch_all(&self, urls: &[String]) -> Vec<SitemapEntry> {
86 let mut all = Vec::new();
87 for url in urls {
88 all.extend(self.fetch(url).await);
89 }
90 all
91 }
92
93 pub fn cached_count(&self) -> usize {
95 self.cache.iter().map(|e| e.value().len()).sum()
96 }
97}
98
99enum SitemapContent {
101 UrlSet(Vec<SitemapEntry>),
102 SitemapIndex(Vec<String>),
103}
104
105fn parse_sitemap(xml: &str) -> Result<SitemapContent, SitemapError> {
110 let mut reader = Reader::from_str(xml);
111 reader.config_mut().trim_text(true);
112
113 let mut buf = Vec::new();
114 let mut root_tag = None;
115
116 loop {
118 match reader.read_event_into(&mut buf) {
119 Ok(Event::Start(e)) => {
120 root_tag = Some(String::from_utf8_lossy(e.name().as_ref()).to_string());
121 break;
122 }
123 Ok(Event::Eof) => break,
124 Ok(_) => continue,
125 Err(_) => return Err(SitemapError::InvalidXml),
126 }
127 }
128
129 let root = root_tag.ok_or(SitemapError::InvalidXml)?;
130
131 match root.as_str() {
132 "urlset" => parse_urlset(reader, &mut buf),
133 "sitemapindex" => parse_sitemapindex(reader, &mut buf),
134 _ => Err(SitemapError::UnknownRoot),
135 }
136}
137
138fn parse_urlset(
140 mut reader: Reader<&[u8]>,
141 buf: &mut Vec<u8>,
142) -> Result<SitemapContent, SitemapError> {
143 let mut entries = Vec::new();
144 let mut current_url: Option<String> = None;
145 let mut current_lastmod: Option<String> = None;
146 let mut current_changefreq: Option<String> = None;
147 let mut current_priority: Option<f64> = None;
148 let mut in_url = false;
149 let mut current_tag = String::new();
150
151 loop {
152 buf.clear();
153 match reader.read_event_into(buf) {
154 Ok(Event::Start(e)) => {
155 current_tag = String::from_utf8_lossy(e.name().as_ref()).to_string();
156 if current_tag == "url" {
157 in_url = true;
158 current_url = None;
159 current_lastmod = None;
160 current_changefreq = None;
161 current_priority = None;
162 }
163 }
164 Ok(Event::Text(e)) => {
165 let text = String::from_utf8_lossy(e.as_ref()).to_string();
166 if in_url {
167 match current_tag.as_str() {
168 "loc" => current_url = Some(text),
169 "lastmod" => current_lastmod = Some(text),
170 "changefreq" => current_changefreq = Some(text),
171 "priority" => {
172 current_priority = text.parse().ok();
173 }
174 _ => {}
175 }
176 }
177 }
178 Ok(Event::End(e)) => {
179 let name = String::from_utf8_lossy(e.name().as_ref()).to_string();
180 if name == "url" && in_url {
181 if let Some(loc) = current_url.take() {
182 entries.push(SitemapEntry {
183 url: loc,
184 lastmod: current_lastmod.take(),
185 changefreq: current_changefreq.take(),
186 priority: current_priority.take(),
187 });
188 }
189 in_url = false;
190 }
191 current_tag.clear();
192 }
193 Ok(Event::Eof) => break,
194 Err(_) => break,
195 _ => {}
196 }
197 }
198
199 Ok(SitemapContent::UrlSet(entries))
200}
201
202fn parse_sitemapindex(
204 mut reader: Reader<&[u8]>,
205 buf: &mut Vec<u8>,
206) -> Result<SitemapContent, SitemapError> {
207 let mut child_urls = Vec::new();
208 let mut current_tag = String::new();
209
210 loop {
211 buf.clear();
212 match reader.read_event_into(buf) {
213 Ok(Event::Start(e)) => {
214 current_tag = String::from_utf8_lossy(e.name().as_ref()).to_string();
215 }
216 Ok(Event::Text(e)) => {
217 if current_tag == "loc" {
218 child_urls.push(String::from_utf8_lossy(e.as_ref()).to_string());
219 }
220 }
221 Ok(Event::End(_)) => {
222 current_tag.clear();
223 }
224 Ok(Event::Eof) => break,
225 Err(_) => break,
226 _ => {}
227 }
228 }
229
230 Ok(SitemapContent::SitemapIndex(child_urls))
231}
232
233#[derive(Debug)]
235pub enum SitemapError {
236 InvalidXml,
237 UnknownRoot,
238}
239
240#[cfg(test)]
241#[allow(clippy::unwrap_used)]
242mod tests {
243 use super::*;
244
245 #[test]
246 fn test_parse_urlset() {
247 let xml = r#"<?xml version="1.0" encoding="UTF-8"?>
248<urlset xmlns="http://www.sitemaps.org/schemas/sitemap/0.9">
249 <url>
250 <loc>https://example.com/</loc>
251 <lastmod>2024-01-15</lastmod>
252 <changefreq>daily</changefreq>
253 <priority>1.0</priority>
254 </url>
255 <url>
256 <loc>https://example.com/about</loc>
257 <lastmod>2024-01-10</lastmod>
258 <changefreq>monthly</changefreq>
259 <priority>0.8</priority>
260 </url>
261</urlset>"#;
262
263 match parse_sitemap(xml).unwrap() {
264 SitemapContent::UrlSet(entries) => {
265 assert_eq!(entries.len(), 2);
266 assert_eq!(entries[0].url, "https://example.com/");
267 assert_eq!(entries[0].lastmod, Some("2024-01-15".to_string()));
268 assert_eq!(entries[0].changefreq, Some("daily".to_string()));
269 assert_eq!(entries[0].priority, Some(1.0));
270 assert_eq!(entries[1].url, "https://example.com/about");
271 }
272 _ => panic!("Expected UrlSet"),
273 }
274 }
275
276 #[test]
277 fn test_parse_sitemapindex() {
278 let xml = r#"<?xml version="1.0" encoding="UTF-8"?>
279<sitemapindex xmlns="http://www.sitemaps.org/schemas/sitemap/0.9">
280 <sitemap>
281 <loc>https://example.com/sitemap-pages.xml</loc>
282 </sitemap>
283 <sitemap>
284 <loc>https://example.com/sitemap-posts.xml</loc>
285 </sitemap>
286</sitemapindex>"#;
287
288 match parse_sitemap(xml).unwrap() {
289 SitemapContent::SitemapIndex(urls) => {
290 assert_eq!(urls.len(), 2);
291 assert_eq!(urls[0], "https://example.com/sitemap-pages.xml");
292 assert_eq!(urls[1], "https://example.com/sitemap-posts.xml");
293 }
294 _ => panic!("Expected SitemapIndex"),
295 }
296 }
297
298 #[test]
299 fn test_parse_empty_urlset() {
300 let xml = r#"<?xml version="1.0" encoding="UTF-8"?>
301<urlset xmlns="http://www.sitemaps.org/schemas/sitemap/0.9">
302</urlset>"#;
303
304 match parse_sitemap(xml).unwrap() {
305 SitemapContent::UrlSet(entries) => {
306 assert_eq!(entries.len(), 0);
307 }
308 _ => panic!("Expected UrlSet"),
309 }
310 }
311}