1 files changed, 20 insertions, 8 deletions
diff --git a/zoned/src/ws.rs b/zoned/src/ws.rs
index 7c49079..fc85d2e 100644
--- a/zoned/src/ws.rs
+++ b/zoned/src/ws.rs
@@ -11,11 +11,14 @@ use tokio::{
     io::{AsyncReadExt, AsyncWriteExt},
     sync::mpsc,
 };
-use tracing::{instrument, warn};
+use tracing::{error, instrument, warn};
 use wspty::PtyMaster;
 use zone_core::{Container, FilterContainer, Runtime};
 
-use crate::{error::PtyError, Error, Result, State};
+use crate::{
+    error::{PtyError, WebSocketError},
+    Error, Result, State,
+};
 
 #[derive(Deserialize, Debug)]
 struct WindowSize {
@@ -23,19 +26,28 @@ struct WindowSize {
     rows: u16,
 }
 
-// #[instrument(err, skip_all)]
-pub async fn handler(mut ws_stream: WebSocket, state: Arc<State>) -> Result<()> {
+pub async fn handler(ws_stream: WebSocket, state: Arc<State>) {
+    if let Err(err) = spawn(ws_stream, state).await {
+        error!("{}", err);
+        // if let Err(_) = ws_stream.send(Message::Text(format!("{}", err))).await {
+        //     error!("Failed to send WebSocket error message response: {}", err);
+        // }
+    }
+}
+
+#[instrument(err, skip_all)]
+pub async fn spawn(mut ws_stream: WebSocket, state: Arc<State>) -> Result<()> {
     let msg = loop {
         match ws_stream.next().await {
             Some(Ok(Message::Text(msg))) => break msg,
-            Some(Ok(_)) => return Err(Error::WebSocketMissingInit),
-            Some(Err(err)) => return Err(err.into()),
+            Some(Err(err)) => Err(Error::from(err)),
+            Some(_) => Err(WebSocketError::Initialization.into()),
             None => continue,
-        };
+        }?;
     };
 
     let container =
-        serde_json::from_str::<Container>(&msg).map_err(|_| Error::WebSocketMissingInit)?;
+        serde_json::from_str::<Container>(&msg).or(Err(WebSocketError::Initialization))?;
 
     state
         .nspawn