Fix ws-upstream relay: defer upstream connection to main runtime

The host_ws_upgrade host function was connecting to the upstream
WebSocket on a temporary single-threaded tokio runtime that was
destroyed immediately after block_on returned. The TcpStream inside
the WebSocket was bound to that dead runtime's I/O driver, so all
reads/writes in the relay task failed instantly (close code 1006).

The fix stores the WsUpgradeRequest in PluginState and defers the
actual TCP connection to the main async context where the relay
task runs, ensuring the I/O driver stays alive for the lifetime
of the connection.

Author: ndreno

Cargo.lock

@@ -18,6 +18,8 @@ rcgen = { workspace = true }
 rustls = { workspace = true }
 base64 = { workspace = true }
 wiremock = "0.6"
+tokio-tungstenite = { workspace = true }
+futures-util = { workspace = true }
 assert_cmd = { workspace = true }
 predicates = { workspace = true }
 

crates/barbacane-test/Cargo.toml

@@ -1234,3 +1234,195 @@ async fn test_ws_upstream_upstream_unavailable() {
         resp.status()
     );
 }
+
+/// Start a simple WebSocket echo server on a random port.
+/// Returns the `(join_handle, "ws://127.0.0.1:PORT")`.
+fn start_ws_echo_server() -> (tokio::task::JoinHandle<()>, String) {
+    let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
+    let port = listener.local_addr().unwrap().port();
+    let url = format!("ws://127.0.0.1:{}", port);
+
+    // Convert to a tokio TcpListener inside the spawned task
+    listener.set_nonblocking(true).unwrap();
+
+    let handle = tokio::spawn(async move {
+        use futures_util::{SinkExt, StreamExt};
+
+        let listener = tokio::net::TcpListener::from_std(listener).unwrap();
+        while let Ok((stream, _)) = listener.accept().await {
+            tokio::spawn(async move {
+                let ws = tokio_tungstenite::accept_async(stream).await.unwrap();
+                let (mut tx, mut rx) = ws.split();
+                while let Some(Ok(msg)) = rx.next().await {
+                    if msg.is_close() {
+                        break;
+                    }
+                    if msg.is_text() || msg.is_binary() {
+                        if tx.send(msg).await.is_err() {
+                            break;
+                        }
+                    }
+                }
+            });
+        }
+    });
+
+    (handle, url)
+}
+
+/// Create a temporary spec + manifest for ws-upstream relay tests.
+fn create_ws_spec(upstream_url: &str) -> (tempfile::TempDir, std::path::PathBuf) {
+    let temp_dir = tempfile::TempDir::new().expect("failed to create temp dir");
+    let spec_path = temp_dir.path().join("ws-relay.yaml");
+
+    let ws_wasm = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR"))
+        .parent()
+        .unwrap()
+        .parent()
+        .unwrap()
+        .join("plugins/ws-upstream/ws-upstream.wasm");
+    let mock_wasm = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR"))
+        .parent()
+        .unwrap()
+        .parent()
+        .unwrap()
+        .join("plugins/mock/mock.wasm");
+
+    let manifest = format!(
+        "plugins:\n  ws-upstream:\n    path: {ws}\n  mock:\n    path: {mock}\n",
+        ws = ws_wasm.display(),
+        mock = mock_wasm.display(),
+    );
+    std::fs::write(temp_dir.path().join("barbacane.yaml"), manifest).expect("write manifest");
+
+    let spec = format!(
+        r#"openapi: "3.1.0"
+info:
+  title: WS Relay Test
+  version: "1.0.0"
+
+paths:
+  /health:
+    get:
+      operationId: health
+      x-barbacane-dispatch:
+        name: mock
+        config:
+          status: 200
+          body: '{{"ok":true}}'
+      responses:
+        "200":
+          description: OK
+
+  /ws/echo:
+    get:
+      operationId: wsEcho
+      x-barbacane-dispatch:
+        name: ws-upstream
+        config:
+          url: "{upstream}"
+      responses:
+        "101":
+          description: Switching Protocols
+"#,
+        upstream = upstream_url,
+    );
+    std::fs::write(&spec_path, spec).expect("write spec");
+
+    (temp_dir, spec_path)
+}
+
+#[tokio::test]
+async fn test_ws_relay_stays_alive_and_echoes_frames() {
+    use futures_util::{SinkExt, StreamExt};
+    use tokio_tungstenite::tungstenite::Message;
+
+    // Start a WS echo server
+    let (_echo_handle, echo_url) = start_ws_echo_server();
+
+    // Create spec pointing to the echo server and boot the gateway
+    let (_temp_dir, spec_path) = create_ws_spec(&echo_url);
+    let gateway = TestGateway::from_spec(spec_path.to_str().unwrap())
+        .await
+        .expect("failed to start gateway");
+
+    // Connect through the gateway
+    let ws_url = format!("ws://127.0.0.1:{}/ws/echo", gateway.port());
+    let (mut ws, _) = tokio_tungstenite::connect_async(&ws_url)
+        .await
+        .expect("WebSocket connection through gateway failed");
+
+    // Send multiple text frames and verify each echoes back
+    for i in 0..5 {
+        let msg = format!("hello {}", i);
+        ws.send(Message::Text(msg.clone().into()))
+            .await
+            .expect("send failed");
+
+        let reply = tokio::time::timeout(std::time::Duration::from_secs(5), ws.next())
+            .await
+            .expect("timed out waiting for echo")
+            .expect("stream ended")
+            .expect("read error");
+
+        assert_eq!(reply, Message::Text(msg.into()), "frame {} did not echo", i);
+    }
+
+    // Send a binary frame too
+    ws.send(Message::Binary(vec![1, 2, 3].into()))
+        .await
+        .expect("binary send failed");
+
+    let reply = tokio::time::timeout(std::time::Duration::from_secs(5), ws.next())
+        .await
+        .expect("timed out waiting for binary echo")
+        .expect("stream ended")
+        .expect("read error");
+
+    assert_eq!(reply, Message::Binary(vec![1, 2, 3].into()));
+
+    // Clean close
+    ws.close(None).await.expect("close failed");
+}
+
+#[tokio::test]
+async fn test_ws_relay_upstream_close_propagates_to_client() {
+    use futures_util::{SinkExt, StreamExt};
+    use tokio_tungstenite::tungstenite::Message;
+
+    let (_echo_handle, echo_url) = start_ws_echo_server();
+    let (_temp_dir, spec_path) = create_ws_spec(&echo_url);
+    let gateway = TestGateway::from_spec(spec_path.to_str().unwrap())
+        .await
+        .expect("failed to start gateway");
+
+    let ws_url = format!("ws://127.0.0.1:{}/ws/echo", gateway.port());
+    let (mut ws, _) = tokio_tungstenite::connect_async(&ws_url)
+        .await
+        .expect("WebSocket connection through gateway failed");
+
+    // Verify the connection works
+    ws.send(Message::Text("ping".into()))
+        .await
+        .expect("send failed");
+
+    let reply = tokio::time::timeout(std::time::Duration::from_secs(5), ws.next())
+        .await
+        .expect("timed out")
+        .expect("stream ended")
+        .expect("read error");
+
+    assert_eq!(reply, Message::Text("ping".into()));
+
+    // Send close and verify the stream ends cleanly
+    ws.send(Message::Close(None)).await.expect("close failed");
+
+    // The next read should be a close frame, end-of-stream, or a reset
+    // (the proxy may drop the connection without a close handshake).
+    let final_msg = tokio::time::timeout(std::time::Duration::from_secs(5), ws.next()).await;
+    match final_msg {
+        Ok(Some(Ok(Message::Close(_)))) | Ok(None) => {} // Clean close
+        Ok(Some(Err(_))) => {}                           // Reset (proxy dropped connection)
+        other => panic!("expected close or end-of-stream, got {:?}", other),
+    }
+}

