a3s_use_browser/
renderer.rs

1use std::time::{Duration, Instant};
2
3use a3s_use_core::{Artifact, UseError, UseResult};
4use async_trait::async_trait;
5use chromiumoxide::cdp::browser_protocol::network::SetUserAgentOverrideParams;
6use chromiumoxide::cdp::browser_protocol::page::CaptureScreenshotFormat;
7use chromiumoxide::page::ScreenshotParams;
8use sha2::{Digest, Sha256};
9use tracing::warn;
10use url::Url;
11
12use crate::pool::{browser_error, BrowserPool};
13use crate::{PageRenderer, RenderRequest, RenderedPage, WaitCondition};
14
15#[async_trait]
16impl PageRenderer for BrowserPool {
17    async fn render(&self, request: RenderRequest) -> UseResult<RenderedPage> {
18        let timeout = request.timeout();
19        match tokio::time::timeout(timeout, self.render_inner(request)).await {
20            Ok(result) => result,
21            Err(_) => Err(UseError::new(
22                "use.browser.timeout",
23                format!("Browser rendering exceeded {} ms.", timeout.as_millis()),
24            )),
25        }
26    }
27}
28
29impl BrowserPool {
30    async fn render_inner(&self, request: RenderRequest) -> UseResult<RenderedPage> {
31        let started = Instant::now();
32        let _permit = self
33            .tab_semaphore()
34            .acquire()
35            .await
36            .map_err(|error| browser_error(format!("Tab limit is closed: {error}")))?;
37        let browser = self.acquire_browser().await?;
38        let page = browser
39            .new_page("about:blank")
40            .await
41            .map_err(|error| browser_error(format!("Failed to open browser tab: {error}")))?;
42        let guard = PageGuard::new(page);
43        let page = guard.page()?;
44
45        if let Some(user_agent) = &request.user_agent {
46            page.set_user_agent(SetUserAgentOverrideParams::new(user_agent))
47                .await
48                .map_err(|error| {
49                    browser_error(format!("Failed to set browser user agent: {error}"))
50                })?;
51        }
52
53        page.goto(request.url.as_str())
54            .await
55            .map_err(|error| browser_error(format!("Browser navigation failed: {error}")))?;
56        apply_wait_condition(page, &request.wait).await?;
57        let html = page
58            .content()
59            .await
60            .map_err(|error| browser_error(format!("Failed to read rendered HTML: {error}")))?;
61        let final_url = page
62            .url()
63            .await
64            .ok()
65            .flatten()
66            .and_then(|value| Url::parse(&value).ok())
67            .unwrap_or_else(|| request.url.clone());
68        let artifacts = match &request.screenshot_path {
69            Some(path) => vec![capture_screenshot(page, path).await?],
70            None => Vec::new(),
71        };
72        guard.close().await;
73
74        Ok(RenderedPage {
75            requested_url: request.url,
76            final_url,
77            status: None,
78            content_type: Some("text/html".to_string()),
79            html,
80            elapsed_ms: started.elapsed().as_millis().try_into().unwrap_or(u64::MAX),
81            artifacts,
82        })
83    }
84}
85
86pub(crate) async fn apply_wait_condition(
87    page: &chromiumoxide::Page,
88    condition: &WaitCondition,
89) -> UseResult<()> {
90    match condition {
91        // `Page::goto` resolves after the requested page is loaded.
92        WaitCondition::Load | WaitCondition::DomContentLoaded => {}
93        WaitCondition::NetworkIdle { idle_ms } => {
94            tokio::time::sleep(Duration::from_millis(*idle_ms)).await;
95        }
96        WaitCondition::Selector { css, timeout_ms } => {
97            match tokio::time::timeout(
98                Duration::from_millis(*timeout_ms),
99                page.find_element(css.as_str()),
100            )
101            .await
102            {
103                Ok(Ok(_)) => {}
104                Ok(Err(error)) => {
105                    return Err(browser_error(format!(
106                        "Browser selector '{css}' failed: {error}"
107                    )))
108                }
109                Err(_) => {
110                    return Err(UseError::new(
111                        "use.browser.wait_timeout",
112                        format!("Selector '{css}' was not found within {timeout_ms} ms."),
113                    ))
114                }
115            }
116        }
117        WaitCondition::Delay { ms } => {
118            tokio::time::sleep(Duration::from_millis(*ms)).await;
119        }
120    }
121    Ok(())
122}
123
124pub(crate) async fn capture_screenshot(
125    page: &chromiumoxide::Page,
126    path: &std::path::Path,
127) -> UseResult<Artifact> {
128    let bytes = page
129        .save_screenshot(
130            ScreenshotParams::builder()
131                .format(CaptureScreenshotFormat::Png)
132                .full_page(true)
133                .build(),
134            path,
135        )
136        .await
137        .map_err(|error| browser_error(format!("Failed to save browser screenshot: {error}")))?;
138    let sha256 = Sha256::digest(&bytes)
139        .iter()
140        .map(|byte| format!("{byte:02x}"))
141        .collect();
142    Ok(Artifact {
143        path: path.to_path_buf(),
144        media_type: "image/png".to_string(),
145        size: bytes.len().try_into().unwrap_or(u64::MAX),
146        sha256,
147    })
148}
149
150struct PageGuard {
151    page: Option<chromiumoxide::Page>,
152}
153
154impl PageGuard {
155    fn new(page: chromiumoxide::Page) -> Self {
156        Self { page: Some(page) }
157    }
158
159    fn page(&self) -> UseResult<&chromiumoxide::Page> {
160        self.page.as_ref().ok_or_else(|| {
161            UseError::new(
162                "use.browser.page_closed",
163                "The browser page was closed before rendering completed.",
164            )
165        })
166    }
167
168    async fn close(mut self) {
169        if let Some(page) = self.page.take() {
170            if let Err(error) = page.close().await {
171                warn!("Failed to close browser tab: {error}");
172            }
173        }
174    }
175}
176
177impl Drop for PageGuard {
178    fn drop(&mut self) {
179        if let Some(page) = self.page.take() {
180            match tokio::runtime::Handle::try_current() {
181                Ok(runtime) => {
182                    runtime.spawn(async move {
183                        if let Err(error) = page.close().await {
184                            warn!("Failed to close browser tab after cancellation: {error}");
185                        }
186                    });
187                }
188                Err(error) => warn!("Cannot schedule browser tab cleanup: {error}"),
189            }
190        }
191    }
192}
193
194#[cfg(all(test, not(windows)))]
195mod tests {
196    use super::*;
197    use crate::{BrowserPoolConfig, BrowserProvider};
198
199    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
200    async fn discovered_chrome_renders_a_network_free_page_when_available() {
201        let Some(executable) = crate::detect_chrome() else {
202            return;
203        };
204        let pool = BrowserPool::new(BrowserPoolConfig {
205            provider: BrowserProvider::ChromeExecutable(executable),
206            ..BrowserPoolConfig::default()
207        });
208        let request = RenderRequest {
209            url: Url::parse("data:text/html,<main id='fixture'>a3s-use</main>").unwrap(),
210            timeout_ms: 10_000,
211            wait: WaitCondition::Load,
212            user_agent: Some("a3s-use-browser-test".to_string()),
213            screenshot_path: None,
214        };
215
216        let rendered = pool.render(request).await;
217        pool.shutdown().await;
218
219        let rendered = rendered.unwrap();
220        assert!(rendered.html.contains("a3s-use"));
221        assert_eq!(rendered.content_type.as_deref(), Some("text/html"));
222    }
223}
a3s_use_browser/renderer.rs

a3s_use_browser/
renderer.rs
