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 #[cfg(feature = "lightpanda")]
19 if self.uses_lightpanda() {
20 return self.render_with_lightpanda(request).await;
21 }
22
23 let timeout = request.timeout();
24 match tokio::time::timeout(timeout, self.render_inner(request)).await {
25 Ok(result) => result,
26 Err(_) => Err(UseError::new(
27 "use.browser.timeout",
28 format!("Browser rendering exceeded {} ms.", timeout.as_millis()),
29 )),
30 }
31 }
32}
33
34impl BrowserPool {
35 async fn render_inner(&self, request: RenderRequest) -> UseResult<RenderedPage> {
36 let started = Instant::now();
37 let _permit = self
38 .tab_semaphore()
39 .acquire()
40 .await
41 .map_err(|error| browser_error(format!("Tab limit is closed: {error}")))?;
42 let browser = self.acquire_browser().await?;
43 let page = browser
44 .new_page("about:blank")
45 .await
46 .map_err(|error| browser_error(format!("Failed to open browser tab: {error}")))?;
47 let guard = PageGuard::new(page);
48 let page = guard.page()?;
49
50 if let Some(user_agent) = &request.user_agent {
51 page.set_user_agent(SetUserAgentOverrideParams::new(user_agent))
52 .await
53 .map_err(|error| {
54 browser_error(format!("Failed to set browser user agent: {error}"))
55 })?;
56 }
57
58 page.goto(request.url.as_str())
59 .await
60 .map_err(|error| browser_error(format!("Browser navigation failed: {error}")))?;
61 apply_wait_condition(page, &request.wait).await?;
62 let html = page
63 .content()
64 .await
65 .map_err(|error| browser_error(format!("Failed to read rendered HTML: {error}")))?;
66 let final_url = page
67 .url()
68 .await
69 .ok()
70 .flatten()
71 .and_then(|value| Url::parse(&value).ok())
72 .unwrap_or_else(|| request.url.clone());
73 let artifacts = match &request.screenshot_path {
74 Some(path) => vec![capture_screenshot(page, path).await?],
75 None => Vec::new(),
76 };
77 guard.close().await;
78
79 Ok(RenderedPage {
80 requested_url: request.url,
81 final_url,
82 status: None,
83 content_type: Some("text/html".to_string()),
84 html,
85 elapsed_ms: started.elapsed().as_millis().try_into().unwrap_or(u64::MAX),
86 artifacts,
87 })
88 }
89}
90
91pub(crate) async fn apply_wait_condition(
92 page: &chromiumoxide::Page,
93 condition: &WaitCondition,
94) -> UseResult<()> {
95 match condition {
96 WaitCondition::Load | WaitCondition::DomContentLoaded => {}
98 WaitCondition::NetworkIdle { idle_ms } => {
99 tokio::time::sleep(Duration::from_millis(*idle_ms)).await;
100 }
101 WaitCondition::Selector { css, timeout_ms } => {
102 match tokio::time::timeout(
103 Duration::from_millis(*timeout_ms),
104 page.find_element(css.as_str()),
105 )
106 .await
107 {
108 Ok(Ok(_)) => {}
109 Ok(Err(error)) => {
110 return Err(browser_error(format!(
111 "Browser selector '{css}' failed: {error}"
112 )))
113 }
114 Err(_) => {
115 return Err(UseError::new(
116 "use.browser.wait_timeout",
117 format!("Selector '{css}' was not found within {timeout_ms} ms."),
118 ))
119 }
120 }
121 }
122 WaitCondition::Delay { ms } => {
123 tokio::time::sleep(Duration::from_millis(*ms)).await;
124 }
125 }
126 Ok(())
127}
128
129pub(crate) async fn capture_screenshot(
130 page: &chromiumoxide::Page,
131 path: &std::path::Path,
132) -> UseResult<Artifact> {
133 let bytes = page
134 .save_screenshot(
135 ScreenshotParams::builder()
136 .format(CaptureScreenshotFormat::Png)
137 .full_page(true)
138 .build(),
139 path,
140 )
141 .await
142 .map_err(|error| browser_error(format!("Failed to save browser screenshot: {error}")))?;
143 let sha256 = Sha256::digest(&bytes)
144 .iter()
145 .map(|byte| format!("{byte:02x}"))
146 .collect();
147 Ok(Artifact {
148 path: path.to_path_buf(),
149 media_type: "image/png".to_string(),
150 size: bytes.len().try_into().unwrap_or(u64::MAX),
151 sha256,
152 })
153}
154
155struct PageGuard {
156 page: Option<chromiumoxide::Page>,
157}
158
159impl PageGuard {
160 fn new(page: chromiumoxide::Page) -> Self {
161 Self { page: Some(page) }
162 }
163
164 fn page(&self) -> UseResult<&chromiumoxide::Page> {
165 self.page.as_ref().ok_or_else(|| {
166 UseError::new(
167 "use.browser.page_closed",
168 "The browser page was closed before rendering completed.",
169 )
170 })
171 }
172
173 async fn close(mut self) {
174 if let Some(page) = self.page.take() {
175 if let Err(error) = page.close().await {
176 warn!("Failed to close browser tab: {error}");
177 }
178 }
179 }
180}
181
182impl Drop for PageGuard {
183 fn drop(&mut self) {
184 if let Some(page) = self.page.take() {
185 match tokio::runtime::Handle::try_current() {
186 Ok(runtime) => {
187 runtime.spawn(async move {
188 if let Err(error) = page.close().await {
189 warn!("Failed to close browser tab after cancellation: {error}");
190 }
191 });
192 }
193 Err(error) => warn!("Cannot schedule browser tab cleanup: {error}"),
194 }
195 }
196 }
197}
198
199#[cfg(all(test, not(windows)))]
200mod tests {
201 use super::*;
202 use crate::{BrowserPoolConfig, BrowserProvider};
203 #[cfg(feature = "lightpanda")]
204 use std::sync::Arc;
205
206 #[cfg(feature = "lightpanda")]
207 fn executable_fixture(contents: &str) -> (tempfile::TempDir, std::path::PathBuf) {
208 use std::os::unix::fs::PermissionsExt;
209
210 let directory = tempfile::tempdir().unwrap();
211 let executable = directory.path().join("lightpanda");
212 std::fs::write(&executable, contents).unwrap();
213 let mut permissions = std::fs::metadata(&executable).unwrap().permissions();
214 permissions.set_mode(0o755);
215 std::fs::set_permissions(&executable, permissions).unwrap();
216 (directory, executable)
217 }
218
219 #[cfg(feature = "lightpanda")]
220 #[tokio::test]
221 async fn lightpanda_renderer_uses_the_fetch_command_for_html() {
222 let (_directory, executable) = executable_fixture(
223 "#!/bin/sh\nprintf '<!DOCTYPE html><html><body>cli fixture</body></html>'\n",
224 );
225 let pool = BrowserPool::new(BrowserPoolConfig {
226 provider: BrowserProvider::LightpandaExecutable(executable),
227 ..BrowserPoolConfig::default()
228 });
229 let request = RenderRequest {
230 url: Url::parse("https://example.test/search?q=rust").unwrap(),
231 timeout_ms: 5_000,
232 wait: WaitCondition::Load,
233 user_agent: None,
234 screenshot_path: None,
235 };
236
237 let rendered = pool.render(request).await;
238 pool.shutdown().await;
239
240 let rendered = rendered.unwrap();
241 assert!(rendered.html.contains("cli fixture"));
242 assert_eq!(rendered.content_type.as_deref(), Some("text/html"));
243 }
244
245 #[cfg(feature = "lightpanda")]
246 #[tokio::test]
247 async fn lightpanda_renderer_forwards_url_deadline_and_proxy_as_arguments() {
248 let directory = tempfile::tempdir().unwrap();
249 let arguments = directory.path().join("arguments.txt");
250 let script = format!(
251 "#!/bin/sh\nprintf '%s\\n' \"$@\" > '{}'\nprintf '<html>arguments fixture</html>'\n",
252 arguments.display()
253 );
254 let (_executable_directory, executable) = executable_fixture(&script);
255 let proxy = "http://user:secret@proxy.example:8080";
256 let pool = BrowserPool::new(BrowserPoolConfig {
257 provider: BrowserProvider::LightpandaExecutable(executable),
258 proxy_url: Some(proxy.to_string()),
259 ..BrowserPoolConfig::default()
260 });
261 let request = RenderRequest {
262 url: Url::parse("https://example.test/search?q=rust").unwrap(),
263 timeout_ms: 5_000,
264 wait: WaitCondition::Load,
265 user_agent: None,
266 screenshot_path: None,
267 };
268
269 pool.render(request).await.unwrap();
270 pool.shutdown().await;
271
272 let arguments = std::fs::read_to_string(arguments).unwrap();
273 assert!(arguments.contains("fetch\n"));
274 assert!(arguments.contains("--dump\nhtml\n"));
275 assert!(arguments.contains("--http_connect_timeout\n5000\n"));
276 assert!(arguments.contains("--http_timeout\n5000\n"));
277 assert!(arguments.contains(&format!("--http_proxy\n{proxy}\n")));
278 assert!(arguments.ends_with("https://example.test/search?q=rust\n"));
279 }
280
281 #[cfg(feature = "lightpanda")]
282 #[tokio::test]
283 async fn lightpanda_renderer_kills_and_reaps_a_timed_out_fetch() {
284 let directory = tempfile::tempdir().unwrap();
285 let pid_file = directory.path().join("pid.txt");
286 let script = format!(
287 "#!/bin/sh\nprintf '%s' \"$$\" > '{}'\nexec sleep 30\n",
288 pid_file.display()
289 );
290 let (_executable_directory, executable) = executable_fixture(&script);
291 let pool = BrowserPool::new(BrowserPoolConfig {
292 provider: BrowserProvider::LightpandaExecutable(executable),
293 ..BrowserPoolConfig::default()
294 });
295 let request = RenderRequest {
296 url: Url::parse("https://example.test/").unwrap(),
297 timeout_ms: 2_000,
298 wait: WaitCondition::Load,
299 user_agent: None,
300 screenshot_path: None,
301 };
302
303 let error = pool.render(request).await.unwrap_err();
304 pool.shutdown().await;
305
306 assert_eq!(error.code, "use.browser.timeout");
307 let pid = std::fs::read_to_string(pid_file).unwrap();
308 let still_running = std::process::Command::new("kill")
309 .args(["-0", pid.trim()])
310 .output()
311 .unwrap()
312 .status
313 .success();
314 assert!(
315 !still_running,
316 "timed-out Lightpanda process {pid} survived"
317 );
318 }
319
320 #[cfg(feature = "lightpanda")]
321 #[tokio::test]
322 async fn cancelling_lightpanda_render_still_kills_and_reaps_the_fetch() {
323 let directory = tempfile::tempdir().unwrap();
324 let pid_file = directory.path().join("cancelled-pid.txt");
325 let script = format!(
326 "#!/bin/sh\nprintf '%s' \"$$\" > '{}'\nexec sleep 30\n",
327 pid_file.display()
328 );
329 let (_executable_directory, executable) = executable_fixture(&script);
330 let pool = Arc::new(BrowserPool::new(BrowserPoolConfig {
331 provider: BrowserProvider::LightpandaExecutable(executable),
332 ..BrowserPoolConfig::default()
333 }));
334 let request = RenderRequest {
335 url: Url::parse("https://example.test/").unwrap(),
336 timeout_ms: 30_000,
337 wait: WaitCondition::Load,
338 user_agent: None,
339 screenshot_path: None,
340 };
341 let render_pool = Arc::clone(&pool);
342 let render = tokio::spawn(async move { render_pool.render(request).await });
343 tokio::time::timeout(std::time::Duration::from_secs(2), async {
344 while !pid_file.is_file() {
345 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
346 }
347 })
348 .await
349 .expect("Lightpanda fixture did not start");
350
351 render.abort();
352 let _ = render.await;
353 let pid = std::fs::read_to_string(pid_file).unwrap();
354 tokio::time::timeout(std::time::Duration::from_secs(2), async {
355 loop {
356 let still_running = std::process::Command::new("kill")
357 .args(["-0", pid.trim()])
358 .output()
359 .unwrap()
360 .status
361 .success();
362 if !still_running {
363 break;
364 }
365 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
366 }
367 })
368 .await
369 .expect("cancelled Lightpanda process was not reaped");
370 pool.shutdown().await;
371 }
372
373 #[cfg(feature = "lightpanda")]
374 #[tokio::test]
375 async fn lightpanda_renderer_rejects_unsupported_exact_user_agent_without_spawning() {
376 let pool = BrowserPool::new(BrowserPoolConfig {
377 provider: BrowserProvider::LightpandaExecutable("/not/spawned".into()),
378 ..BrowserPoolConfig::default()
379 });
380 let request = RenderRequest {
381 url: Url::parse("https://example.test/").unwrap(),
382 timeout_ms: 1_000,
383 wait: WaitCondition::Load,
384 user_agent: Some("exact-agent".to_string()),
385 screenshot_path: None,
386 };
387
388 let error = pool.render(request).await.unwrap_err();
389 pool.shutdown().await;
390
391 assert_eq!(error.code, "use.browser.unsupported");
392 assert!(error.message.contains("user-agent"));
393 }
394
395 #[cfg(feature = "lightpanda")]
396 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
397 async fn installed_lightpanda_renders_a_local_http_page_when_available() {
398 use tokio::io::AsyncWriteExt;
399
400 let Some(executable) = crate::detect_lightpanda() else {
401 return;
402 };
403 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
404 let address = listener.local_addr().unwrap();
405 let server = tokio::spawn(async move {
406 let (mut stream, _) = listener.accept().await.unwrap();
407 stream
408 .write_all(
409 b"HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=utf-8\r\nContent-Length: 41\r\nConnection: close\r\n\r\n<html><body>runtime fixture</body></html>",
410 )
411 .await
412 .unwrap();
413 });
414 let pool = BrowserPool::new(BrowserPoolConfig {
415 provider: BrowserProvider::LightpandaExecutable(executable),
416 ..BrowserPoolConfig::default()
417 });
418 let request = RenderRequest {
419 url: Url::parse(&format!("http://{address}/fixture")).unwrap(),
420 timeout_ms: 5_000,
421 wait: WaitCondition::Load,
422 user_agent: None,
423 screenshot_path: None,
424 };
425
426 let rendered = pool.render(request).await;
427 pool.shutdown().await;
428 server.abort();
429
430 let rendered = rendered.unwrap();
431 assert!(rendered.html.contains("runtime fixture"));
432 }
433
434 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
435 async fn discovered_chrome_renders_a_network_free_page_when_available() {
436 let _guard = crate::test_support::lock_chrome_integration_test().await;
437 let Some(executable) = crate::detect_chrome() else {
438 return;
439 };
440 let pool = BrowserPool::new(BrowserPoolConfig {
441 provider: BrowserProvider::ChromeExecutable(executable),
442 ..BrowserPoolConfig::default()
443 });
444 let request = RenderRequest {
445 url: Url::parse("data:text/html,<main id='fixture'>a3s-use</main>").unwrap(),
446 timeout_ms: crate::test_support::CHROME_OPERATION_TIMEOUT_MS,
447 wait: WaitCondition::Load,
448 user_agent: Some("a3s-use-browser-test".to_string()),
449 screenshot_path: None,
450 };
451
452 let rendered = pool.render(request).await;
453 pool.shutdown().await;
454
455 let rendered = rendered.unwrap();
456 assert!(rendered.html.contains("a3s-use"));
457 assert_eq!(rendered.content_type.as_deref(), Some("text/html"));
458 }
459}