Merge pull request #778 from hyperware-ai/hf/dont-crash-on-web-socket-close

dont crash on web socket close

Author: nick1udwig

Cargo.lock

@@ -1,8 +1,9 @@
 use crate::http::server_types::{
-    HttpResponse, HttpServerAction, HttpServerError, HttpServerRequest, IncomingHttpRequest,
-    MessageType, RpcResponseBody, WsMessageType,
+    HttpResponse, HttpResponseSenders, HttpSender, HttpServerAction, HttpServerError,
+    HttpServerRequest, IncomingHttpRequest, MessageType, RpcResponseBody, WebSocketSender,
+    WebSocketSenders, WsMessageType,
 };
-use crate::http::utils;
+use crate::http::utils::{self, send_action_response};
 use crate::keygen;
 use base64::{engine::general_purpose::STANDARD as base64_standard, Engine};
 use dashmap::DashMap;
@@ -39,18 +40,6 @@ const WS_SELF_IMPOSED_MAX_CONNECTIONS: u32 = 128;
 
 const LOGIN_HTML: &str = include_str!("login.html");
 
-/// mapping from a given HTTP request (assigned an ID) to the oneshot
-/// channel that will get a response from the app that handles the request,
-/// and a string which contains the path that the request was made to.
-type HttpResponseSenders = Arc<DashMap<u64, (String, HttpSender)>>;
-type HttpSender = tokio::sync::oneshot::Sender<(HttpResponse, Vec<u8>)>;
-
-/// mapping from an open websocket connection to a channel that will ingest
-/// WebSocketPush messages from the app that handles the connection, and
-/// send them to the connection.
-type WebSocketSenders = Arc<DashMap<u32, (ProcessId, WebSocketSender)>>;
-type WebSocketSender = tokio::sync::mpsc::Sender<warp::ws::Message>;
-
 type PathBindings = Arc<RwLock<Router<BoundPath>>>;
 type WsPathBindings = Arc<RwLock<Router<BoundWsPath>>>;
 type SecureSubdomains = Arc<RwLock<HashSet<String>>>;
@@ -121,7 +110,14 @@ async fn send_push(
             }
         }
         WsMessageType::Close => {
-            unreachable!();
+            return utils::handle_close_websocket(
+                id,
+                &source,
+                send_to_loop,
+                ws_senders,
+                channel_id,
+            )
+            .await;
         }
     };
     // Send to the websocket if registered
@@ -1546,19 +1542,17 @@ async fn handle_app_message(
                     return;
                 }
                 HttpServerAction::WebSocketClose(channel_id) => {
-                    if let Some(got) = ws_senders.get(&channel_id) {
-                        if got.value().0 != km.source.process {
-                            send_action_response(
-                                km.id,
-                                km.source,
-                                &send_to_loop,
-                                Err(HttpServerError::WsChannelNotFound),
-                            )
-                            .await;
-                            return;
-                        }
-                        let _ = got.value().1.send(warp::ws::Message::close()).await;
-                        ws_senders.remove(&channel_id);
+                    let is_return = utils::handle_close_websocket(
+                        km.id,
+                        &km.source,
+                        &send_to_loop,
+                        ws_senders,
+                        channel_id,
+                    )
+                    .await;
+
+                    if is_return {
+                        return;
                     }
                 }
             }
@@ -1569,28 +1563,3 @@ async fn handle_app_message(
         }
     }
 }
-
-pub async fn send_action_response(
-    id: u64,
-    target: Address,
-    send_to_loop: &MessageSender,
-    result: Result<(), HttpServerError>,
-) {
-    KernelMessage::builder()
-        .id(id)
-        .source(("our", HTTP_SERVER_PROCESS_ID.clone()))
-        .target(target)
-        .message(Message::Response((
-            Response {
-                inherit: false,
-                body: serde_json::to_vec(&result).unwrap(),
-                metadata: None,
-                capabilities: vec![],
-            },
-            None,
-        )))
-        .build()
-        .unwrap()
-        .send(send_to_loop)
-        .await;
-}

