Initial commit

src/browser.rs

@@ -0,0 +1,135 @@
+use std::path::{Path, PathBuf};
+
+use inotify::{Inotify, WatchMask};
+use mozprofile::preferences::Pref;
+use mozprofile::profile::Profile;
+use mozrunner::runner::{
+    FirefoxProcess, FirefoxRunner, Runner, RunnerError, RunnerProcess,
+};
+use tokio_stream::StreamExt;
+use tracing::{debug, error, trace, warn};
+
+#[derive(Debug)]
+pub enum BrowserError {
+    Io(std::io::Error),
+    Runner(RunnerError),
+}
+
+impl std::fmt::Display for BrowserError {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        match self {
+            Self::Io(e) => write!(f, "{}", e),
+            Self::Runner(e) => write!(f, "{}", e),
+        }
+    }
+}
+
+impl std::error::Error for BrowserError {
+    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
+        match self {
+            Self::Io(e) => Some(e),
+            Self::Runner(e) => Some(e),
+        }
+    }
+}
+
+impl From<std::io::Error> for BrowserError {
+    fn from(e: std::io::Error) -> Self {
+        Self::Io(e)
+    }
+}
+
+impl From<RunnerError> for BrowserError {
+    fn from(e: RunnerError) -> Self {
+        Self::Runner(e)
+    }
+}
+
+pub struct Browser {
+    profile: PathBuf,
+    process: FirefoxProcess,
+}
+
+impl Browser {
+    pub fn launch() -> Result<Self, BrowserError> {
+        Self::launch_with_binary(None)
+    }
+
+    pub fn launch_with_binary(
+        binary: Option<&Path>,
+    ) -> Result<Self, BrowserError> {
+        let mut profile = Profile::new(None)?;
+        match profile.user_prefs() {
+            Ok(prefs) => {
+                prefs.insert("marionette.port", Pref::new(0));
+            }
+            Err(e) => {
+                error!("Could not get profile user prefs: {}", e);
+            }
+        }
+        let binary = match binary {
+            Some(b) => b,
+            None => Path::new("firefox"),
+        };
+        let profile_path = profile.path.clone();
+        let mut runner = FirefoxRunner::new(binary, Some(profile));
+        runner.arg("--marionette");
+        let process = runner.start()?;
+        Ok(Self {
+            profile: profile_path,
+            process,
+        })
+    }
+
+    pub fn marionette_port(&self) -> Option<u16> {
+        let mut path = self.profile.clone();
+        path.push("MarionetteActivePort");
+        let port = match std::fs::read_to_string(path) {
+            Ok(v) => match v.parse() {
+                Ok(v) => v,
+                Err(e) => {
+                    warn!("Failed to parse Marionette port: {}", e);
+                    return None;
+                }
+            },
+            Err(e) => {
+                warn!("Could not read active Marionette port: {}", e);
+                return None;
+            }
+        };
+        Some(port)
+    }
+
+    pub async fn wait_ready(&self) -> Result<(), std::io::Error> {
+        if self.marionette_port().is_some() {
+            debug!("Marionette port has already been assigned");
+            return Ok(());
+        }
+        let mut inotify = Inotify::init()?;
+        debug!("Starting inotify watch for {}", self.profile.display());
+        inotify.add_watch(self.profile.clone(), WatchMask::CREATE)?;
+        let mut buffer = [0; 1024];
+        let mut stream = inotify.event_stream(&mut buffer)?;
+        while let Some(evt) = stream.next().await {
+            trace!("inotify event: {:?}", evt);
+            if let Some(name) = evt?.name {
+                if name == "MarionetteActivePort" {
+                    debug!("Marionette port has been assigned");
+                    break;
+                }
+            }
+        }
+        Ok(())
+    }
+}
+
+impl Drop for Browser {
+    fn drop(&mut self) {
+        debug!("Stopping browser process");
+        if let Err(e) =
+            self.process.wait(std::time::Duration::from_millis(500))
+        {
+            error!("Error stopping browser process: {}", e);
+        }
+    }
+}

src/main.rs

