From 8431a33f20806f553c6af22b65047b7edb1a66c9 Mon Sep 17 00:00:00 2001
From: "Dustin C. Hatch" <dustin@hatch.name>
Date: Sat, 31 Dec 2022 10:26:11 -0600
Subject: [PATCH] marionette: Separate commands from connection

The `MarionetteConnection` structure provides the low-level interface to
communicate with the Marionette server via the TCP socket.  In contrast,
the `Marionette` structure provides the high-level interface to execute
Marionette commands.  Separating these will make the code a bit cleaner,
in my opinion.
---
 src/marionette/mod.rs | 82 +++++++++++++++++++++++++------------------
 src/session.rs        |  5 +--
 2 files changed, 50 insertions(+), 37 deletions(-)

diff --git a/src/marionette/mod.rs b/src/marionette/mod.rs
index 11dacf8..a8f049e 100644
--- a/src/marionette/mod.rs
+++ b/src/marionette/mod.rs
@@ -37,13 +37,13 @@ type CommandResult = Result<Option<serde_json::Value>, ErrorResponse>;
 
 type SenderMap = HashMap<u32, oneshot::Sender<CommandResult>>;
 
-pub struct Marionette {
+pub struct MarionetteConnection {
     ts: Instant,
     stream: BufWriter<OwnedWriteHalf>,
     sender: Arc<Mutex<SenderMap>>,
 }
 
-impl Marionette {
+impl MarionetteConnection {
     pub async fn connect<A>(addr: A) -> Result<Self, ConnectionError>
     where
         A: ToSocketAddrs,
@@ -61,39 +61,6 @@ impl Marionette {
         Ok(Self { ts, stream, sender })
     }
 
-    pub async fn get_title(&mut self) -> Result<String, CommandError> {
-        let res: GetTitleResponse =
-            self.send_message(Command::GetTitle).await?.unwrap();
-        debug!("Received message: {:?}", res);
-        Ok(res.value)
-    }
-
-    pub async fn navigate<U>(&mut self, url: U) -> Result<(), CommandError>
-    where
-        U: Into<String>,
-    {
-        let res: Option<serde_json::Value> = self
-            .send_message(Command::Navigate(NavigateParams {
-                url: url.into(),
-            }))
-            .await?;
-        debug!("Received message: {:?}", res);
-        Ok(())
-    }
-
-    pub async fn new_session(
-        &mut self,
-    ) -> Result<NewSessionResponse, CommandError> {
-        let res = self
-            .send_message(Command::NewSession(NewSessionParams {
-                strict_file_interactability: true,
-            }))
-            .await?
-            .unwrap();
-        debug!("Received message: {:?}", res);
-        Ok(res)
-    }
-
     pub async fn send_message<T>(
         &mut self,
         command: Command,
@@ -190,3 +157,48 @@ impl Marionette {
         Ok(buf)
     }
 }
+
+pub struct Marionette {
+    conn: MarionetteConnection,
+}
+
+impl Marionette {
+    pub fn new(conn: MarionetteConnection) -> Self {
+        Self { conn }
+    }
+
+    pub async fn get_title(&mut self) -> Result<String, CommandError> {
+        let res: GetTitleResponse =
+            self.conn.send_message(Command::GetTitle).await?.unwrap();
+        debug!("Received message: {:?}", res);
+        Ok(res.value)
+    }
+
+    pub async fn navigate<U>(&mut self, url: U) -> Result<(), CommandError>
+    where
+        U: Into<String>,
+    {
+        let res: Option<serde_json::Value> = self
+            .conn
+            .send_message(Command::Navigate(NavigateParams {
+                url: url.into(),
+            }))
+            .await?;
+        debug!("Received message: {:?}", res);
+        Ok(())
+    }
+
+    pub async fn new_session(
+        &mut self,
+    ) -> Result<NewSessionResponse, CommandError> {
+        let res = self
+            .conn
+            .send_message(Command::NewSession(NewSessionParams {
+                strict_file_interactability: true,
+            }))
+            .await?
+            .unwrap();
+        debug!("Received message: {:?}", res);
+        Ok(res)
+    }
+}
diff --git a/src/session.rs b/src/session.rs
index faaf1a9..59820de 100644
--- a/src/session.rs
+++ b/src/session.rs
@@ -5,7 +5,7 @@ use tracing::{debug, error, info, warn};
 use crate::browser::{Browser, BrowserError};
 use crate::config::Configuration;
 use crate::marionette::error::{CommandError, ConnectionError};
-use crate::marionette::Marionette;
+use crate::marionette::{Marionette, MarionetteConnection};
 use crate::mqtt::{Message, MqttClient, MqttPublisher};
 
 #[derive(Debug)]
@@ -81,7 +81,8 @@ impl Session {
             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?;
+        let conn = MarionetteConnection::connect(("127.0.0.1", port)).await?;
+        let mut marionette = Marionette::new(conn);
         info!("Successfully connected to Firefox Marionette");
         let ses = marionette.new_session().await?;
         debug!("Started Marionette session {}", ses.session_id);