From d86a990cfb58bcf7c41a15df793018e19a7b422d Mon Sep 17 00:00:00 2001 From: npub1mprnacetjua2xx3p5eddmhxyk6wv929ymm5py8kd2xfxurxahspqqlgyta Date: Sat, 18 Jul 2026 13:39:53 -0400 Subject: [PATCH 1/4] fix(desktop): own native websocket lifecycle Co-authored-by: npub1mprnacetjua2xx3p5eddmhxyk6wv929ymm5py8kd2xfxurxahspqqlgyta Signed-off-by: npub1mprnacetjua2xx3p5eddmhxyk6wv929ymm5py8kd2xfxurxahspqqlgyta --- desktop/src-tauri/Cargo.lock | 56 --- desktop/src-tauri/Cargo.toml | 3 +- desktop/src-tauri/build.rs | 10 +- desktop/src-tauri/src/lib.rs | 4 +- desktop/src-tauri/src/native_websocket.rs | 442 ++++++++++++++++++ desktop/src/app/useReloadShortcut.ts | 30 +- .../shared/api/relayWebSocketClose.test.mjs | 62 +-- desktop/src/shared/api/relayWebSocketClose.ts | 28 +- desktop/src/testing/e2eBridge.ts | 19 + 9 files changed, 514 insertions(+), 140 deletions(-) create mode 100644 desktop/src-tauri/src/native_websocket.rs diff --git a/desktop/src-tauri/Cargo.lock b/desktop/src-tauri/Cargo.lock index ab536fbfe0..ea3348f93c 100644 --- a/desktop/src-tauri/Cargo.lock +++ b/desktop/src-tauri/Cargo.lock @@ -1035,7 +1035,6 @@ dependencies = [ "tauri-plugin-process", "tauri-plugin-single-instance", "tauri-plugin-updater", - "tauri-plugin-websocket", "tauri-plugin-window-state", "tempfile", "tokio", @@ -9770,26 +9769,6 @@ dependencies = [ "zip 4.6.1", ] -[[package]] -name = "tauri-plugin-websocket" -version = "2.4.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fe037be7e1c30be639fe12dc3077e8ec4709e6f30ca4a6016b7edc76232ae9ee" -dependencies = [ - "futures-util", - "http", - "log", - "rand 0.9.4", - "rustls", - "serde", - "serde_json", - "tauri", - "tauri-plugin", - "thiserror 2.0.18", - "tokio", - "tokio-tungstenite 0.28.0", -] - [[package]] name = "tauri-plugin-window-state" version = "2.4.1" @@ -10243,22 +10222,6 @@ dependencies = [ "webpki-roots 0.26.11", ] -[[package]] -name = "tokio-tungstenite" -version = "0.28.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d25a406cddcc431a75d3d9afc6a7c0f7428d4891dd973e4d54c56b46127bf857" -dependencies = [ - "futures-util", - "log", - "rustls", - "rustls-pki-types", - "tokio", - "tokio-rustls", - "tungstenite 0.28.0", - "webpki-roots 0.26.11", -] - [[package]] name = "tokio-tungstenite" version = "0.29.0" @@ -10691,25 +10654,6 @@ dependencies = [ "utf-8", ] -[[package]] -name = "tungstenite" -version = "0.28.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8628dcc84e5a09eb3d8423d6cb682965dea9133204e8fb3efee74c2a0c259442" -dependencies = [ - "bytes", - "data-encoding", - "http", - "httparse", - "log", - "rand 0.9.4", - "rustls", - "rustls-pki-types", - "sha1 0.10.6", - "thiserror 2.0.18", - "utf-8", -] - [[package]] name = "tungstenite" version = "0.29.0" diff --git a/desktop/src-tauri/Cargo.toml b/desktop/src-tauri/Cargo.toml index 0225e6930b..633d3fe9d0 100644 --- a/desktop/src-tauri/Cargo.toml +++ b/desktop/src-tauri/Cargo.toml @@ -63,14 +63,13 @@ tauri-plugin-deep-link = "2" tauri-plugin-opener = "2" tauri-plugin-single-instance = { version = "2", features = ["deep-link"] } tauri-plugin-window-state = "2" -tauri-plugin-websocket = "2" tauri-plugin-dialog = "2" tauri-plugin-updater = "2" tauri-plugin-process = "2" infer = "0.19" hex = "0.4" ed25519-dalek = "=3.0.0-rc.0" -tokio = { version = "1", features = ["fs", "sync", "rt", "macros", "time"] } +tokio = { version = "1", features = ["fs", "sync", "rt", "macros", "time", "net", "io-util"] } tokio-tungstenite = { version = "0.29", features = ["rustls-tls-webpki-roots"] } tokio-util = { version = "0.7", features = ["rt"] } bytes = "1" diff --git a/desktop/src-tauri/build.rs b/desktop/src-tauri/build.rs index c4d7b71f49..dfad9c2fec 100644 --- a/desktop/src-tauri/build.rs +++ b/desktop/src-tauri/build.rs @@ -117,5 +117,13 @@ fn main() { ); } - tauri_build::build() + tauri_build::try_build( + tauri_build::Attributes::new().plugin( + "websocket", + tauri_build::InlinedPlugin::new() + .commands(&["connect", "send", "disconnect", "disconnect_all"]) + .default_permission(tauri_build::DefaultPermissionRule::AllowAllCommands), + ), + ) + .expect("failed to build Tauri application"); } diff --git a/desktop/src-tauri/src/lib.rs b/desktop/src-tauri/src/lib.rs index da1ea50f7c..fe9b39899a 100644 --- a/desktop/src-tauri/src/lib.rs +++ b/desktop/src-tauri/src/lib.rs @@ -15,6 +15,7 @@ mod migration; #[cfg(test)] mod model_tests; mod models; +mod native_websocket; mod nostr_bind; pub mod nostr_convert; mod prevent_sleep; @@ -62,7 +63,6 @@ use tauri_plugin_window_state::StateFlags; #[cfg(target_os = "macos")] const INITIAL_RENDER_READY_EVENT: &str = "initial-render-ready"; - #[tauri::command] fn perform_sidebar_default_haptic() { #[cfg(target_os = "macos")] @@ -331,7 +331,7 @@ pub fn run() { }) .build(), ) - .plugin(tauri_plugin_websocket::init()) + .plugin(native_websocket::init()) .plugin(tauri_plugin_dialog::init()) .plugin(tauri_plugin_process::init()); diff --git a/desktop/src-tauri/src/native_websocket.rs b/desktop/src-tauri/src/native_websocket.rs new file mode 100644 index 0000000000..f1ac4d041c --- /dev/null +++ b/desktop/src-tauri/src/native_websocket.rs @@ -0,0 +1,442 @@ +use std::{collections::HashMap, sync::Arc, time::Duration}; + +use futures_util::{SinkExt, StreamExt}; +use serde::{Deserialize, Serialize}; +use tauri::{ipc::Channel, plugin::TauriPlugin, Manager, Runtime}; +use tokio::sync::{mpsc, oneshot, Mutex}; +use tokio_tungstenite::{ + connect_async, + tungstenite::protocol::{frame::coding::CloseCode, CloseFrame, Message}, +}; +use tokio_util::sync::CancellationToken; + +const CONNECT_TIMEOUT: Duration = Duration::from_secs(10); +const WRITE_TIMEOUT: Duration = Duration::from_secs(10); +const SHUTDOWN_TIMEOUT: Duration = Duration::from_millis(250); +const SEND_QUEUE_CAPACITY: usize = 64; + +type Id = u32; + +#[derive(Debug, Deserialize)] +#[serde(tag = "type", content = "data")] +enum WebSocketMessage { + Text(String), + Binary(Vec), + Ping(Vec), + Pong(Vec), + Close(Option), +} + +#[derive(Debug, Deserialize)] +struct CloseFramePayload { + code: u16, + reason: String, +} + +impl From for Message { + fn from(message: WebSocketMessage) -> Self { + match message { + WebSocketMessage::Text(value) => Message::Text(value.into()), + WebSocketMessage::Binary(value) => Message::Binary(value.into()), + WebSocketMessage::Ping(value) => Message::Ping(value.into()), + WebSocketMessage::Pong(value) => Message::Pong(value.into()), + WebSocketMessage::Close(frame) => Message::Close(frame.map(|frame| CloseFrame { + code: CloseCode::from(frame.code), + reason: frame.reason.into(), + })), + } + } +} + +#[derive(Serialize)] +#[serde(tag = "type", content = "data")] +enum OutboundMessage { + Text(String), + Binary(Vec), + Ping(Vec), + Pong(Vec), + Close(Option), + Error(String), +} + +#[derive(Serialize)] +struct CloseFramePayloadOut { + code: u16, + reason: String, +} + +struct SendRequest { + message: Message, + result: oneshot::Sender>, +} + +struct ConnectionHandle { + sender: mpsc::Sender, + cancel: CancellationToken, + task: Mutex>>, +} + +#[derive(Clone)] +struct WebSocketManager { + connections: Arc>>>, + connect_cancel: Arc>, +} + +impl Default for WebSocketManager { + fn default() -> Self { + Self { + connections: Arc::default(), + connect_cancel: Arc::new(Mutex::new(CancellationToken::new())), + } + } +} + +impl WebSocketManager { + async fn remove(&self, id: Id) -> Option> { + self.connections.lock().await.remove(&id) + } + + async fn disconnect_handle(handle: Arc) { + handle.cancel.cancel(); + if let Some(mut task) = handle.task.lock().await.take() { + if tokio::time::timeout(SHUTDOWN_TIMEOUT, &mut task) + .await + .is_err() + { + task.abort(); + } + } + } + + async fn disconnect(&self, id: Id) { + if let Some(handle) = self.remove(id).await { + Self::disconnect_handle(handle).await; + } + } +} + +#[tauri::command] +async fn connect( + manager: tauri::State<'_, WebSocketManager>, + url: String, + on_message: Channel, + _config: Option, +) -> Result { + let connect_cancel = manager.connect_cancel.lock().await.clone(); + let (socket, _) = tokio::select! { + _ = connect_cancel.cancelled() => return Err("WebSocket connection cancelled".to_string()), + result = tokio::time::timeout(CONNECT_TIMEOUT, connect_async(&url)) => result + .map_err(|_| "WebSocket connection timed out".to_string())? + .map_err(|error| error.to_string())?, + }; + + // Serialize registration with disconnect_all so a reload cannot miss a + // connection that finished its handshake concurrently with teardown. + let current_connect_cancel = manager.connect_cancel.lock().await; + if connect_cancel.is_cancelled() { + return Err("WebSocket connection cancelled".to_string()); + } + + let id = loop { + let candidate = uuid::Uuid::new_v4().as_u128() as u32; + if !manager.connections.lock().await.contains_key(&candidate) { + break candidate; + } + }; + let (sender, receiver) = mpsc::channel(SEND_QUEUE_CAPACITY); + let cancel = CancellationToken::new(); + let handle = Arc::new(ConnectionHandle { + sender, + cancel: cancel.clone(), + task: Mutex::new(None), + }); + let mut task_slot = handle.task.lock().await; + manager.connections.lock().await.insert(id, handle.clone()); + + let task_manager = manager.inner().clone(); + let task = tauri::async_runtime::spawn(run_connection( + id, + socket, + receiver, + cancel, + on_message, + task_manager, + )); + *task_slot = Some(task); + drop(task_slot); + drop(current_connect_cancel); + Ok(id) +} + +#[tauri::command] +async fn send( + manager: tauri::State<'_, WebSocketManager>, + id: Id, + message: WebSocketMessage, +) -> Result<(), String> { + let handle = manager + .connections + .lock() + .await + .get(&id) + .cloned() + .ok_or_else(|| format!("WebSocket connection {id} not found"))?; + let (result_tx, result_rx) = oneshot::channel(); + tokio::time::timeout( + WRITE_TIMEOUT, + handle.sender.send(SendRequest { + message: message.into(), + result: result_tx, + }), + ) + .await + .map_err(|_| "WebSocket send queue timed out".to_string())? + .map_err(|_| "WebSocket connection closed".to_string())?; + + tokio::time::timeout(WRITE_TIMEOUT, result_rx) + .await + .map_err(|_| "WebSocket send timed out".to_string())? + .map_err(|_| "WebSocket connection closed".to_string())? +} + +#[tauri::command] +async fn disconnect(manager: tauri::State<'_, WebSocketManager>, id: Id) -> Result<(), String> { + manager.disconnect(id).await; + Ok(()) +} + +#[tauri::command] +async fn disconnect_all(manager: tauri::State<'_, WebSocketManager>) -> Result<(), String> { + let mut connect_cancel = manager.connect_cancel.lock().await; + connect_cancel.cancel(); + *connect_cancel = CancellationToken::new(); + let handles = { + let mut connections = manager.connections.lock().await; + connections + .drain() + .map(|(_, handle)| handle) + .collect::>() + }; + futures_util::future::join_all(handles.into_iter().map(WebSocketManager::disconnect_handle)) + .await; + Ok(()) +} + +async fn run_connection( + id: Id, + mut socket: tokio_tungstenite::WebSocketStream, + mut receiver: mpsc::Receiver, + cancel: CancellationToken, + on_message: Channel, + manager: WebSocketManager, +) where + S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin, +{ + loop { + tokio::select! { + _ = cancel.cancelled() => { + let _ = tokio::time::timeout( + SHUTDOWN_TIMEOUT, + socket.send(Message::Close(Some(CloseFrame { + code: CloseCode::Normal, + reason: "disconnect".into(), + }))), + ).await; + break; + } + request = receiver.recv() => { + let Some(request) = request else { break }; + let result = tokio::time::timeout(WRITE_TIMEOUT, socket.send(request.message)) + .await + .map_err(|_| "WebSocket send timed out".to_string()) + .and_then(|result| result.map_err(|error| error.to_string())); + let failed = result.is_err(); + let _ = request.result.send(result); + if failed { break; } + } + incoming = socket.next() => { + let message = match incoming { + Some(Ok(message)) => outbound_message(message), + Some(Err(error)) => OutboundMessage::Error(error.to_string()), + None => OutboundMessage::Close(None), + }; + let terminal = matches!(message, OutboundMessage::Close(_) | OutboundMessage::Error(_)); + if let Ok(value) = serde_json::to_value(message) { + let _ = on_message.send(value); + } + if terminal { break; } + } + } + } + manager.remove(id).await; +} + +fn outbound_message(message: Message) -> OutboundMessage { + match message { + Message::Text(value) => OutboundMessage::Text(value.to_string()), + Message::Binary(value) => OutboundMessage::Binary(value.to_vec()), + Message::Ping(value) => OutboundMessage::Ping(value.to_vec()), + Message::Pong(value) => OutboundMessage::Pong(value.to_vec()), + Message::Close(frame) => OutboundMessage::Close(frame.map(|frame| CloseFramePayloadOut { + code: frame.code.into(), + reason: frame.reason.to_string(), + })), + Message::Frame(_) => OutboundMessage::Error("unexpected raw WebSocket frame".to_string()), + } +} + +pub fn init() -> TauriPlugin { + tauri::plugin::Builder::new("websocket") + .invoke_handler(tauri::generate_handler![ + connect, + send, + disconnect, + disconnect_all + ]) + .setup(|app, _api| { + app.manage(WebSocketManager::default()); + Ok(()) + }) + .build() +} + +#[cfg(test)] +mod tests { + use super::*; + use tauri::ipc::InvokeResponseBody; + use tokio::io::duplex; + use tokio_tungstenite::{tungstenite::protocol::Role, WebSocketStream}; + + fn silent_channel() -> Channel { + Channel::new(|_: InvokeResponseBody| Ok(())) + } + + #[tokio::test] + async fn eof_removes_connection() { + let manager = WebSocketManager::default(); + let (client_io, server_io) = duplex(1024); + let (client, server) = tokio::join!( + WebSocketStream::from_raw_socket(client_io, Role::Client, None), + WebSocketStream::from_raw_socket(server_io, Role::Server, None), + ); + let (sender, receiver) = mpsc::channel(SEND_QUEUE_CAPACITY); + let handle = Arc::new(ConnectionHandle { + sender, + cancel: CancellationToken::new(), + task: Mutex::new(None), + }); + manager.connections.lock().await.insert(1, handle.clone()); + let task = tauri::async_runtime::spawn(run_connection( + 1, + client, + receiver, + handle.cancel.clone(), + silent_channel(), + manager.clone(), + )); + *handle.task.lock().await = Some(task); + + drop(server); + tokio::time::timeout(Duration::from_secs(1), async { + while manager.connections.lock().await.contains_key(&1) { + tokio::task::yield_now().await; + } + }) + .await + .expect("EOF should clean up its native connection ID"); + } + + #[tokio::test] + async fn disconnect_removes_before_bounded_task_shutdown() { + let manager = WebSocketManager::default(); + let (sender, _receiver) = mpsc::channel(SEND_QUEUE_CAPACITY); + let handle = Arc::new(ConnectionHandle { + sender, + cancel: CancellationToken::new(), + task: Mutex::new(Some(tauri::async_runtime::spawn(async { + std::future::pending::<()>().await; + }))), + }); + manager.connections.lock().await.insert(7, handle); + + tokio::time::timeout(Duration::from_secs(1), manager.disconnect(7)) + .await + .expect("disconnect should abort an unresponsive task"); + assert!(!manager.connections.lock().await.contains_key(&7)); + + // Repeated teardown is intentionally a no-op. + manager.disconnect(7).await; + } + + #[tokio::test] + async fn teardown_gate_stays_closed_until_tasks_stop() { + let manager = WebSocketManager::default(); + let gate = manager.connect_cancel.lock().await; + let (sender, _receiver) = mpsc::channel(SEND_QUEUE_CAPACITY); + let handle = Arc::new(ConnectionHandle { + sender, + cancel: CancellationToken::new(), + task: Mutex::new(Some(tauri::async_runtime::spawn(async { + std::future::pending::<()>().await; + }))), + }); + manager.connections.lock().await.insert(1, handle); + gate.cancel(); + let handles = { + let mut connections = manager.connections.lock().await; + connections + .drain() + .map(|(_, handle)| handle) + .collect::>() + }; + + let shutdown = futures_util::future::join_all( + handles.into_iter().map(WebSocketManager::disconnect_handle), + ); + assert!(manager.connect_cancel.try_lock().is_err()); + shutdown.await; + drop(gate); + assert!(manager.connect_cancel.try_lock().is_ok()); + } + + #[tokio::test] + async fn one_connection_does_not_block_another_send_queue() { + let manager = WebSocketManager::default(); + let (blocked_sender, blocked_receiver) = mpsc::channel(1); + blocked_sender + .send(SendRequest { + message: Message::Text("blocked".into()), + result: oneshot::channel().0, + }) + .await + .unwrap(); + let blocked = Arc::new(ConnectionHandle { + sender: blocked_sender, + cancel: CancellationToken::new(), + task: Mutex::new(None), + }); + manager.connections.lock().await.insert(1, blocked); + + let (healthy_sender, mut healthy_receiver) = mpsc::channel(1); + let healthy = Arc::new(ConnectionHandle { + sender: healthy_sender.clone(), + cancel: CancellationToken::new(), + task: Mutex::new(None), + }); + manager.connections.lock().await.insert(2, healthy); + + let (result, _) = oneshot::channel(); + tokio::time::timeout( + Duration::from_millis(50), + healthy_sender.send(SendRequest { + message: Message::Text("healthy".into()), + result, + }), + ) + .await + .expect("a full queue on one connection must not block another") + .unwrap(); + assert!(healthy_receiver.recv().await.is_some()); + drop(blocked_receiver); + } +} diff --git a/desktop/src/app/useReloadShortcut.ts b/desktop/src/app/useReloadShortcut.ts index 74bfd4f2fe..fb93f77f6f 100644 --- a/desktop/src/app/useReloadShortcut.ts +++ b/desktop/src/app/useReloadShortcut.ts @@ -1,36 +1,34 @@ import * as React from "react"; +import { closeAllWebSockets } from "@/shared/api/relayWebSocketClose"; import { hasPrimaryShortcutModifier } from "@/shared/lib/platform"; -/** - * Reloads the webview on the platform's reload shortcut (Cmd+R on macOS, - * Ctrl+R elsewhere), matching browser behavior. - * - * `window.location.reload()` is the app's existing reload primitive (see - * App.tsx, useCommunityInit.ts): it triggers a full reinit that re-reads - * localStorage and reconnects relays. - */ +const RELOAD_TEARDOWN_TIMEOUT_MS = 500; + +/** Reloads the webview after bounded native WebSocket teardown. */ export function useReloadShortcut() { React.useEffect(() => { - function handleKeyDown(event: KeyboardEvent) { + async function handleKeyDown(event: KeyboardEvent) { if ( !hasPrimaryShortcutModifier(event) || event.altKey || - event.shiftKey + event.shiftKey || + event.key.toLowerCase() !== "r" ) { return; } - if (event.key.toLowerCase() !== "r") { - return; - } event.preventDefault(); + await Promise.race([ + closeAllWebSockets(), + new Promise((resolve) => + window.setTimeout(resolve, RELOAD_TEARDOWN_TIMEOUT_MS), + ), + ]); window.location.reload(); } window.addEventListener("keydown", handleKeyDown); - return () => { - window.removeEventListener("keydown", handleKeyDown); - }; + return () => window.removeEventListener("keydown", handleKeyDown); }, []); } diff --git a/desktop/src/shared/api/relayWebSocketClose.test.mjs b/desktop/src/shared/api/relayWebSocketClose.test.mjs index 59b646385f..2cb34f604c 100644 --- a/desktop/src/shared/api/relayWebSocketClose.test.mjs +++ b/desktop/src/shared/api/relayWebSocketClose.test.mjs @@ -1,67 +1,29 @@ import assert from "node:assert/strict"; -import fs from "node:fs"; -import path from "node:path"; import test from "node:test"; -import { fileURLToPath } from "node:url"; -import { closeWebSocket } from "./relayWebSocketClose.ts"; +import { closeAllWebSockets, closeWebSocket } from "./relayWebSocketClose.ts"; -test("closeWebSocket sends a Close frame through plugin:websocket|send", async () => { +test("closeWebSocket invokes authoritative native disconnect", async () => { const calls = []; await closeWebSocket(42, "community switch", async (cmd, args) => { calls.push({ cmd, args }); }); - assert.equal(calls.length, 1); - assert.equal(calls[0].cmd, "plugin:websocket|send"); - assert.deepEqual(calls[0].args, { - id: 42, - message: { - type: "Close", - data: { code: 1000, reason: "community switch" }, - }, - }); + assert.deepEqual(calls, [ + { cmd: "plugin:websocket|disconnect", args: { id: 42 } }, + ]); }); -test("closeWebSocket swallows send failures (socket already gone)", async () => { +test("closeWebSocket is idempotent when the native socket is gone", async () => { await closeWebSocket(7, "connection reset", async () => { throw new Error("WebSocket connection not found"); }); }); -// Regression guard: tauri-plugin-websocket registers only `connect` and -// `send` — there is no `disconnect` command. Invoking one rejects silently -// and leaks the socket (relay zombie pile, community-switch disconnects). -// Any socket teardown must go through closeWebSocket. -test("no source file invokes the nonexistent plugin:websocket|disconnect command", () => { - const srcRoot = path.resolve( - path.dirname(fileURLToPath(import.meta.url)), - "..", - "..", - ); - const offenders = []; - - const walk = (dir) => { - for (const entry of fs.readdirSync(dir, { withFileTypes: true })) { - const full = path.join(dir, entry.name); - if (entry.isDirectory()) { - walk(full); - continue; - } - if (!/\.(ts|tsx|js|jsx|mjs)$/.test(entry.name)) continue; - if (full === fileURLToPath(import.meta.url)) continue; - if ( - fs.readFileSync(full, "utf8").includes("plugin:websocket|disconnect") - ) { - offenders.push(path.relative(srcRoot, full)); - } - } - }; - walk(srcRoot); - - assert.deepEqual( - offenders, - [], - "plugin:websocket|disconnect does not exist in tauri-plugin-websocket — use closeWebSocket (Close frame via plugin:websocket|send) instead", - ); +test("closeAllWebSockets invokes native process-wide teardown", async () => { + const calls = []; + await closeAllWebSockets(async (cmd, args) => calls.push({ cmd, args })); + assert.deepEqual(calls, [ + { cmd: "plugin:websocket|disconnect_all", args: undefined }, + ]); }); diff --git a/desktop/src/shared/api/relayWebSocketClose.ts b/desktop/src/shared/api/relayWebSocketClose.ts index b4346f0656..9a31f8c8ee 100644 --- a/desktop/src/shared/api/relayWebSocketClose.ts +++ b/desktop/src/shared/api/relayWebSocketClose.ts @@ -1,28 +1,30 @@ import { invoke } from "@tauri-apps/api/core"; /** - * tauri-plugin-websocket 2.4.2 registers only `connect` and `send` — there is - * no `disconnect` command, so invoking one rejects and the socket leaks. Close - * the way the plugin's own JS API does: send a Close frame; the plugin's read - * loop drops the connection when the peer echoes the Close (or the TCP read - * stream terminates). + * Remove the connection from the native manager before waiting for its socket + * task to stop. The command is bounded and idempotent; failures mean the + * process is already tearing down or the socket is already gone. */ export function closeWebSocket( id: number, reason: string, invokeFn: typeof invoke = invoke, ): Promise { - return invokeFn("plugin:websocket|send", { - id, - message: { - type: "Close", - data: { code: 1000, reason }, - }, - }).then( + return invokeFn("plugin:websocket|disconnect", { id }).then( () => undefined, (err) => { - // Expected when the socket is already gone; greppable for anything else. console.debug(`closeWebSocket(${id}, ${reason}) rejected:`, err); }, ); } + +export function closeAllWebSockets( + invokeFn: typeof invoke = invoke, +): Promise { + return invokeFn("plugin:websocket|disconnect_all").then( + () => undefined, + (err) => { + console.debug("closeAllWebSockets() rejected:", err); + }, + ); +} diff --git a/desktop/src/testing/e2eBridge.ts b/desktop/src/testing/e2eBridge.ts index 7d491ba81b..03513e67e7 100644 --- a/desktop/src/testing/e2eBridge.ts +++ b/desktop/src/testing/e2eBridge.ts @@ -9850,6 +9850,25 @@ export function maybeInstallE2eTauriMocks() { return connectMockSocket( payload as Parameters[0], ); + case "plugin:websocket|disconnect": { + const { id } = payload as { id: number }; + if (isRelayMode(activeConfig)) { + realSockets.get(id)?.close(); + realSockets.delete(id); + } else { + const socket = mockSockets.get(id); + mockSockets.delete(id); + if (socket) sendWsClose(socket.handler); + } + return null; + } + case "plugin:websocket|disconnect_all": + for (const socket of realSockets.values()) socket.close(); + realSockets.clear(); + for (const socket of mockSockets.values()) sendWsClose(socket.handler); + mockSockets.clear(); + mockWebsocketSendMutexWedged = false; + return null; case "plugin:websocket|send": if (isRelayMode(activeConfig)) { return sendToRealSocket( From 3633663fe9a304956d25aacfd104a9b3509b1abb Mon Sep 17 00:00:00 2001 From: npub1mprnacetjua2xx3p5eddmhxyk6wv929ymm5py8kd2xfxurxahspqqlgyta Date: Sat, 18 Jul 2026 13:51:16 -0400 Subject: [PATCH 2/4] test(desktop): prove websocket teardown completes Co-authored-by: npub1mprnacetjua2xx3p5eddmhxyk6wv929ymm5py8kd2xfxurxahspqqlgyta Signed-off-by: npub1mprnacetjua2xx3p5eddmhxyk6wv929ymm5py8kd2xfxurxahspqqlgyta --- desktop/src-tauri/src/native_websocket.rs | 93 ++++++++++++++++++++--- 1 file changed, 83 insertions(+), 10 deletions(-) diff --git a/desktop/src-tauri/src/native_websocket.rs b/desktop/src-tauri/src/native_websocket.rs index f1ac4d041c..1162820200 100644 --- a/desktop/src-tauri/src/native_websocket.rs +++ b/desktop/src-tauri/src/native_websocket.rs @@ -104,6 +104,7 @@ impl WebSocketManager { .is_err() { task.abort(); + let _ = task.await; } } } @@ -115,17 +116,15 @@ impl WebSocketManager { } } -#[tauri::command] -async fn connect( - manager: tauri::State<'_, WebSocketManager>, - url: String, +async fn open_connection( + manager: &WebSocketManager, + url: &str, on_message: Channel, - _config: Option, ) -> Result { let connect_cancel = manager.connect_cancel.lock().await.clone(); let (socket, _) = tokio::select! { _ = connect_cancel.cancelled() => return Err("WebSocket connection cancelled".to_string()), - result = tokio::time::timeout(CONNECT_TIMEOUT, connect_async(&url)) => result + result = tokio::time::timeout(CONNECT_TIMEOUT, connect_async(url)) => result .map_err(|_| "WebSocket connection timed out".to_string())? .map_err(|error| error.to_string())?, }; @@ -153,7 +152,7 @@ async fn connect( let mut task_slot = handle.task.lock().await; manager.connections.lock().await.insert(id, handle.clone()); - let task_manager = manager.inner().clone(); + let task_manager = manager.clone(); let task = tauri::async_runtime::spawn(run_connection( id, socket, @@ -169,8 +168,17 @@ async fn connect( } #[tauri::command] -async fn send( +async fn connect( manager: tauri::State<'_, WebSocketManager>, + url: String, + on_message: Channel, + _config: Option, +) -> Result { + open_connection(manager.inner(), &url, on_message).await +} + +async fn send_message( + manager: &WebSocketManager, id: Id, message: WebSocketMessage, ) -> Result<(), String> { @@ -199,6 +207,15 @@ async fn send( .map_err(|_| "WebSocket connection closed".to_string())? } +#[tauri::command] +async fn send( + manager: tauri::State<'_, WebSocketManager>, + id: Id, + message: WebSocketMessage, +) -> Result<(), String> { + send_message(manager.inner(), id, message).await +} + #[tauri::command] async fn disconnect(manager: tauri::State<'_, WebSocketManager>, id: Id) -> Result<(), String> { manager.disconnect(id).await; @@ -303,6 +320,8 @@ pub fn init() -> TauriPlugin { #[cfg(test)] mod tests { use super::*; + use std::sync::atomic::{AtomicBool, Ordering}; + use tauri::ipc::InvokeResponseBody; use tokio::io::duplex; use tokio_tungstenite::{tungstenite::protocol::Role, WebSocketStream}; @@ -311,6 +330,46 @@ mod tests { Channel::new(|_: InvokeResponseBody| Ok(())) } + #[tokio::test] + async fn live_tcp_server_connect_send_and_disconnect() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let (received_tx, received_rx) = oneshot::channel(); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut socket = tokio_tungstenite::accept_async(stream).await.unwrap(); + let message = socket.next().await.unwrap().unwrap(); + received_tx.send(message).unwrap(); + while let Some(message) = socket.next().await { + if matches!(message, Ok(Message::Close(_))) { + break; + } + } + }); + + let manager = WebSocketManager::default(); + let id = open_connection(&manager, &format!("ws://{address}"), silent_channel()) + .await + .unwrap(); + send_message(&manager, id, WebSocketMessage::Text("live-probe".into())) + .await + .unwrap(); + assert_eq!( + tokio::time::timeout(Duration::from_secs(1), received_rx) + .await + .unwrap() + .unwrap(), + Message::Text("live-probe".into()) + ); + + manager.disconnect(id).await; + assert!(!manager.connections.lock().await.contains_key(&id)); + tokio::time::timeout(Duration::from_secs(1), server) + .await + .expect("live server should observe native socket shutdown") + .unwrap(); + } + #[tokio::test] async fn eof_removes_connection() { let manager = WebSocketManager::default(); @@ -347,22 +406,36 @@ mod tests { } #[tokio::test] - async fn disconnect_removes_before_bounded_task_shutdown() { + async fn disconnect_removes_and_drops_task_before_returning() { + struct DropGuard(Arc); + impl Drop for DropGuard { + fn drop(&mut self) { + self.0.store(true, Ordering::SeqCst); + } + } + let manager = WebSocketManager::default(); + let dropped = Arc::new(AtomicBool::new(false)); + let task_dropped = dropped.clone(); + let (ready_tx, ready_rx) = oneshot::channel(); let (sender, _receiver) = mpsc::channel(SEND_QUEUE_CAPACITY); let handle = Arc::new(ConnectionHandle { sender, cancel: CancellationToken::new(), - task: Mutex::new(Some(tauri::async_runtime::spawn(async { + task: Mutex::new(Some(tauri::async_runtime::spawn(async move { + let _guard = DropGuard(task_dropped); + ready_tx.send(()).unwrap(); std::future::pending::<()>().await; }))), }); manager.connections.lock().await.insert(7, handle); + ready_rx.await.unwrap(); tokio::time::timeout(Duration::from_secs(1), manager.disconnect(7)) .await .expect("disconnect should abort an unresponsive task"); assert!(!manager.connections.lock().await.contains_key(&7)); + assert!(dropped.load(Ordering::SeqCst)); // Repeated teardown is intentionally a no-op. manager.disconnect(7).await; From fd2558871d4a98c34cdfce65697317b8e8755eec Mon Sep 17 00:00:00 2001 From: npub1mprnacetjua2xx3p5eddmhxyk6wv929ymm5py8kd2xfxurxahspqqlgyta Date: Sat, 18 Jul 2026 14:16:21 -0400 Subject: [PATCH 3/4] fix(desktop): select rustls crypto provider Co-authored-by: npub1mprnacetjua2xx3p5eddmhxyk6wv929ymm5py8kd2xfxurxahspqqlgyta Signed-off-by: npub1mprnacetjua2xx3p5eddmhxyk6wv929ymm5py8kd2xfxurxahspqqlgyta --- desktop/src-tauri/Cargo.lock | 1 + desktop/src-tauri/Cargo.toml | 1 + desktop/src-tauri/src/native_websocket.rs | 26 +++++++++++++++++++++++ 3 files changed, 28 insertions(+) diff --git a/desktop/src-tauri/Cargo.lock b/desktop/src-tauri/Cargo.lock index ea3348f93c..9c56b6c6d5 100644 --- a/desktop/src-tauri/Cargo.lock +++ b/desktop/src-tauri/Cargo.lock @@ -1017,6 +1017,7 @@ dependencies = [ "rodio", "rubato", "rusqlite", + "rustls", "security-framework 3.7.0", "serde", "serde_json", diff --git a/desktop/src-tauri/Cargo.toml b/desktop/src-tauri/Cargo.toml index 633d3fe9d0..87b2ded1c2 100644 --- a/desktop/src-tauri/Cargo.toml +++ b/desktop/src-tauri/Cargo.toml @@ -83,6 +83,7 @@ toml = "0.8" nostr = { version = "0.44", features = ["nip44"] } zeroize = "1" reqwest = { version = "0.13", features = ["json", "query", "stream", "blocking"] } +rustls = { version = "0.23", default-features = false, features = ["aws_lc_rs", "std"] } url = "2" buzz_core_pkg = { package = "buzz-core", path = "../../crates/buzz-core" } buzz_persona_pkg = { package = "buzz-persona", path = "../../crates/buzz-persona" } diff --git a/desktop/src-tauri/src/native_websocket.rs b/desktop/src-tauri/src/native_websocket.rs index 1162820200..c0cf2e76f1 100644 --- a/desktop/src-tauri/src/native_websocket.rs +++ b/desktop/src-tauri/src/native_websocket.rs @@ -15,6 +15,11 @@ const WRITE_TIMEOUT: Duration = Duration::from_secs(10); const SHUTDOWN_TIMEOUT: Duration = Duration::from_millis(250); const SEND_QUEUE_CAPACITY: usize = 64; +pub(crate) fn install_crypto_provider() { + // Dependencies enable both rustls providers; choose one before TLS setup. + let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); +} + type Id = u32; #[derive(Debug, Deserialize)] @@ -303,6 +308,7 @@ fn outbound_message(message: Message) -> OutboundMessage { } pub fn init() -> TauriPlugin { + install_crypto_provider(); tauri::plugin::Builder::new("websocket") .invoke_handler(tauri::generate_handler![ connect, @@ -320,6 +326,7 @@ pub fn init() -> TauriPlugin { #[cfg(test)] mod tests { use super::*; + use futures_util::FutureExt; use std::sync::atomic::{AtomicBool, Ordering}; use tauri::ipc::InvokeResponseBody; @@ -330,6 +337,25 @@ mod tests { Channel::new(|_: InvokeResponseBody| Ok(())) } + #[tokio::test] + async fn secure_websocket_reaches_tls_without_panicking() { + install_crypto_provider(); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (_stream, _) = listener.accept().await.unwrap(); + tokio::time::sleep(Duration::from_millis(100)).await; + }); + let result = std::panic::AssertUnwindSafe(tokio_tungstenite::connect_async(format!( + "wss://{address}" + ))) + .catch_unwind() + .await; + + assert!(result.is_ok(), "TLS setup must not panic"); + server.await.unwrap(); + } + #[tokio::test] async fn live_tcp_server_connect_send_and_disconnect() { let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); From 3cab1abefb3d7d66c5383ef68363121f67c3a21b Mon Sep 17 00:00:00 2001 From: npub1mprnacetjua2xx3p5eddmhxyk6wv929ymm5py8kd2xfxurxahspqqlgyta Date: Sat, 18 Jul 2026 14:40:45 -0400 Subject: [PATCH 4/4] test(desktop): expect native websocket disconnect Co-authored-by: npub1mprnacetjua2xx3p5eddmhxyk6wv929ymm5py8kd2xfxurxahspqqlgyta Signed-off-by: npub1mprnacetjua2xx3p5eddmhxyk6wv929ymm5py8kd2xfxurxahspqqlgyta --- desktop/tests/e2e/onboarding.spec.ts | 18 +++++------------- 1 file changed, 5 insertions(+), 13 deletions(-) diff --git a/desktop/tests/e2e/onboarding.spec.ts b/desktop/tests/e2e/onboarding.spec.ts index 5602b51b8f..533a0dae8b 100644 --- a/desktop/tests/e2e/onboarding.spec.ts +++ b/desktop/tests/e2e/onboarding.spec.ts @@ -1663,26 +1663,18 @@ test("membership denial can import a different invited key", async ({ // Alice already has a relay profile with a display name, so after the // identity swap the onboarding gate auto-completes. - // The identity swap must tear down the old relay socket. There is no - // `plugin:websocket|disconnect` command in tauri-plugin-websocket — closing - // is a Close frame sent through `plugin:websocket|send`. + // The identity swap must tear down the old relay socket through the native + // disconnect command before the replacement identity connects. await expect .poll(() => page.evaluate( () => ( window as Window & { - __BUZZ_E2E_COMMAND_PAYLOADS__?: Array<{ - command: string; - payload: unknown; - }>; + __BUZZ_E2E_COMMANDS__?: string[]; } - ).__BUZZ_E2E_COMMAND_PAYLOADS__?.some( - (entry) => - entry.command === "plugin:websocket|send" && - (entry.payload as { message?: { type?: string } })?.message - ?.type === "Close", - ) ?? false, + ).__BUZZ_E2E_COMMANDS__?.includes("plugin:websocket|disconnect") ?? + false, ), ) .toBe(true);