@@ -0,0 +1,27 @@
+mod browser;
+mod marionette;
+mod session;
+
+use tokio::signal::unix::{self, SignalKind};
+use tracing::info;
+use tracing_subscriber::filter::EnvFilter;
+
+use session::Session;
+
+#[tokio::main(flavor = "current_thread")]
+async fn main() {
+    tracing_subscriber::fmt()
+        .with_env_filter(EnvFilter::from_default_env())
+        .with_writer(std::io::stderr)
+        .init();
+
+    let mut sig_term = unix::signal(SignalKind::terminate()).unwrap();
+    let mut sig_int = unix::signal(SignalKind::interrupt()).unwrap();
+
+    let session = Session::begin().await.unwrap();
+
+    tokio::select! {
+        _ = sig_term.recv() => info!("Received SIGTERM"),
+        _ = sig_int.recv() => info!("Received SIGINT"),
+    };
+}

src/marionette/error.rs

@@ -0,0 +1,58 @@
+use std::str::Utf8Error;
+use std::num::ParseIntError;
+
+#[derive(Debug)]
+pub enum MessageError {
+    Io(std::io::Error),
+    Parse(ParseIntError),
+    Utf8(Utf8Error),
+}
+
+impl From<std::io::Error> for MessageError {
+    fn from(e: std::io::Error) -> Self {
+        Self::Io(e)
+    }
+}
+
+impl From<ParseIntError> for MessageError {
+    fn from(e: ParseIntError) -> Self {
+        Self::Parse(e)
+    }
+}
+
+impl From<Utf8Error> for MessageError {
+    fn from(e: Utf8Error) -> Self {
+        Self::Utf8(e)
+    }
+}
+
+#[derive(Debug)]
+pub enum HandshakeError {
+    Io(std::io::Error),
+    Parse(ParseIntError),
+    Utf8(Utf8Error),
+    Json(serde_json::Error),
+}
+
+impl From<MessageError> for HandshakeError {
+    fn from(e: MessageError) -> Self {
+        match e {
+            MessageError::Io(e) => Self::Io(e),
+            MessageError::Parse(e) => Self::Parse(e),
+            MessageError::Utf8(e) => Self::Utf8(e),
+        }
+    }
+}
+
+impl From<std::io::Error> for HandshakeError {
+    fn from(e: std::io::Error) -> Self {
+        Self::Io(e)
+    }
+}
+
+impl From<serde_json::Error> for HandshakeError {
+    fn from(e: serde_json::Error) -> Self {
+        Self::Json(e)
+    }
+}
+

src/marionette/message.rs

@@ -0,0 +1,70 @@
+use serde::{Deserialize, Serialize};
+
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct Hello {
+    pub application_type: String,
+    pub marionette_protocol: u16,
+}
+
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase")]
+#[allow(dead_code)]
+pub struct Timeouts {
+    implicit: u32,
+    page_load: u32,
+    script: u32,
+}
+
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase")]
+#[allow(dead_code)]
+pub struct Capabilities {
+    pub accept_insecure_certs: bool,
+    pub browser_name: String,
+    pub browser_version: String,
+    #[serde(rename = "moz:accessibilityChecks")]
+    pub accessibility_checks: bool,
+    #[serde(rename = "moz:buildID")]
+    pub build_id: String,
+    #[serde(rename = "moz:headless")]
+    pub headless: bool,
+    #[serde(rename = "moz:processID")]
+    pub process_id: u128,
+    #[serde(rename = "moz:profile")]
+    pub profile: String,
+    #[serde(rename = "moz:shutdownTimeout")]
+    pub shutdown_timeout: u32,
+    #[serde(rename = "moz:useNonSpecCompliantPointerOrigin")]
+    pub use_non_spec_compliant_pointer_origin: bool,
+    #[serde(rename = "moz:webdriverClick")]
+    pub webdriver_click: bool,
+    pub page_load_strategy: String,
+    pub platform_name: String,
+    pub platform_version: String,
+    // proxy:
+    pub set_window_rect: bool,
+    pub timeouts: Timeouts,
+    pub unhandled_prompt_behavior: String,
+}
+
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase")]
+#[allow(dead_code)]
+pub struct NewSessionResponse {
+    pub capabilities: Capabilities,
+    pub session_id: String,
+}
+
+#[derive(Debug, Serialize)]
+#[serde(rename_all = "camelCase")]
+pub struct NewSessionParams {
+    pub strict_file_interactability: bool,
+}
+
+#[derive(Debug, Serialize)]
+#[serde(tag = "command", content = "params")]
+pub enum Command {
+    #[serde(rename = "WebDriver:NewSession")]
+    NewSession(NewSessionParams),
+}

src/marionette/mod.rs

