Add headless Chrome as a toolproof backend

toolproof/Cargo.toml

@@ -25,6 +25,7 @@ console = "0.15"
 dialoguer = { version = "0.11", features = ["fuzzy-select"] }
 async-trait = "0.1.78"
 pagebrowse = "0.1.0"
+chromiumoxide = "0.5"
 clap = { version = "4", features = ["cargo"] }
 schematic = { version = "0.12.0", features = ["yaml"] }
 strip-ansi-escapes = "0.2.0"

toolproof/src/civilization.rs

@@ -3,19 +3,21 @@ use std::{
     fs,
     io::{Read, Write},
     path::PathBuf,
-    process::{Command, ExitStatus},
+    process::{ExitStatus, Stdio},
     str::from_utf8,
     sync::Arc,
+    time::Duration,
 };
 
 use actix_web::dev::ServerHandle;
-use pagebrowse::{Pagebrowser, PagebrowserWindow};
 use portpicker::pick_unused_port;
 use tempfile::tempdir;
-use tokio::task::JoinHandle;
+use tokio::{process::Command, task::JoinHandle};
 use wax::Glob;
 
-use crate::{errors::ToolproofTestFailure, options::ToolproofParams, universe::Universe};
+use crate::{
+    definitions::browser::BrowserWindow, errors::ToolproofTestFailure, universe::Universe,
+};
 
 #[derive(Debug)]
 pub struct CommandOutput {
@@ -27,7 +29,7 @@ pub struct Civilization<'u> {
     pub tmp_dir: Option<tempfile::TempDir>,
     pub last_command_output: Option<CommandOutput>,
     pub assigned_server_port: Option<u16>,
-    pub window: Option<PagebrowserWindow>,
+    pub window: Option<BrowserWindow>,
     pub threads: Vec<JoinHandle<Result<(), std::io::Error>>>,
     pub handles: Vec<ServerHandle>,
     pub env_vars: HashMap<String, String>,
@@ -150,7 +152,7 @@ impl<'u> Civilization<'u> {
         self.env_vars.insert(name, value);
     }
 
-    pub fn run_command(&mut self, cmd: String) -> Result<ExitStatus, ToolproofTestFailure> {
+    pub async fn run_command(&mut self, cmd: String) -> Result<ExitStatus, ToolproofTestFailure> {
         let mut command = Command::new("sh");
         command
             .arg("-c")
@@ -161,7 +163,25 @@ impl<'u> Civilization<'u> {
             command.env(key, val);
         }
 
-        let Ok(output) = command.output() else {
+        command.stdout(Stdio::piped());
+        command.stderr(Stdio::piped());
+        let running = command.spawn().map_err(|_| ToolproofTestFailure::Custom {
+            msg: format!("Failed to run command: {cmd}"),
+        })?;
+
+        let Ok(output) = (match tokio::time::timeout(
+            Duration::from_secs(30),
+            running.wait_with_output(),
+        )
+        .await
+        {
+            Ok(out) => out,
+            Err(_) => {
+                return Err(ToolproofTestFailure::Custom {
+                    msg: format!("Failed to run command due to timeout: {cmd}"),
+                });
+            }
+        }) else {
             return Err(ToolproofTestFailure::Custom {
                 msg: format!("Failed to run command: {cmd}"),
             });

toolproof/src/definitions/browser/mod.rs

@@ -1,15 +1,138 @@
-use std::collections::HashMap;
+use std::sync::Arc;
 
 use async_trait::async_trait;
+use chromiumoxide::cdp::browser_protocol::target::CreateTargetParams;
+use futures::StreamExt;
+use tokio::task::JoinHandle;
 
 use crate::civilization::Civilization;
 use crate::errors::{ToolproofInputError, ToolproofStepError};
+use crate::options::ToolproofParams;
 
 use super::{SegmentArgs, ToolproofInstruction, ToolproofRetriever};
 
+use chromiumoxide::browser::{Browser, BrowserConfig};
 use pagebrowse::{PagebrowseBuilder, Pagebrowser, PagebrowserWindow};
 
 const HARNESS: &'static str = include_str!("./harness.js");
+const INIT_SCRIPT: &'static str = include_str!("./init.js");
+
+fn harnessed(js: String) -> String {
+    HARNESS.replace("// insert_toolproof_inner_js", &js)
+}
+
+pub enum BrowserTester {
+    Pagebrowse(Arc<Pagebrowser>),
+    Chrome {
+        browser: Arc<Browser>,
+        event_thread: Arc<JoinHandle<Result<(), std::io::Error>>>,
+    },
+}
+
+impl BrowserTester {
+    async fn initialize(params: &ToolproofParams) -> Self {
+        match params.browser {
+            crate::options::ToolproofBrowserImpl::Chrome => {
+                let (browser, mut handler) =
+                    Browser::launch(BrowserConfig::builder().build().unwrap())
+                        .await
+                        .unwrap();
+
+                BrowserTester::Chrome {
+                    browser: Arc::new(browser),
+                    event_thread: Arc::new(tokio::task::spawn(async move {
+                        loop {
+                            let _ = handler.next().await.unwrap();
+                        }
+                    })),
+                }
+            }
+            crate::options::ToolproofBrowserImpl::Pagebrowse => {
+                let pagebrowser = PagebrowseBuilder::new(params.concurrency)
+                    .visible(false)
+                    .manager_path(format!(
+                        "{}/../bin/pagebrowse_manager",
+                        env!("CARGO_MANIFEST_DIR")
+                    ))
+                    .init_script(INIT_SCRIPT.to_string())
+                    .build()
+                    .await
+                    .expect("Can't build the pagebrowser");
+
+                BrowserTester::Pagebrowse(Arc::new(pagebrowser))
+            }
+        }
+    }
+
+    async fn get_window(&self) -> BrowserWindow {
+        match self {
+            BrowserTester::Pagebrowse(pb) => {
+                BrowserWindow::Pagebrowse(pb.get_window().await.unwrap())
+            }
+            BrowserTester::Chrome { browser, .. } => {
+                let page = browser
+                    .new_page(CreateTargetParams {
+                        url: "about:blank".to_string(),
+                        width: None,
+                        height: None,
+                        browser_context_id: None,
+                        enable_begin_frame_control: None,
+                        new_window: None,
+                        background: None,
+                    })
+                    .await
+                    .unwrap();
+                page.evaluate_on_new_document(INIT_SCRIPT.to_string())
+                    .await
+                    .expect("Could not set initialization js");
+                BrowserWindow::Chrome(page)
+            }
+        }
+    }
+}
+
+pub enum BrowserWindow {
+    Chrome(chromiumoxide::Page),
+    Pagebrowse(PagebrowserWindow),
+}
+
+impl BrowserWindow {
+    async fn navigate(&self, url: String, wait_for_load: bool) -> Result<(), ToolproofStepError> {
+        match self {
+            BrowserWindow::Chrome(page) => {
+                // TODO: This is implicitly always wait_for_load: true
+                page.goto(url)
+                    .await
+                    .map(|_| ())
+                    .map_err(|inner| ToolproofStepError::Internal(inner.into()))
+            }
+            BrowserWindow::Pagebrowse(window) => window
+                .navigate(url, wait_for_load)
+                .await
+                .map_err(|inner| ToolproofStepError::Internal(inner.into())),
+        }
+    }
+
+    async fn evaluate_script(
+        &self,
+        script: String,
+    ) -> Result<Option<serde_json::Value>, ToolproofStepError> {
+        match self {
+            BrowserWindow::Chrome(page) => {
+                let res = page
+                    .evaluate_function(format!("async function() {{{}}}", harnessed(script)))
+                    .await
+                    .map_err(|inner| ToolproofStepError::Internal(inner.into()))?;
+
+                Ok(res.object().value.clone())
+            }
+            BrowserWindow::Pagebrowse(window) => window
+                .evaluate_script(harnessed(script))
+                .await
+                .map_err(|inner| ToolproofStepError::Internal(inner.into())),
+        }
+    }
+}
 
 mod load_page {
     use super::*;
@@ -37,12 +160,15 @@ mod load_page {
                 args.get_string("url")?
             );
 
-            let window = civ.universe.pagebrowser.get_window().await.unwrap();
+            let browser = civ
+                .universe
+                .browser
+                .get_or_init(|| async { BrowserTester::initialize(&civ.universe.ctx.params).await })
+                .await;
 
-            window
-                .navigate(url.to_string(), true)
-                .await
-                .map_err(|inner| ToolproofStepError::Internal(inner.into()))?;
+            let window = browser.get_window().await;
+
+            window.navigate(url.to_string(), true).await?;
 
             civ.window = Some(window);
 
@@ -61,10 +187,6 @@ mod eval_js {
 
     use super::*;
 
-    fn harnessed(js: String) -> String {
-        HARNESS.replace("// insert_toolproof_inner_js", &js)
-    }
-
     async fn eval_and_return_js(
         js: String,
         civ: &mut Civilization<'_>,
@@ -77,10 +199,7 @@ mod eval_js {
             ));
         };
 
-        let value = window
-            .evaluate_script(harnessed(js))
-            .await
-            .map_err(|inner| ToolproofStepError::Internal(inner.into()))?;
+        let value = window.evaluate_script(js).await?;
 
         let Some(serde_json::Value::Object(map)) = &value else {
             return Err(ToolproofStepError::External(

toolproof/src/definitions/mod.rs

@@ -11,7 +11,7 @@ use crate::{
 };
 
 mod assertions;
-mod browser;
+pub mod browser;
 mod filesystem;
 mod hosting;
 mod process;

toolproof/src/definitions/process/mod.rs

@@ -61,7 +61,7 @@ mod run {
         ) -> Result<(), ToolproofStepError> {
             let command = args.get_string("command")?;
 
-            let exit_status = civ.run_command(command.to_string())?;
+            let exit_status = civ.run_command(command.to_string()).await?;
 
             if !exit_status.success() {
                 return Err(ToolproofTestFailure::Custom {
@@ -97,7 +97,7 @@ mod run {
         ) -> Result<(), ToolproofStepError> {
             let command = args.get_string("command")?;
 
-            let exit_status = civ.run_command(command.to_string())?;
+            let exit_status = civ.run_command(command.to_string()).await?;
 
             if exit_status.success() {
                 return Err(ToolproofTestFailure::Custom {

toolproof/src/errors.rs

@@ -1,5 +1,6 @@
 use std::path::PathBuf;
 
+use chromiumoxide::error::CdpError;
 use pagebrowse::PagebrowseError;
 use thiserror::Error;
 
@@ -45,6 +46,8 @@ pub enum ToolproofInternalError {
     Custom { msg: String },
     #[error("{0}")]
     PagebrowseError(#[from] PagebrowseError),
+    #[error("{0}")]
+    ChromeError(#[from] CdpError),
 }
 
 #[derive(Error, Debug)]