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 if format!("{e:#}").contains("timed out") {
309 Error::Timeout {
310 url: opts.url.clone(),
311 timeout: opts.timeout,
312 }
313 } else {
314 Error::engine(e, Some(opts.url.clone()))
315 }
316 })?;
317
318 let mut page = Page::from_servo(servo_page);
319 page.visibility_policy = opts.visibility;
320 if let Some(schema) = opts.extract_schema.as_ref() {
321 page.extracted = Some(schema.extract_from(&page.html));
322 }
323 Ok(page)
324}
325
326pub fn markdown(url: &str) -> crate::error::Result<String> {
328 fetch(FetchOptions::new(url))?.markdown_with_url(url)
329}
330
331pub fn extract_json(url: &str) -> crate::error::Result<String> {
333 fetch(FetchOptions::new(url))?.extract_json_with_url(url)
334}
335
336pub fn text(url: &str) -> crate::error::Result<String> {
338 Ok(fetch(FetchOptions::new(url))?.inner_text)
339}
340
341#[cfg(test)]
342mod tests {
343 use super::*;
344
345 #[test]
346 fn fetch_options_defaults() {
347 let opts = FetchOptions::new("https://example.com");
348 assert_eq!(opts.url, "https://example.com");
349 assert_eq!(opts.timeout, Duration::from_secs(30));
350 assert_eq!(opts.settle, Duration::ZERO);
351 assert!(matches!(opts.mode, FetchMode::Content));
352 }
353
354 #[test]
355 fn fetch_options_screenshot() {
356 let opts = FetchOptions::screenshot("https://example.com", true);
357 assert!(matches!(opts.mode, FetchMode::Screenshot { full_page: true }));
358 }
359
360 #[test]
361 fn fetch_options_javascript() {
362 let opts = FetchOptions::javascript("https://example.com", "document.title");
363 assert!(matches!(opts.mode, FetchMode::JavaScript(ref e) if e == "document.title"));
364 }
365
366 #[test]
367 fn fetch_options_chaining() {
368 let opts = FetchOptions::new("https://example.com")
369 .timeout(Duration::from_secs(60))
370 .settle(Duration::from_millis(500));
371 assert_eq!(opts.timeout, Duration::from_secs(60));
372 assert_eq!(opts.settle, Duration::from_millis(500));
373 }
374
375 #[test]
376 fn fetch_user_agent_set() {
377 let opts = FetchOptions::new("https://example.com").user_agent("MyBot/1.0");
378 assert_eq!(opts.user_agent.as_deref(), Some("MyBot/1.0"));
379 }
380
381 #[test]
382 fn fetch_user_agent_default_is_none() {
383 let opts = FetchOptions::new("https://example.com");
384 assert!(opts.user_agent.is_none());
385 }
386
387 #[test]
388 fn fetch_user_agent_sanitizes_crlf() {
389 let opts = FetchOptions::new("https://example.com").user_agent("Bot\r\nX-Evil: yes");
390 assert_eq!(opts.user_agent.as_deref(), Some("Bot X-Evil: yes"));
391 }
392
393 #[test]
394 fn fetch_user_agent_sanitizes_null() {
395 let opts = FetchOptions::new("https://example.com").user_agent("Bot\0/1.0");
396 assert_eq!(opts.user_agent.as_deref(), Some("Bot /1.0"));
397 }
398
399 #[test]
400 fn fetch_user_agent_empty_string() {
401 let opts = FetchOptions::new("https://example.com").user_agent("");
402 assert_eq!(opts.user_agent.as_deref(), Some(""));
403 }
404
405 #[test]
406 fn page_markdown_from_html() {
407 let page = Page {
408 html: "<html><head><title>Test</title></head><body><p>hello world</p></body></html>".into(),
409 inner_text: "hello world".into(),
410 ..Page::default()
411 };
412 let md = page.markdown().unwrap();
413 assert!(md.contains("hello world"));
414 }
415
416 #[test]
417 fn page_extract_json_produces_valid_json() {
418 let page = Page {
419 html: "<html><head><title>Test</title></head><body><p>content</p></body></html>".into(),
420 inner_text: "content".into(),
421 ..Page::default()
422 };
423 let json = page.extract_json().unwrap();
424 let _: Value = serde_json::from_str(&json).expect("valid JSON");
425 }
426
427 #[test]
428 fn page_screenshot_png_none_by_default() {
429 let page = Page::default();
430 assert!(page.screenshot_png().is_none());
431 }
432
433 #[test]
434 fn page_markdown_with_selector_scopes_to_subtree() {
435 let page = Page {
436 html: "<html><body><article>keep</article><aside>drop</aside></body></html>".into(),
437 ..Page::default()
438 };
439 let md = page.markdown_with_selector("https://example.com", "article").unwrap();
440 assert!(md.contains("keep"));
441 assert!(!md.contains("drop"));
442 }
443
444 #[test]
445 fn page_extract_json_with_selector_includes_url() {
446 let page = Page {
447 html: "<html><body><article>scoped</article></body></html>".into(),
448 ..Page::default()
449 };
450 let json = page
451 .extract_json_with_selector("https://example.com/page", "article")
452 .unwrap();
453 let parsed: Value = serde_json::from_str(&json).expect("valid JSON");
454 assert_eq!(parsed["url"].as_str(), Some("https://example.com/page"));
455 assert!(parsed["text_content"].as_str().unwrap().contains("scoped"));
456 }
457
458 #[test]
459 fn page_markdown_with_selector_no_match_returns_empty() {
460 let page = Page {
461 html: "<html><body><article>x</article></body></html>".into(),
462 ..Page::default()
463 };
464 let md = page.markdown_with_selector("", ".nonexistent").unwrap();
465 assert!(md.is_empty());
466 }
467
468 #[test]
469 fn page_markdown_with_invalid_selector_returns_error() {
470 let page = Page {
471 html: "<html><body><p>x</p></body></html>".into(),
472 ..Page::default()
473 };
474 let err = page.markdown_with_selector("", "###invalid[[[").unwrap_err();
475 assert!(err.to_string().contains("invalid CSS selector"));
476 }
477
478 #[test]
479 fn page_markdown_with_empty_selector_returns_error() {
480 let page = Page {
481 html: "<html><body><p>x</p></body></html>".into(),
482 ..Page::default()
483 };
484 assert!(page.markdown_with_selector("", "").is_err());
485 }
486
487 #[test]
488 fn fetch_rejects_invalid_url() {
489 let result = fetch(FetchOptions::new("not a url"));
490 assert!(result.is_err());
491 let err = result.unwrap_err();
492 assert!(matches!(err, Error::InvalidUrl { .. }));
493 }
494
495 #[test]
496 fn fetch_rejects_private_ip() {
497 let result = fetch(FetchOptions::new("http://127.0.0.1/"));
498 assert!(result.is_err());
499 }
500
501 #[test]
502 fn fetch_rejects_file_scheme() {
503 let result = fetch(FetchOptions::new("file:///etc/passwd"));
504 assert!(result.is_err());
505 }
506
507 mod page_from_servo {
508 use crate::bridge;
509 use crate::fetch::{ConsoleLevel, Page};
510
511 fn synthetic_image(w: u32, h: u32) -> image::RgbaImage {
512 image::RgbaImage::from_pixel(w, h, image::Rgba([255, 0, 0, 255]))
513 }
514
515 fn empty_servo_page() -> bridge::ServoPage {
516 bridge::ServoPage::default()
517 }
518
519 #[test]
520 fn extracts_title_from_html() {
521 let mut sp = empty_servo_page();
522 sp.html = "<html><head><title>Hello World</title></head></html>".into();
523 let page = Page::from_servo(sp);
524 assert_eq!(page.title.as_deref(), Some("Hello World"));
525 }
526
527 #[test]
528 fn title_is_none_when_tag_missing() {
529 let mut sp = empty_servo_page();
530 sp.html = "<html><body>no title here</body></html>".into();
531 let page = Page::from_servo(sp);
532 assert!(page.title.is_none());
533 }
534
535 #[test]
536 fn title_is_none_when_tag_empty() {
537 let mut sp = empty_servo_page();
538 sp.html = "<html><head><title></title></head></html>".into();
539 let page = Page::from_servo(sp);
540 assert!(page.title.is_none());
541 }
542
543 #[test]
544 fn title_is_none_for_empty_html() {
545 let page = Page::from_servo(empty_servo_page());
546 assert!(page.title.is_none());
547 }
548
549 #[test]
550 fn inner_text_none_becomes_empty_string() {
551 let sp = empty_servo_page();
552 assert!(sp.inner_text.is_none());
553 let page = Page::from_servo(sp);
554 assert_eq!(page.inner_text, "");
555 }
556
557 #[test]
558 fn screenshot_is_encoded_as_png() {
559 let mut sp = empty_servo_page();
560 sp.screenshot = Some(synthetic_image(8, 8));
561 let page = Page::from_servo(sp);
562 let bytes = page.screenshot_png().expect("screenshot encoded");
563 assert_eq!(&bytes[..8], b"\x89PNG\r\n\x1a\n", "PNG magic bytes");
564 }
565
566 #[test]
567 fn console_messages_empty_by_default() {
568 let page = Page::from_servo(empty_servo_page());
569 assert!(page.console_messages.is_empty());
570 }
571
572 #[test]
573 fn console_messages_preserve_all_six_levels() {
574 let cases = [
575 (bridge::ConsoleLevel::Log, ConsoleLevel::Log),
576 (bridge::ConsoleLevel::Debug, ConsoleLevel::Debug),
577 (bridge::ConsoleLevel::Info, ConsoleLevel::Info),
578 (bridge::ConsoleLevel::Warn, ConsoleLevel::Warn),
579 (bridge::ConsoleLevel::Error, ConsoleLevel::Error),
580 (bridge::ConsoleLevel::Trace, ConsoleLevel::Trace),
581 ];
582 for (src, expected) in cases {
583 let mut sp = empty_servo_page();
584 sp.console_messages = vec![bridge::ConsoleMessage {
585 level: src,
586 message: "msg".into(),
587 }];
588 let page = Page::from_servo(sp);
589 assert_eq!(
590 page.console_messages.len(),
591 1,
592 "console message lost for source level {src:?}",
593 );
594 assert_eq!(
595 page.console_messages[0].level, expected,
596 "level mapping wrong for source {src:?}",
597 );
598 }
599 }
600
601 #[test]
602 fn console_messages_preserve_ordering_across_levels() {
603 let mut sp = empty_servo_page();
604 sp.console_messages = vec![
605 bridge::ConsoleMessage {
606 level: bridge::ConsoleLevel::Info,
607 message: "first".into(),
608 },
609 bridge::ConsoleMessage {
610 level: bridge::ConsoleLevel::Error,
611 message: "second".into(),
612 },
613 bridge::ConsoleMessage {
614 level: bridge::ConsoleLevel::Warn,
615 message: "third".into(),
616 },
617 ];
618 let page = Page::from_servo(sp);
619 assert_eq!(page.console_messages.len(), 3);
620 assert_eq!(page.console_messages[0].message, "first");
621 assert_eq!(page.console_messages[1].message, "second");
622 assert_eq!(page.console_messages[2].message, "third");
623 assert_eq!(page.console_messages[0].level, ConsoleLevel::Info);
624 assert_eq!(page.console_messages[1].level, ConsoleLevel::Error);
625 assert_eq!(page.console_messages[2].level, ConsoleLevel::Warn);
626 }
627
628 #[test]
629 fn extracted_starts_as_none_until_schema_applied() {
630 let page = Page::from_servo(empty_servo_page());
631 assert!(page.extracted.is_none());
632 }
633
634 #[test]
635 fn full_round_trip_preserves_every_field() {
636 let sp = bridge::ServoPage {
637 html: "<html><head><title>T</title></head><body>B</body></html>".into(),
638 inner_text: Some("B".into()),
639 layout_json: Some("[]".into()),
640 visibility_json: Some("[]".into()),
641 screenshot: Some(synthetic_image(2, 2)),
642 js_result: Some("42".into()),
643 accessibility_tree: Some("{}".into()),
644 a11y: None,
645 console_messages: vec![bridge::ConsoleMessage {
646 level: bridge::ConsoleLevel::Log,
647 message: "x".into(),
648 }],
649 };
650 let page = Page::from_servo(sp);
651 assert_eq!(page.html, "<html><head><title>T</title></head><body>B</body></html>");
652 assert_eq!(page.inner_text, "B");
653 assert_eq!(page.title.as_deref(), Some("T"));
654 assert_eq!(page.layout_json.as_deref(), Some("[]"));
655 assert_eq!(page.js_result.as_deref(), Some("42"));
656 assert_eq!(page.accessibility_tree.as_deref(), Some("{}"));
657 assert_eq!(page.console_messages.len(), 1);
658 assert!(page.screenshot_png().is_some());
659 assert!(page.extracted.is_none());
660 }
661 }
662}