@@ -0,0 +1,78 @@
+pub mod error;
+pub mod message;
+
+use std::time::Instant;
+
+use serde::{Deserialize, Serialize};
+use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufStream};
+use tokio::net::{TcpStream, ToSocketAddrs};
+use tracing::{debug, trace};
+
+pub use error::{HandshakeError, MessageError};
+use message::{Command, Hello, NewSessionParams, NewSessionResponse};
+
+#[derive(Debug, Deserialize, Serialize)]
+struct Message(u8, u32, Option<String>, serde_json::Value);
+
+pub struct Marionette {
+    ts: Instant,
+    stream: BufStream<TcpStream>,
+}
+
+impl Marionette {
+    pub async fn connect<A>(addr: A) -> Result<Self, std::io::Error>
+    where
+        A: ToSocketAddrs,
+    {
+        let conn = TcpStream::connect(addr).await?;
+        let stream = BufStream::new(conn);
+        let ts = Instant::now();
+        Ok(Self { ts, stream })
+    }
+
+    pub async fn handshake(&mut self) -> Result<(), HandshakeError> {
+        let buf = self.next_message().await?;
+        let hello: Hello = serde_json::from_slice(&buf)?;
+        debug!("Received hello: {:?}", hello);
+        self.send_message(Command::NewSession(NewSessionParams {
+            strict_file_interactability: true,
+        }))
+        .await?;
+        let buf = self.next_message().await?;
+        let msg: Message = serde_json::from_slice(&buf)?;
+        let res: NewSessionResponse = serde_json::from_value(msg.3)?;
+        debug!("Received message: {:?}", res);
+        Ok(())
+    }
+
+    async fn next_message(&mut self) -> Result<Vec<u8>, MessageError> {
+        let mut buf = vec![];
+        self.stream.read_until(b':', &mut buf).await?;
+        let length: usize =
+            std::str::from_utf8(&buf[..buf.len() - 1])?.parse()?;
+        trace!("Message length: {:?}", length);
+        let mut buf = vec![0; length];
+        self.stream.read_exact(&mut buf[..]).await?;
+        trace!("Received message: {:?}", buf);
+        Ok(buf)
+    }
+
+    async fn send_message(
+        &mut self,
+        command: Command,
+    ) -> Result<(), std::io::Error> {
+        let value = serde_json::to_value(command)?;
+        let (command, params) = (
+            value.get("command").unwrap().as_str().unwrap().into(),
+            value.get("params").unwrap().clone(),
+        );
+        let msgid = (self.ts.elapsed().as_millis() % u32::MAX as u128) as u32;
+        let message = Message(0, msgid, Some(command), params);
+        let message = serde_json::to_string(&message)?;
+        let message = format!("{}:{}", message.len(), message);
+        trace!("Sending message: {}", message);
+        self.stream.write_all(message.as_bytes()).await?;
+        self.stream.flush().await?;
+        Ok(())
+    }
+}

src/session.rs

@@ -0,0 +1,57 @@
+use tracing::{debug, info};
+
+use crate::browser::{Browser, BrowserError};
+use crate::marionette::Marionette;
+use crate::marionette::error::HandshakeError;
+
+#[derive(Debug)]
+pub enum SessionError {
+    Browser(BrowserError),
+    Io(std::io::Error),
+    Handshake(HandshakeError),
+    InvalidState(String),
+}
+
+impl From<BrowserError> for SessionError {
+    fn from(e: BrowserError) -> Self {
+        Self::Browser(e)
+    }
+}
+
+impl From<std::io::Error> for SessionError {
+    fn from(e: std::io::Error) -> Self {
+        Self::Io(e)
+    }
+}
+
+impl From<HandshakeError> for SessionError {
+    fn from(e: HandshakeError) -> Self {
+        Self::Handshake(e)
+    }
+}
+
+
+pub struct Session {
+    browser: Browser,
+    marionette: Marionette,
+}
+
+impl Session {
+    pub async fn begin() -> Result<Self, SessionError> {
+        debug!("Launching Firefox");
+        let browser = Browser::launch()?;
+        browser.wait_ready().await?;
+        debug!("Firefox Marionette is now ready");
+        let Some(port) = browser.marionette_port() else {
+            return Err(SessionError::InvalidState("No active Marionette port".into()));
+        };
+        debug!("Connecting to Firefox Marionette on port {}", port);
+        let mut marionette = Marionette::connect(("127.0.0.1", port)).await?;
+        info!("Successfully connected to Firefox Marionette");
+        marionette.handshake().await?;
+        Ok(Self {
+            browser,
+            marionette,
+        })
+    }
+}