crates/barbacane-test/tests/plugins.rs

@@ -137,11 +137,13 @@ pub struct PluginState {
     /// each body chunk, then drops the sender to signal end-of-stream.
     pub stream_sender: Option<Arc<tokio::sync::mpsc::UnboundedSender<StreamEvent>>>,
 
-    /// Upstream WebSocket connection established by `host_ws_upgrade` (ADR-0026).
+    /// Upstream WebSocket upgrade request from `host_ws_upgrade` (ADR-0026).
     ///
-    /// After a successful `host_ws_upgrade`, the connected stream is stored here
-    /// for the data plane to pick up and relay frames bidirectionally.
-    pub ws_upstream: Option<crate::ws_client::UpstreamWsStream>,
+    /// After a successful `host_ws_upgrade`, the request params are stored here.
+    /// The actual connection is deferred to the async relay task on the main
+    /// runtime, because the TcpStream must be created on the runtime that will
+    /// drive it (a temporary runtime's I/O driver dies when the runtime drops).
+    pub ws_upgrade_request: Option<crate::ws_client::WsUpgradeRequest>,
 }
 
 #[allow(dead_code)] // Constructors used by different pool configurations
@@ -167,7 +169,7 @@ impl PluginState {
             last_broker_result: None,
             last_uuid_result: None,
             stream_sender: None,
-            ws_upstream: None,
+            ws_upgrade_request: None,
         }
     }
 
@@ -196,7 +198,7 @@ impl PluginState {
             last_broker_result: None,
             last_uuid_result: None,
             stream_sender: None,
-            ws_upstream: None,
+            ws_upgrade_request: None,
         }
     }
 
@@ -226,7 +228,7 @@ impl PluginState {
             last_broker_result: None,
             last_uuid_result: None,
             stream_sender: None,
-            ws_upstream: None,
+            ws_upgrade_request: None,
         }
     }
 
@@ -261,7 +263,7 @@ impl PluginState {
             last_broker_result: None,
             last_uuid_result: None,
             stream_sender: None,
-            ws_upstream: None,
+            ws_upgrade_request: None,
         }
     }
 