hyperdrive/src/http/server.rs

@@ -1,6 +1,12 @@
+use crate::http::server_types::{HttpServerError, WebSocketSenders};
 use hmac::{Hmac, Mac};
 use jwt::VerifyWithKey;
-use lib::{core::ProcessId, types::http_server};
+use lib::types::{
+    core::{
+        Address, KernelMessage, Message, MessageSender, ProcessId, Response, HTTP_SERVER_PROCESS_ID,
+    },
+    http_server,
+};
 use serde::{Deserialize, Serialize};
 use sha2::Sha256;
 use std::collections::HashMap;
@@ -158,3 +164,56 @@ pub fn is_behind_reverse_proxy(headers: &warp::http::HeaderMap) -> bool {
     }
     return false;
 }
+
+pub async fn handle_close_websocket(
+    id: u64,
+    source: &Address,
+    send_to_loop: &MessageSender,
+    ws_senders: WebSocketSenders,
+    channel_id: u32,
+) -> bool {
+    let Some(got) = ws_senders.get(&channel_id) else {
+        return false;
+    };
+
+    if got.value().0 != source.process {
+        send_action_response(
+            id,
+            source.clone(),
+            send_to_loop,
+            Err(HttpServerError::WsChannelNotFound),
+        )
+        .await;
+        return true;
+    }
+
+    let _ = got.value().1.send(warp::ws::Message::close()).await;
+    ws_senders.remove(&channel_id);
+
+    return false;
+}
+
+pub async fn send_action_response(
+    id: u64,
+    target: Address,
+    send_to_loop: &MessageSender,
+    result: Result<(), HttpServerError>,
+) {
+    KernelMessage::builder()
+        .id(id)
+        .source(("our", HTTP_SERVER_PROCESS_ID.clone()))
+        .target(target)
+        .message(Message::Response((
+            Response {
+                inherit: false,
+                body: serde_json::to_vec(&result).unwrap(),
+                metadata: None,
+                capabilities: vec![],
+            },
+            None,
+        )))
+        .build()
+        .unwrap()
+        .send(send_to_loop)
+        .await;
+}

hyperdrive/src/http/utils.rs

@@ -23,6 +23,7 @@ alloy = { version = "0.8.1", features = [
     "rpc-types",
     "rpc-types-eth",
 ] }
+dashmap = "5.5.3"
 lazy_static = "1.4.0"
 rand = "0.8.4"
 regex = "1.11.0"
@@ -32,4 +33,5 @@ serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 thiserror = "1.0"
 tokio = { version = "1.28", features = ["sync"] }
+warp = "0.3.5"
 wasmtime = { version = "33.0.0", features = ["component-model"] }

lib/Cargo.toml

@@ -1,6 +1,8 @@
-use crate::core::LazyLoadBlob;
+use crate::core::{LazyLoadBlob, ProcessId};
+use dashmap::DashMap;
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
+use std::sync::Arc;
 use thiserror::Error;
 
 /// HTTP Request received from the `http-server:distro:sys` service as a
@@ -201,3 +203,15 @@ pub struct JwtClaims {
     pub subdomain: Option<String>,
     pub expiration: u64,
 }
+
+/// mapping from a given HTTP request (assigned an ID) to the oneshot
+/// channel that will get a response from the app that handles the request,
+/// and a string which contains the path that the request was made to.
+pub type HttpResponseSenders = Arc<DashMap<u64, (String, HttpSender)>>;
+pub type HttpSender = tokio::sync::oneshot::Sender<(HttpResponse, Vec<u8>)>;
+
+/// mapping from an open websocket connection to a channel that will ingest
+/// WebSocketPush messages from the app that handles the connection, and
+/// send them to the connection.
+pub type WebSocketSenders = Arc<DashMap<u32, (ProcessId, WebSocketSender)>>;
+pub type WebSocketSender = tokio::sync::mpsc::Sender<warp::ws::Message>;