From 6e8ba8b7b6ad215c18d9e6b57f12390d7dd3fe7e Mon Sep 17 00:00:00 2001
From: "Dustin C. Hatch" <dustin@hatch.name>
Date: Sat, 7 Jan 2023 22:28:15 -0600
Subject: [PATCH] marionette: Handle unexpected disconnects

If the Marionette server closes the connection unexpectedly, it will
manifest as an empty message from the socket.  We need to handle this
and return an error, instead of causing a panic by attempting to read
from the empty buffer (by trying to create a slice with a negative
length).
---
 src/marionette/error.rs | 3 +++
 src/marionette/mod.rs   | 3 +++
 2 files changed, 6 insertions(+)

diff --git a/src/marionette/error.rs b/src/marionette/error.rs
index 17e019c..f1fccc4 100644
--- a/src/marionette/error.rs
+++ b/src/marionette/error.rs
@@ -8,6 +8,7 @@ pub enum MessageError {
     Io(std::io::Error),
     Parse(ParseIntError),
     Utf8(Utf8Error),
+    Disconnected,
 }
 
 impl From<std::io::Error> for MessageError {
@@ -34,6 +35,7 @@ impl std::fmt::Display for MessageError {
             Self::Io(e) => write!(f, "I/O error: {}", e),
             Self::Parse(e) => write!(f, "Error parsing message: {}", e),
             Self::Utf8(e) => write!(f, "Error parsing message: {}", e),
+            Self::Disconnected => write!(f, "Disconnected"),
         }
     }
 }
@@ -44,6 +46,7 @@ impl std::error::Error for MessageError {
             Self::Io(e) => Some(e),
             Self::Parse(e) => Some(e),
             Self::Utf8(e) => Some(e),
+            Self::Disconnected => None,
         }
     }
 }
diff --git a/src/marionette/mod.rs b/src/marionette/mod.rs
index 0ccbcf7..f044c83 100644
--- a/src/marionette/mod.rs
+++ b/src/marionette/mod.rs
@@ -150,6 +150,9 @@ impl MarionetteConnection {
     {
         let mut buf = vec![];
         stream.read_until(b':', &mut buf).await?;
+        if buf.is_empty() {
+            return Err(MessageError::Disconnected);
+        }
         let length: usize =
             std::str::from_utf8(&buf[..buf.len() - 1])?.parse()?;
         trace!("Message length: {:?}", length);