diff --git a/client/transport/src/ws/mod.rs b/client/transport/src/ws/mod.rs index 9008d522d8..09ce80ea72 100644 --- a/client/transport/src/ws/mod.rs +++ b/client/transport/src/ws/mod.rs @@ -35,6 +35,7 @@ use futures_util::io::{BufReader, BufWriter}; use jsonrpsee_core::client::{MaybeSend, ReceivedMessage, TransportReceiverT, TransportSenderT}; use jsonrpsee_core::TEN_MB_SIZE_BYTES; use jsonrpsee_core::{async_trait, Cow}; +use soketto::connection::CloseReason; use soketto::connection::Error::Utf8; use soketto::data::ByteSlice125; use soketto::handshake::client::{Client as WsHandshakeClient, ServerResponse}; @@ -230,6 +231,9 @@ pub enum WsError { /// Message was too large. #[error("The message was too large")] MessageTooLarge, + /// Connection was closed. + #[error("Connection was closed: {0:?}")] + Closed(CloseReason), } #[async_trait] @@ -280,19 +284,16 @@ where /// Returns a `Future` resolving when the server sent us something back. async fn receive(&mut self) -> Result { - loop { - let mut message = Vec::new(); - let recv = self.inner.receive(&mut message).await?; - - match recv { - Incoming::Data(Data::Text(_)) => { - let s = String::from_utf8(message).map_err(|err| WsError::Connection(Utf8(err.utf8_error())))?; - break Ok(ReceivedMessage::Text(s)); - } - Incoming::Data(Data::Binary(_)) => break Ok(ReceivedMessage::Bytes(message)), - Incoming::Pong(_) => break Ok(ReceivedMessage::Pong), - _ => continue, + let mut message = Vec::new(); + + match self.inner.receive(&mut message).await? { + Incoming::Data(Data::Text(_)) => { + let s = String::from_utf8(message).map_err(|err| WsError::Connection(Utf8(err.utf8_error())))?; + Ok(ReceivedMessage::Text(s)) } + Incoming::Data(Data::Binary(_)) => Ok(ReceivedMessage::Bytes(message)), + Incoming::Pong(_) => Ok(ReceivedMessage::Pong), + Incoming::Closed(c) => Err(WsError::Closed(c)), } } }