Properly handle websocket close during handshake

src/client/connection.rs

@@ -2,6 +2,7 @@ use std::collections::{HashMap, VecDeque};
 
 use futures_util::{Sink, SinkExt, Stream, StreamExt};
 use tokio::sync::{oneshot, Mutex};
+pub use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode;
 use tokio_tungstenite::tungstenite::Message;
 use tracing::debug;
 
@@ -86,7 +87,7 @@ impl ReidentifyReceiverList {
 pub enum HandshakeError {
     /// The connection to obs-websocket was interrupted while trying to read a message.
     #[error("connection to obs-websocket was closed")]
-    ConnectionClosed,
+    ConnectionClosed(Option<CloseDetails>),
     /// Receiving a message did not succeed.
     #[error("failed reading websocket message")]
     Receive(#[source] tokio_tungstenite::tungstenite::Error),
@@ -110,6 +111,19 @@ pub enum HandshakeError {
     NoIdentified,
 }
 
+/// Description about the reason of why the web-socket connection was closed.
+#[derive(Debug)]
+pub struct CloseDetails {
+    /// Close code to precisely identify the reason.
+    ///
+    /// This can be turned into a [`u16`] and compared against the
+    /// [`WebSocketCloseCode`](crate::responses::WebSocketCloseCode) to further identify the close
+    /// reason, if related to `obs-websocket`.
+    pub code: CloseCode,
+    /// Textual representation of the close code, or additional details for it.
+    pub reason: String,
+}
+
 pub(super) async fn handshake(
     write: &mut (impl Sink<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin),
     read: &mut (impl Stream<Item = tokio_tungstenite::tungstenite::Result<Message>> + Unpin),
@@ -119,13 +133,22 @@ pub(super) async fn handshake(
     async fn read_message(
         read: &mut (impl Stream<Item = tokio_tungstenite::tungstenite::Result<Message>> + Unpin),
     ) -> Result<ServerMessage, HandshakeError> {
-        let message = read
+        let mut message = read
             .next()
             .await
-            .ok_or(HandshakeError::ConnectionClosed)?
-            .map_err(HandshakeError::Receive)?
-            .into_text()
-            .map_err(HandshakeError::IntoText)?;
+            .ok_or(HandshakeError::ConnectionClosed(None))?
+            .map_err(HandshakeError::Receive)?;
+
+        if let Message::Close(info) = &mut message {
+            return Err(HandshakeError::ConnectionClosed(info.take().map(|i| {
+                CloseDetails {
+                    code: i.code,
+                    reason: i.reason.into_owned(),
+                }
+            })));
+        }
+
+        let message = message.into_text().map_err(HandshakeError::IntoText)?;
 
         serde_json::from_str::<ServerMessage>(&message).map_err(HandshakeError::DeserializeMessage)
     }

src/responses/mod.rs

@@ -260,3 +260,40 @@ pub enum StatusCode {
     /// The combination of request fields cannot be used to perform an action.
     CannotAct = 703,
 }
+
+/// Additional close codes, defined by `obs-websocket`. These can be used to further pin down the
+/// details of why the web-socket connection was closed.
+#[derive(Debug)]
+#[repr(u16)]
+pub enum WebSocketCloseCode {
+    /// Unknown reason, should never be used.
+    UnknownReason = 4000,
+    /// The server was unable to decode the incoming websocket message.
+    MessageDecodeError = 4002,
+    /// A data field is required but missing from the payload.
+    MissingDataField = 4003,
+    /// A data field's value type is invalid.
+    InvalidDataFieldType = 4004,
+    /// A data field's value is invalid.
+    InvalidDataFieldValue = 4005,
+    /// The specified `op` was invalid or missing.
+    UnknownOpCode = 4006,
+    /// The client sent a websocket message without first sending `Identify` message.
+    NotIdentified = 4007,
+    /// The client sent an `Identify` message while already identified.
+    ///
+    /// **Note:** Once a client has identified, only `Reidentify` may be used to change session
+    /// parameters.
+    AlreadyIdentified = 4008,
+    /// The authentication attempt (via `Identify`) failed.
+    AuthenticationFailed = 4009,
+    /// The server detected the usage of an old version of the obs-websocket RPC protocol.
+    UnsupportedRpcVersion = 4010,
+    /// The websocket session has been invalidated by the obs-websocket server.
+    ///
+    /// **Note:** This is the code used by the `Kick` button in the UI Session List. If you receive
+    /// this code, you must not automatically reconnect.
+    SessionInvalidated = 4011,
+    /// A requested feature is not supported due to hardware/software limitations.
+    UnsupportedFeature = 4012,
+}