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).
dev/ci
Dustin 2023-01-07 22:28:15 -06:00
parent e92d8c9cef
commit 6e8ba8b7b6
2 changed files with 6 additions and 0 deletions

View File

@ -8,6 +8,7 @@ pub enum MessageError {
Io(std::io::Error), Io(std::io::Error),
Parse(ParseIntError), Parse(ParseIntError),
Utf8(Utf8Error), Utf8(Utf8Error),
Disconnected,
} }
impl From<std::io::Error> for MessageError { 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::Io(e) => write!(f, "I/O error: {}", e),
Self::Parse(e) => write!(f, "Error parsing message: {}", e), Self::Parse(e) => write!(f, "Error parsing message: {}", e),
Self::Utf8(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::Io(e) => Some(e),
Self::Parse(e) => Some(e), Self::Parse(e) => Some(e),
Self::Utf8(e) => Some(e), Self::Utf8(e) => Some(e),
Self::Disconnected => None,
} }
} }
} }

View File

@ -150,6 +150,9 @@ impl MarionetteConnection {
{ {
let mut buf = vec![]; let mut buf = vec![];
stream.read_until(b':', &mut buf).await?; stream.read_until(b':', &mut buf).await?;
if buf.is_empty() {
return Err(MessageError::Disconnected);
}
let length: usize = let length: usize =
std::str::from_utf8(&buf[..buf.len() - 1])?.parse()?; std::str::from_utf8(&buf[..buf.len() - 1])?.parse()?;
trace!("Message length: {:?}", length); trace!("Message length: {:?}", length);