diff --git a/rust/src/client_api/browser.rs b/rust/src/client_api/browser.rs index ae386d4..137fa9c 100644 --- a/rust/src/client_api/browser.rs +++ b/rust/src/client_api/browser.rs @@ -13,6 +13,13 @@ pub struct WebApi { error_handler: Box, } +impl Drop for WebApi { + fn drop(&mut self) { + // Close with normal closure code when dropped + let _ = self.conn.close_with_code(1000); + } +} + impl WebApi { pub fn start( conn: Connection, diff --git a/rust/src/client_api/client_events.rs b/rust/src/client_api/client_events.rs index 7c5d7f9..87bc70d 100644 --- a/rust/src/client_api/client_events.rs +++ b/rust/src/client_api/client_events.rs @@ -247,9 +247,15 @@ impl ContractError { pub enum ClientRequest<'a> { DelegateOp(#[serde(borrow)] DelegateRequest<'a>), ContractOp(#[serde(borrow)] ContractRequest<'a>), - Disconnect { cause: Option> }, - Authenticate { token: String }, + Disconnect { + cause: Option>, + }, + Authenticate { + token: String, + }, NodeQueries(ConnectedPeers), + /// Gracefully disconnect from the host. + Close, } #[derive(Serialize, Deserialize, Debug, Clone)] @@ -297,6 +303,7 @@ impl ClientRequest<'_> { ClientRequest::Disconnect { cause } => ClientRequest::Disconnect { cause }, ClientRequest::Authenticate { token } => ClientRequest::Authenticate { token }, ClientRequest::NodeQueries(query) => ClientRequest::NodeQueries(query), + ClientRequest::Close => ClientRequest::Close, } } @@ -600,6 +607,7 @@ impl Display for ClientRequest<'_> { ClientRequest::Disconnect { .. } => write!(f, "client disconnected"), ClientRequest::Authenticate { .. } => write!(f, "authenticate"), ClientRequest::NodeQueries(query) => write!(f, "node queries: {:?}", query), + ClientRequest::Close => write!(f, "close"), } } } diff --git a/rust/src/client_api/regular.rs b/rust/src/client_api/regular.rs index cfaa1c7..ae00723 100644 --- a/rust/src/client_api/regular.rs +++ b/rust/src/client_api/regular.rs @@ -9,7 +9,13 @@ use tokio::{ net::TcpStream, sync::mpsc::{self, Receiver, Sender}, }; -use tokio_tungstenite::{tungstenite::Message, MaybeTlsStream, WebSocketStream}; +use tokio_tungstenite::{ + tungstenite::{ + protocol::{frame::coding::CloseCode, CloseFrame}, + Message, + }, + MaybeTlsStream, WebSocketStream, +}; type Connection = WebSocketStream>; @@ -19,6 +25,15 @@ pub struct WebApi { queue: Vec>, } +impl Drop for WebApi { + fn drop(&mut self) { + let req = self.request_tx.clone(); + tokio::spawn(async move { + let _ = req.send(ClientRequest::Close).await; + }); + } +} + impl Stream for WebApi { type Item = HostResult; @@ -114,6 +129,7 @@ impl WebApi { res.ok_or_else(|| ClientError::from(ErrorKind::ChannelClosed))? } + #[doc(hidden)] pub async fn disconnect(self, cause: impl Into>) { let _ = self .request_tx @@ -148,6 +164,7 @@ async fn request_handler( tracing::debug!(?error, "request handler error"); let error = match error { Error::ChannelClosed => ErrorKind::ChannelClosed.into(), + Error::ConnectionClosed => ErrorKind::Disconnect.into(), other => ClientError::from(format!("{other}")), }; let _ = response_tx.send(Err(error)).await; @@ -163,6 +180,17 @@ async fn process_request( .map_err(Into::into) .map_err(Error::OtherError)?; conn.send(Message::Binary(msg.into())).await?; + if let ClientRequest::Disconnect { cause } = req { + conn.close(cause.map(|c| CloseFrame { + code: CloseCode::Normal, + reason: format!("{c}").into(), + })) + .await?; + return Err(Error::ConnectionClosed); + } else if let ClientRequest::Close = req { + conn.close(None).await?; + return Err(Error::ConnectionClosed); + } Ok(()) }