servo_fetch/
client.rs

1//! Async client with reusable defaults.
2
3use std::sync::Arc;
4use std::time::Duration;
5
6use crate::error::{Error, Result};
7use crate::fetch::{self, FetchOptions, Page};
8use crate::net::sanitize_user_agent;
9use crate::visibility::VisibilityPolicy;
10
11/// Async client carrying reusable fetch defaults.
12#[derive(Debug, Clone)]
13pub struct Client {
14    inner: Arc<ClientInner>,
15}
16
17#[derive(Debug)]
18pub(crate) struct ClientInner {
19    pub(crate) timeout: Duration,
20    pub(crate) settle: Duration,
21    pub(crate) user_agent: Option<String>,
22    pub(crate) visibility: VisibilityPolicy,
23    pub(crate) headers: http::HeaderMap,
24}
25
26impl ClientInner {
27    /// Unset fields fall back to these defaults; values already set on `opts` win.
28    pub(crate) fn apply_defaults(&self, mut opts: FetchOptions) -> FetchOptions {
29        opts.timeout.get_or_insert(self.timeout);
30        opts.settle.get_or_insert(self.settle);
31        opts.visibility.get_or_insert(self.visibility);
32        if let Some(ua) = &self.user_agent {
33            opts.user_agent.get_or_insert_with(|| ua.clone());
34        }
35        if opts.headers.is_empty() && !self.headers.is_empty() {
36            opts.headers = self.headers.clone();
37        }
38        opts
39    }
40
41    pub(crate) fn options(&self, url: &str) -> FetchOptions {
42        self.apply_defaults(FetchOptions::new(url))
43    }
44}
45
46impl Default for Client {
47    fn default() -> Self {
48        Self::new()
49    }
50}
51
52impl Client {
53    /// Construct a client with default options.
54    #[must_use]
55    pub fn new() -> Self {
56        Self::builder().build()
57    }
58
59    /// Begin building a client with custom options.
60    pub fn builder() -> ClientBuilder {
61        ClientBuilder::default()
62    }
63
64    /// Fetch a page using the client's default options.
65    pub async fn fetch(&self, url: &str) -> Result<Page> {
66        fetch::fetch(&self.inner.options(url)).await
67    }
68
69    /// Fetch a page with explicit options.
70    pub async fn fetch_with(&self, opts: &FetchOptions) -> Result<Page> {
71        fetch::fetch(&self.inner.apply_defaults(opts.clone())).await
72    }
73
74    /// Fetch and extract readable Markdown.
75    pub async fn markdown(&self, url: &str) -> Result<String> {
76        self.fetch(url).await?.markdown_with_url(url)
77    }
78
79    /// Fetch and extract plain text (`document.body.innerText`).
80    pub async fn text(&self, url: &str) -> Result<String> {
81        Ok(self.fetch(url).await?.inner_text)
82    }
83
84    /// Fetch and extract structured JSON.
85    pub async fn extract_json(&self, url: &str) -> Result<String> {
86        self.fetch(url).await?.extract_json_with_url(url)
87    }
88
89    /// Capture a PNG screenshot of the page.
90    pub async fn screenshot(&self, url: &str, opts: &ScreenshotOptions) -> Result<Vec<u8>> {
91        let fopts = self.inner.apply_defaults(FetchOptions::screenshot(url, opts.full_page));
92        let page = fetch::fetch(&fopts).await?;
93        page.screenshot_png()
94            .map(<[u8]>::to_vec)
95            .ok_or_else(|| Error::screenshot(anyhow::anyhow!("screenshot returned no data"), Some(url.to_string())))
96    }
97
98    /// Execute a JavaScript expression after page load and return the result.
99    pub async fn execute_js(&self, url: &str, expression: impl Into<String>) -> Result<String> {
100        let fopts = self.inner.apply_defaults(FetchOptions::javascript(url, expression));
101        let page = fetch::fetch(&fopts).await?;
102        page.js_result
103            .ok_or_else(|| Error::javascript(anyhow::anyhow!("execute_js returned no result"), Some(url.to_string())))
104    }
105}
106
107/// Options for [`Client::screenshot`].
108#[derive(Debug, Default, Clone)]
109#[non_exhaustive]
110pub struct ScreenshotOptions {
111    /// Capture the full scrollable page (rather than only the viewport).
112    pub full_page: bool,
113}
114
115/// Builder for [`Client`].
116#[must_use = "ClientBuilder does nothing until .build() is called"]
117#[derive(Debug, Default)]
118pub struct ClientBuilder {
119    timeout: Option<Duration>,
120    settle: Option<Duration>,
121    user_agent: Option<String>,
122    visibility: Option<VisibilityPolicy>,
123    headers: http::HeaderMap,
124}
125
126impl ClientBuilder {
127    /// Default page-load timeout (default 30s).
128    pub fn timeout(mut self, timeout: Duration) -> Self {
129        self.timeout = Some(timeout);
130        self
131    }
132
133    /// Default settle wait after the load event (default 0).
134    pub fn settle(mut self, settle: Duration) -> Self {
135        self.settle = Some(settle);
136        self
137    }
138
139    /// Default User-Agent header.
140    pub fn user_agent(mut self, ua: impl Into<String>) -> Self {
141        self.user_agent = Some(sanitize_user_agent(ua.into()));
142        self
143    }
144
145    /// Default visibility policy applied during extraction.
146    pub fn visibility(mut self, policy: VisibilityPolicy) -> Self {
147        self.visibility = Some(policy);
148        self
149    }
150
151    /// Default custom request headers sent with every navigation request.
152    pub fn headers(mut self, headers: http::HeaderMap) -> Self {
153        self.headers = headers;
154        self
155    }
156
157    /// Build the [`Client`].
158    #[must_use]
159    pub fn build(self) -> Client {
160        Client {
161            inner: Arc::new(self.build_inner()),
162        }
163    }
164
165    pub(crate) fn build_inner(self) -> ClientInner {
166        ClientInner {
167            timeout: self.timeout.unwrap_or(FetchOptions::DEFAULT_TIMEOUT),
168            settle: self.settle.unwrap_or_default(),
169            user_agent: self.user_agent,
170            visibility: self.visibility.unwrap_or_default(),
171            headers: self.headers,
172        }
173    }
174}
175
176#[cfg(test)]
177mod tests {
178    use super::*;
179
180    #[test]
181    fn client_default_uses_30s_timeout() {
182        assert_eq!(Client::default().inner.timeout, FetchOptions::DEFAULT_TIMEOUT);
183    }
184
185    #[test]
186    fn client_builder_sets_timeout() {
187        let client = Client::builder().timeout(Duration::from_secs(60)).build();
188        assert_eq!(client.inner.timeout, Duration::from_secs(60));
189    }
190
191    #[test]
192    fn client_builder_sets_settle() {
193        let client = Client::builder().settle(Duration::from_millis(500)).build();
194        assert_eq!(client.inner.settle, Duration::from_millis(500));
195    }
196
197    #[test]
198    fn client_builder_sanitizes_user_agent() {
199        let client = Client::builder().user_agent("Bot\r\nX-Evil: yes").build();
200        assert_eq!(client.inner.user_agent.as_deref(), Some("Bot  X-Evil: yes"));
201    }
202
203    #[test]
204    fn client_options_propagates_defaults() {
205        let client = Client::builder()
206            .timeout(Duration::from_secs(60))
207            .settle(Duration::from_millis(500))
208            .user_agent("MyBot")
209            .build();
210        let opts = client.inner.options("https://example.com");
211        assert_eq!(opts.timeout, Some(Duration::from_secs(60)));
212        assert_eq!(opts.settle, Some(Duration::from_millis(500)));
213        assert_eq!(opts.user_agent.as_deref(), Some("MyBot"));
214    }
215
216    #[test]
217    fn client_apply_defaults_caller_value_wins() {
218        let client = Client::builder()
219            .timeout(Duration::from_secs(60))
220            .user_agent("ClientBot")
221            .build();
222        let user_opts = FetchOptions::new("https://example.com")
223            .timeout(Duration::from_secs(10))
224            .user_agent("UserBot");
225        let merged = client.inner.apply_defaults(user_opts);
226        // Caller's explicit values win
227        assert_eq!(merged.timeout, Some(Duration::from_secs(10)));
228        assert_eq!(merged.user_agent.as_deref(), Some("UserBot"));
229    }
230
231    #[test]
232    fn client_apply_defaults_fills_unset_fields() {
233        let client = Client::builder()
234            .timeout(Duration::from_secs(60))
235            .settle(Duration::from_millis(750))
236            .user_agent("ClientBot")
237            .visibility(VisibilityPolicy::off())
238            .build();
239        let user_opts = FetchOptions::new("https://example.com");
240        let merged = client.inner.apply_defaults(user_opts);
241        // None fields filled with client defaults
242        assert_eq!(merged.timeout, Some(Duration::from_secs(60)));
243        assert_eq!(merged.settle, Some(Duration::from_millis(750)));
244        assert_eq!(merged.user_agent.as_deref(), Some("ClientBot"));
245        assert_eq!(merged.visibility, Some(VisibilityPolicy::off()));
246    }
247
248    #[test]
249    fn client_clone_shares_inner() {
250        let client = Client::new();
251        assert!(Arc::ptr_eq(&client.inner, &client.clone().inner));
252    }
253
254    #[test]
255    fn screenshot_options_default_is_viewport() {
256        assert!(!ScreenshotOptions::default().full_page);
257    }
258
259    #[test]
260    fn assert_send_sync() {
261        fn check<T: Send + Sync>() {}
262        check::<Client>();
263        check::<ClientBuilder>();
264        check::<ScreenshotOptions>();
265    }
266
267    #[test]
268    fn client_builder_sets_visibility() {
269        let client = Client::builder().visibility(VisibilityPolicy::off()).build();
270        assert_eq!(client.inner.visibility, VisibilityPolicy::off());
271    }
272
273    #[tokio::test]
274    async fn client_fetch_invalid_url_returns_invalid_url_error() {
275        let client = Client::new();
276        let err = client.fetch("not a url").await.unwrap_err();
277        assert!(matches!(err, Error::InvalidUrl { .. }), "got: {err:?}");
278    }
279
280    #[tokio::test]
281    async fn client_fetch_private_address_is_rejected() {
282        let client = Client::new();
283        let err = client.fetch("http://127.0.0.1/").await.unwrap_err();
284        assert!(err.is_network(), "got: {err:?}");
285    }
286}
servo_fetch/client.rs

servo_fetch/
client.rs