@@ -297,7 +299,7 @@ impl PluginState {
             last_broker_result: None,
             last_uuid_result: None,
             stream_sender: None,
-            ws_upstream: None,
+            ws_upgrade_request: None,
         }
     }
 
@@ -319,9 +321,9 @@ impl PluginState {
         self.stream_sender = Some(sender);
     }
 
-    /// Take the upstream WebSocket connection established by host_ws_upgrade (ADR-0026).
-    pub fn take_ws_upstream(&mut self) -> Option<crate::ws_client::UpstreamWsStream> {
-        self.ws_upstream.take()
+    /// Take the upstream WebSocket upgrade request from host_ws_upgrade (ADR-0026).
+    pub fn take_ws_upgrade_request(&mut self) -> Option<crate::ws_client::WsUpgradeRequest> {
+        self.ws_upgrade_request.take()
     }
 }
 
@@ -620,12 +622,12 @@ impl PluginInstance {
         self.store.data_mut().set_stream_sender(sender);
     }
 
-    /// Take the upstream WebSocket connection established by `host_ws_upgrade` (ADR-0026).
+    /// Take the upstream WebSocket upgrade request from `host_ws_upgrade` (ADR-0026).
     ///
-    /// Returns `None` if no WebSocket upgrade was performed or the connection
+    /// Returns `None` if no WebSocket upgrade was requested or the request
     /// was already taken.
-    pub fn take_ws_upstream(&mut self) -> Option<crate::ws_client::UpstreamWsStream> {
-        self.store.data_mut().take_ws_upstream()
+    pub fn take_ws_upgrade_request(&mut self) -> Option<crate::ws_client::WsUpgradeRequest> {
+        self.store.data_mut().take_ws_upgrade_request()
     }
 }
 
@@ -1159,13 +1161,14 @@ fn add_host_functions(linker: &mut Linker<PluginState>) -> Result<(), WasmError>
         )
         .map_err(|e| WasmError::Instantiation(format!("failed to add host_http_stream: {}", e)))?;
 
-    // host_ws_upgrade - connect to an upstream WebSocket endpoint (ADR-0026)
+    // host_ws_upgrade - request an upstream WebSocket connection (ADR-0026)
     //
     // The plugin sends a JSON payload: { url, connect_timeout_ms, headers }.
-    // On success: the upstream WebSocket connection is stored in PluginState
-    // for the data plane to pick up, and returns 0.
-    // On failure: the error string is stored in last_http_result (readable via
-    // host_http_read_result), and returns -1.
+    // The request is validated and stored in PluginState. The actual TCP
+    // connection is deferred to the async relay task on the main tokio runtime,
+    // because a TcpStream must be created on the runtime that will drive it
+    // (a temporary runtime's I/O driver dies when the runtime drops).
+    // Returns 0 on success (valid request), -1 on parse failure.
     linker
         .func_wrap(
             "barbacane",
@@ -1199,54 +1202,13 @@ fn add_host_functions(linker: &mut Linker<PluginState>) -> Result<(), WasmError>
                 tracing::debug!(
                     plugin = %plugin_name,
                     url = %ws_request.url,
-                    "host_ws_upgrade: connecting to upstream"
+                    "host_ws_upgrade: storing request for deferred connection"
                 );
 
-                // Connect to upstream WebSocket on a dedicated tokio runtime
-                let result = std::thread::scope(|s| {
-                    let handle = s.spawn(|| {
-                        let rt = match tokio::runtime::Builder::new_current_thread()
-                            .enable_all()
-                            .build()
-                        {
-                            Ok(rt) => rt,
-                            Err(e) => {
-                                tracing::error!("host_ws_upgrade: failed to create runtime: {}", e);
-                                return Err(format!("runtime error: {}", e));
-                            }
-                        };
-
-                        rt.block_on(crate::ws_client::connect_upstream(ws_request))
-                    });
-
-                    match handle.join() {
-                        Ok(result) => result,
-                        Err(e) => {
-                            tracing::error!("host_ws_upgrade: worker thread panicked: {:?}", e);
-                            Err("internal error: worker thread panicked".to_string())
-                        }
-                    }
-                });
-
-                match result {
-                    Ok(ws_stream) => {
-                        tracing::debug!(
-                            plugin = %plugin_name,
-                            "host_ws_upgrade: upstream connection established"
-                        );
-                        caller.data_mut().ws_upstream = Some(ws_stream);
-                        0
-                    }
-                    Err(err_msg) => {
-                        tracing::warn!(
-                            plugin = %plugin_name,
-                            error = %err_msg,
-                            "host_ws_upgrade: upstream connection failed"
-                        );
-                        caller.data_mut().last_http_result = Some(err_msg.into_bytes());
-                        -1
-                    }
-                }
+                // Store the request; the actual connection happens on the main
+                // runtime inside the relay task (see relay_websocket).
+                caller.data_mut().ws_upgrade_request = Some(ws_request);
+                0
             },
         )
         .map_err(|e| WasmError::Instantiation(format!("failed to add host_ws_upgrade: {}", e)))?;