1use std::collections::HashMap;
4use std::fmt;
5use std::sync::Arc;
6use std::time::Duration;
7
8use serde_json::Value;
9use servo::accesskit::{Node, NodeId};
10
11use crate::error::Error;
12use crate::net::sanitize_user_agent;
13
14#[derive(Debug, Clone, Default, serde::Serialize)]
16#[non_exhaustive]
17pub struct Page {
18 pub html: String,
20 pub inner_text: String,
22 pub title: Option<String>,
24 #[serde(skip_serializing_if = "Option::is_none")]
26 pub layout_json: Option<String>,
27 #[serde(skip)]
29 visibility_json: Option<String>,
30 #[serde(skip_serializing_if = "Option::is_none")]
32 pub js_result: Option<String>,
33 pub console_messages: Vec<ConsoleMessage>,
35 #[serde(skip_serializing_if = "Option::is_none")]
37 pub accessibility_tree: Option<String>,
38 #[serde(skip_serializing_if = "Option::is_none")]
40 pub extracted: Option<Value>,
41 #[serde(skip)]
43 screenshot_png: Option<Vec<u8>>,
44 #[serde(skip)]
46 a11y: Option<Arc<HashMap<NodeId, Node>>>,
47 #[serde(skip)]
49 visibility_policy: crate::visibility::VisibilityPolicy,
50}
51
52impl Page {
53 pub fn markdown(&self) -> crate::error::Result<String> {
55 self.markdown_with_url("")
56 }
57
58 pub fn markdown_with_url(&self, url: &str) -> crate::error::Result<String> {
60 Ok(crate::extract::extract_text(&self.extract_input(url, None))?)
61 }
62
63 pub fn extract_json(&self) -> crate::error::Result<String> {
65 self.extract_json_with_url("")
66 }
67
68 pub fn extract_json_with_url(&self, url: &str) -> crate::error::Result<String> {
70 Ok(crate::extract::extract_json(&self.extract_input(url, None))?)
71 }
72
73 pub fn markdown_with_selector(&self, url: &str, selector: &str) -> crate::error::Result<String> {
75 Ok(crate::extract::extract_text(&self.extract_input(url, Some(selector)))?)
76 }
77
78 pub fn extract_json_with_selector(&self, url: &str, selector: &str) -> crate::error::Result<String> {
80 Ok(crate::extract::extract_json(&self.extract_input(url, Some(selector)))?)
81 }
82
83 #[must_use]
85 pub fn screenshot_png(&self) -> Option<&[u8]> {
86 self.screenshot_png.as_deref()
87 }
88
89 fn extract_input<'a>(&'a self, url: &'a str, selector: Option<&'a str>) -> crate::extract::ExtractInput<'a> {
90 crate::extract::ExtractInput::new(&self.html, url)
91 .with_layout_json(self.layout_json.as_deref())
92 .with_visibility_json(self.visibility_json.as_deref())
93 .with_a11y(self.a11y.as_deref())
94 .with_inner_text(Some(&self.inner_text))
95 .with_selector(selector)
96 .with_visibility(self.visibility_policy)
97 }
98
99 pub(crate) fn from_servo(page: crate::bridge::ServoPage) -> Self {
100 let title = {
101 let doc = dom_query::Document::from(page.html.as_str());
102 let t = doc.select("title").text().to_string();
103 if t.is_empty() { None } else { Some(t) }
104 };
105 let screenshot_png = page.screenshot.and_then(|img| {
106 let mut buf = std::io::Cursor::new(Vec::new());
107 img.write_to(&mut buf, image::ImageFormat::Png).ok()?;
108 Some(buf.into_inner())
109 });
110 Self {
111 html: page.html,
112 inner_text: page.inner_text.unwrap_or_default(),
113 title,
114 layout_json: page.layout_json,
115 visibility_json: page.visibility_json,
116 js_result: page.js_result,
117 console_messages: page
118 .console_messages
119 .into_iter()
120 .map(|m| ConsoleMessage {
121 level: match m.level {
122 crate::bridge::ConsoleLevel::Log => ConsoleLevel::Log,
123 crate::bridge::ConsoleLevel::Debug => ConsoleLevel::Debug,
124 crate::bridge::ConsoleLevel::Info => ConsoleLevel::Info,
125 crate::bridge::ConsoleLevel::Warn => ConsoleLevel::Warn,
126 crate::bridge::ConsoleLevel::Error => ConsoleLevel::Error,
127 crate::bridge::ConsoleLevel::Trace => ConsoleLevel::Trace,
128 },
129 message: m.message,
130 })
131 .collect(),
132 screenshot_png,
133 accessibility_tree: page.accessibility_tree,
134 a11y: page.a11y.map(Arc::new),
135 extracted: None,
136 visibility_policy: crate::visibility::VisibilityPolicy::default(),
137 }
138 }
139}
140
141#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
143#[non_exhaustive]
144pub struct ConsoleMessage {
145 pub level: ConsoleLevel,
147 pub message: String,
149}
150
151#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)]
153#[serde(rename_all = "lowercase")]
154#[non_exhaustive]
155pub enum ConsoleLevel {
156 Log,
158 Debug,
160 Info,
162 Warn,
164 Error,
166 Trace,
168}
169
170impl ConsoleLevel {
171 #[must_use]
173 pub fn as_str(&self) -> &'static str {
174 match self {
175 Self::Log => "log",
176 Self::Debug => "debug",
177 Self::Info => "info",
178 Self::Warn => "warn",
179 Self::Error => "error",
180 Self::Trace => "trace",
181 }
182 }
183}
184
185impl fmt::Display for ConsoleLevel {
186 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
187 f.pad(self.as_str())
188 }
189}
190
191#[derive(Debug, Clone, Default)]
192pub(crate) enum FetchMode {
193 #[default]
194 Content,
195 Screenshot {
196 full_page: bool,
197 },
198 JavaScript(String),
199}
200
201#[must_use = "options do nothing until passed to fetch()"]
203#[derive(Debug, Clone)]
204pub struct FetchOptions {
205 pub(crate) url: String,
206 pub(crate) timeout: Duration,
207 pub(crate) settle: Duration,
208 pub(crate) mode: FetchMode,
209 pub(crate) user_agent: Option<String>,
210 pub(crate) extract_schema: Option<crate::schema::ExtractSchema>,
211 pub(crate) visibility: crate::visibility::VisibilityPolicy,
212}
213
214impl FetchOptions {
215 pub fn new(url: &str) -> Self {
217 Self {
218 url: url.into(),
219 timeout: Duration::from_secs(30),
220 settle: Duration::ZERO,
221 mode: FetchMode::Content,
222 user_agent: None,
223 extract_schema: None,
224 visibility: crate::visibility::VisibilityPolicy::default(),
225 }
226 }
227
228 pub fn screenshot(url: &str, full_page: bool) -> Self {
230 Self {
231 mode: FetchMode::Screenshot { full_page },
232 ..Self::new(url)
233 }
234 }
235
236 pub fn javascript(url: &str, expression: impl Into<String>) -> Self {
238 Self {
239 mode: FetchMode::JavaScript(expression.into()),
240 ..Self::new(url)
241 }
242 }
243
244 pub fn timeout(mut self, timeout: Duration) -> Self {
246 self.timeout = timeout;
247 self
248 }
249
250 pub fn settle(mut self, settle: Duration) -> Self {
252 self.settle = settle;
253 self
254 }
255
256 pub fn user_agent(mut self, ua: impl Into<String>) -> Self {
258 self.user_agent = Some(sanitize_user_agent(ua.into()));
259 self
260 }
261
262 pub fn schema(mut self, schema: crate::schema::ExtractSchema) -> Self {
264 self.extract_schema = Some(schema);
265 self
266 }
267
268 pub fn visibility(mut self, policy: crate::visibility::VisibilityPolicy) -> Self {
270 self.visibility = policy;
271 self
272 }
273}
274
275#[allow(clippy::needless_pass_by_value)]
277pub fn fetch(opts: FetchOptions) -> crate::error::Result<Page> {
278 crate::net::ensure_crypto_provider();
279
280 crate::net::validate_url(&opts.url)?;
281
282 if matches!(opts.mode, FetchMode::Content)
283 && let Some(bytes) = crate::pdf::probe(&opts.url, opts.timeout.as_secs().max(1))
284 {
285 let text = crate::extract::extract_pdf(&bytes);
286 return Ok(Page {
287 html: String::new(),
288 inner_text: text,
289 ..Page::default()
290 });
291 }
292
293 let bridge_opts = crate::bridge::FetchOptions {
294 url: &opts.url,
295 timeout_secs: opts.timeout.as_secs().max(1),
296 settle_ms: u64::try_from(opts.settle.as_millis()).unwrap_or(u64::MAX),
297 user_agent: opts.user_agent.as_deref(),
298 mode: match opts.mode {
299 FetchMode::Content => crate::bridge::FetchMode::Content { include_a11y: false },
300 FetchMode::Screenshot { full_page } => crate::bridge::FetchMode::Screenshot { full_page },
301 FetchMode::JavaScript(ref expr) => crate::bridge::FetchMode::ExecuteJs {
302 expression: expr.clone(),
303 },
304 },
305 };
306
307 let servo_page = crate::bridge::fetch_page(bridge_opts).map_err(|e| {
308 let msg = format!("{e:#}");
309 if msg.contains("timed out") {
310 Error::Timeout {
311 url: opts.url.clone(),
312 timeout: opts.timeout,
313 }
314 } else {
315 Error::Engine(msg)
316 }
317 })?;
318
319 let mut page = Page::from_servo(servo_page);
320 page.visibility_policy = opts.visibility;
321 if let Some(schema) = opts.extract_schema.as_ref() {
322 page.extracted = Some(schema.extract_from(&page.html));
323 }
324 Ok(page)
325}
326
327pub fn markdown(url: &str) -> crate::error::Result<String> {
329 fetch(FetchOptions::new(url))?.markdown_with_url(url)
330}
331
332pub fn extract_json(url: &str) -> crate::error::Result<String> {
334 fetch(FetchOptions::new(url))?.extract_json_with_url(url)
335}
336
337pub fn text(url: &str) -> crate::error::Result<String> {
339 Ok(fetch(FetchOptions::new(url))?.inner_text)
340}
341
342#[cfg(test)]
343mod tests {
344 use super::*;
345
346 #[test]
347 fn fetch_options_defaults() {
348 let opts = FetchOptions::new("https://example.com");
349 assert_eq!(opts.url, "https://example.com");
350 assert_eq!(opts.timeout, Duration::from_secs(30));
351 assert_eq!(opts.settle, Duration::ZERO);
352 assert!(matches!(opts.mode, FetchMode::Content));
353 }
354
355 #[test]
356 fn fetch_options_screenshot() {
357 let opts = FetchOptions::screenshot("https://example.com", true);
358 assert!(matches!(opts.mode, FetchMode::Screenshot { full_page: true }));
359 }
360
361 #[test]
362 fn fetch_options_javascript() {
363 let opts = FetchOptions::javascript("https://example.com", "document.title");
364 assert!(matches!(opts.mode, FetchMode::JavaScript(ref e) if e == "document.title"));
365 }
366
367 #[test]
368 fn fetch_options_chaining() {
369 let opts = FetchOptions::new("https://example.com")
370 .timeout(Duration::from_secs(60))
371 .settle(Duration::from_millis(500));
372 assert_eq!(opts.timeout, Duration::from_secs(60));
373 assert_eq!(opts.settle, Duration::from_millis(500));
374 }
375
376 #[test]
377 fn fetch_user_agent_set() {
378 let opts = FetchOptions::new("https://example.com").user_agent("MyBot/1.0");
379 assert_eq!(opts.user_agent.as_deref(), Some("MyBot/1.0"));
380 }
381
382 #[test]
383 fn fetch_user_agent_default_is_none() {
384 let opts = FetchOptions::new("https://example.com");
385 assert!(opts.user_agent.is_none());
386 }
387
388 #[test]
389 fn fetch_user_agent_sanitizes_crlf() {
390 let opts = FetchOptions::new("https://example.com").user_agent("Bot\r\nX-Evil: yes");
391 assert_eq!(opts.user_agent.as_deref(), Some("Bot X-Evil: yes"));
392 }
393
394 #[test]
395 fn fetch_user_agent_sanitizes_null() {
396 let opts = FetchOptions::new("https://example.com").user_agent("Bot\0/1.0");
397 assert_eq!(opts.user_agent.as_deref(), Some("Bot /1.0"));
398 }
399
400 #[test]
401 fn fetch_user_agent_empty_string() {
402 let opts = FetchOptions::new("https://example.com").user_agent("");
403 assert_eq!(opts.user_agent.as_deref(), Some(""));
404 }
405
406 #[test]
407 fn page_markdown_from_html() {
408 let page = Page {
409 html: "<html><head><title>Test</title></head><body><p>hello world</p></body></html>".into(),
410 inner_text: "hello world".into(),
411 ..Page::default()
412 };
413 let md = page.markdown().unwrap();
414 assert!(md.contains("hello world"));
415 }
416
417 #[test]
418 fn page_extract_json_produces_valid_json() {
419 let page = Page {
420 html: "<html><head><title>Test</title></head><body><p>content</p></body></html>".into(),
421 inner_text: "content".into(),
422 ..Page::default()
423 };
424 let json = page.extract_json().unwrap();
425 let _: Value = serde_json::from_str(&json).expect("valid JSON");
426 }
427
428 #[test]
429 fn page_screenshot_png_none_by_default() {
430 let page = Page::default();
431 assert!(page.screenshot_png().is_none());
432 }
433
434 #[test]
435 fn page_markdown_with_selector_scopes_to_subtree() {
436 let page = Page {
437 html: "<html><body><article>keep</article><aside>drop</aside></body></html>".into(),
438 ..Page::default()
439 };
440 let md = page.markdown_with_selector("https://example.com", "article").unwrap();
441 assert!(md.contains("keep"));
442 assert!(!md.contains("drop"));
443 }
444
445 #[test]
446 fn page_extract_json_with_selector_includes_url() {
447 let page = Page {
448 html: "<html><body><article>scoped</article></body></html>".into(),
449 ..Page::default()
450 };
451 let json = page
452 .extract_json_with_selector("https://example.com/page", "article")
453 .unwrap();
454 let parsed: Value = serde_json::from_str(&json).expect("valid JSON");
455 assert_eq!(parsed["url"].as_str(), Some("https://example.com/page"));
456 assert!(parsed["text_content"].as_str().unwrap().contains("scoped"));
457 }
458
459 #[test]
460 fn page_markdown_with_selector_no_match_returns_empty() {
461 let page = Page {
462 html: "<html><body><article>x</article></body></html>".into(),
463 ..Page::default()
464 };
465 let md = page.markdown_with_selector("", ".nonexistent").unwrap();
466 assert!(md.is_empty());
467 }
468
469 #[test]
470 fn page_markdown_with_invalid_selector_returns_error() {
471 let page = Page {
472 html: "<html><body><p>x</p></body></html>".into(),
473 ..Page::default()
474 };
475 let err = page.markdown_with_selector("", "###invalid[[[").unwrap_err();
476 assert!(err.to_string().contains("invalid CSS selector"));
477 }
478
479 #[test]
480 fn page_markdown_with_empty_selector_returns_error() {
481 let page = Page {
482 html: "<html><body><p>x</p></body></html>".into(),
483 ..Page::default()
484 };
485 assert!(page.markdown_with_selector("", "").is_err());
486 }
487
488 #[test]
489 fn fetch_rejects_invalid_url() {
490 let result = fetch(FetchOptions::new("not a url"));
491 assert!(result.is_err());
492 let err = result.unwrap_err();
493 assert!(matches!(err, Error::InvalidUrl { .. }));
494 }
495
496 #[test]
497 fn fetch_rejects_private_ip() {
498 let result = fetch(FetchOptions::new("http://127.0.0.1/"));
499 assert!(result.is_err());
500 }
501
502 #[test]
503 fn fetch_rejects_file_scheme() {
504 let result = fetch(FetchOptions::new("file:///etc/passwd"));
505 assert!(result.is_err());
506 }
507
508 mod page_from_servo {
509 use crate::bridge;
510 use crate::fetch::{ConsoleLevel, Page};
511
512 fn synthetic_image(w: u32, h: u32) -> image::RgbaImage {
513 image::RgbaImage::from_pixel(w, h, image::Rgba([255, 0, 0, 255]))
514 }
515
516 fn empty_servo_page() -> bridge::ServoPage {
517 bridge::ServoPage::default()
518 }
519
520 #[test]
521 fn extracts_title_from_html() {
522 let mut sp = empty_servo_page();
523 sp.html = "<html><head><title>Hello World</title></head></html>".into();
524 let page = Page::from_servo(sp);
525 assert_eq!(page.title.as_deref(), Some("Hello World"));
526 }
527
528 #[test]
529 fn title_is_none_when_tag_missing() {
530 let mut sp = empty_servo_page();
531 sp.html = "<html><body>no title here</body></html>".into();
532 let page = Page::from_servo(sp);
533 assert!(page.title.is_none());
534 }
535
536 #[test]
537 fn title_is_none_when_tag_empty() {
538 let mut sp = empty_servo_page();
539 sp.html = "<html><head><title></title></head></html>".into();
540 let page = Page::from_servo(sp);
541 assert!(page.title.is_none());
542 }
543
544 #[test]
545 fn title_is_none_for_empty_html() {
546 let page = Page::from_servo(empty_servo_page());
547 assert!(page.title.is_none());
548 }
549
550 #[test]
551 fn inner_text_none_becomes_empty_string() {
552 let sp = empty_servo_page();
553 assert!(sp.inner_text.is_none());
554 let page = Page::from_servo(sp);
555 assert_eq!(page.inner_text, "");
556 }
557
558 #[test]
559 fn screenshot_is_encoded_as_png() {
560 let mut sp = empty_servo_page();
561 sp.screenshot = Some(synthetic_image(8, 8));
562 let page = Page::from_servo(sp);
563 let bytes = page.screenshot_png().expect("screenshot encoded");
564 assert_eq!(&bytes[..8], b"\x89PNG\r\n\x1a\n", "PNG magic bytes");
565 }
566
567 #[test]
568 fn console_messages_empty_by_default() {
569 let page = Page::from_servo(empty_servo_page());
570 assert!(page.console_messages.is_empty());
571 }
572
573 #[test]
574 fn console_messages_preserve_all_six_levels() {
575 let cases = [
576 (bridge::ConsoleLevel::Log, ConsoleLevel::Log),
577 (bridge::ConsoleLevel::Debug, ConsoleLevel::Debug),
578 (bridge::ConsoleLevel::Info, ConsoleLevel::Info),
579 (bridge::ConsoleLevel::Warn, ConsoleLevel::Warn),
580 (bridge::ConsoleLevel::Error, ConsoleLevel::Error),
581 (bridge::ConsoleLevel::Trace, ConsoleLevel::Trace),
582 ];
583 for (src, expected) in cases {
584 let mut sp = empty_servo_page();
585 sp.console_messages = vec![bridge::ConsoleMessage {
586 level: src,
587 message: "msg".into(),
588 }];
589 let page = Page::from_servo(sp);
590 assert_eq!(
591 page.console_messages.len(),
592 1,
593 "console message lost for source level {src:?}",
594 );
595 assert_eq!(
596 page.console_messages[0].level, expected,
597 "level mapping wrong for source {src:?}",
598 );
599 }
600 }
601
602 #[test]
603 fn console_messages_preserve_ordering_across_levels() {
604 let mut sp = empty_servo_page();
605 sp.console_messages = vec![
606 bridge::ConsoleMessage {
607 level: bridge::ConsoleLevel::Info,
608 message: "first".into(),
609 },
610 bridge::ConsoleMessage {
611 level: bridge::ConsoleLevel::Error,
612 message: "second".into(),
613 },
614 bridge::ConsoleMessage {
615 level: bridge::ConsoleLevel::Warn,
616 message: "third".into(),
617 },
618 ];
619 let page = Page::from_servo(sp);
620 assert_eq!(page.console_messages.len(), 3);
621 assert_eq!(page.console_messages[0].message, "first");
622 assert_eq!(page.console_messages[1].message, "second");
623 assert_eq!(page.console_messages[2].message, "third");
624 assert_eq!(page.console_messages[0].level, ConsoleLevel::Info);
625 assert_eq!(page.console_messages[1].level, ConsoleLevel::Error);
626 assert_eq!(page.console_messages[2].level, ConsoleLevel::Warn);
627 }
628
629 #[test]
630 fn extracted_starts_as_none_until_schema_applied() {
631 let page = Page::from_servo(empty_servo_page());
632 assert!(page.extracted.is_none());
633 }
634
635 #[test]
636 fn full_round_trip_preserves_every_field() {
637 let sp = bridge::ServoPage {
638 html: "<html><head><title>T</title></head><body>B</body></html>".into(),
639 inner_text: Some("B".into()),
640 layout_json: Some("[]".into()),
641 visibility_json: Some("[]".into()),
642 screenshot: Some(synthetic_image(2, 2)),
643 js_result: Some("42".into()),
644 accessibility_tree: Some("{}".into()),
645 a11y: None,
646 console_messages: vec![bridge::ConsoleMessage {
647 level: bridge::ConsoleLevel::Log,
648 message: "x".into(),
649 }],
650 };
651 let page = Page::from_servo(sp);
652 assert_eq!(page.html, "<html><head><title>T</title></head><body>B</body></html>");
653 assert_eq!(page.inner_text, "B");
654 assert_eq!(page.title.as_deref(), Some("T"));
655 assert_eq!(page.layout_json.as_deref(), Some("[]"));
656 assert_eq!(page.js_result.as_deref(), Some("42"));
657 assert_eq!(page.accessibility_tree.as_deref(), Some("{}"));
658 assert_eq!(page.console_messages.len(), 1);
659 assert!(page.screenshot_png().is_some());
660 assert!(page.extracted.is_none());
661 }
662 }
663}