fix(webocket): Avoid panic when polling quicksink after errors

Pull-Request: libp2p#5482.

Cargo.toml

@@ -112,7 +112,7 @@ libp2p-upnp = { version = "0.2.2", path = "protocols/upnp" }
 libp2p-webrtc = { version = "0.7.1-alpha", path = "transports/webrtc" }
 libp2p-webrtc-utils = { version = "0.2.1", path = "misc/webrtc-utils" }
 libp2p-webrtc-websys = { version = "0.3.0-alpha", path = "transports/webrtc-websys" }
-libp2p-websocket = { version = "0.43.1", path = "transports/websocket" }
+libp2p-websocket = { version = "0.43.2", path = "transports/websocket" }
 libp2p-websocket-websys = { version = "0.3.2", path = "transports/websocket-websys" }
 libp2p-webtransport-websys = { version = "0.3.0", path = "transports/webtransport-websys" }
 libp2p-yamux = { version = "0.45.2", path = "muxers/yamux" }

transports/websocket/CHANGELOG.md

@@ -1,3 +1,9 @@
+## 0.43.2
+
+- fix: Avoid websocket panic on polling after errors. See [PR 5482].
+
+[PR 5482]: https://github.com/libp2p/rust-libp2p/pull/5482
+
 ## 0.43.1
 
 ## 0.43.0

transports/websocket/Cargo.toml

@@ -3,7 +3,7 @@ name = "libp2p-websocket"
 edition = "2021"
 rust-version = { workspace = true }
 description = "WebSocket transport for libp2p"
-version = "0.43.1"
+version = "0.43.2"
 authors = ["Parity Technologies <admin@parity.io>"]
 license = "MIT"
 repository = "https://github.com/libp2p/rust-libp2p"
@@ -21,6 +21,7 @@ pin-project-lite = "0.2.14"
 rw-stream-sink = { workspace = true }
 soketto = "0.8.0"
 tracing = { workspace = true }
+thiserror = "1.0.61"
 url = "2.5"
 webpki-roots = "0.25"
 

transports/websocket/src/framed.rs

@@ -571,7 +571,7 @@ fn location_to_multiaddr<T>(location: &str) -> Result<Multiaddr, Error<T>> {
 /// The websocket connection.
 pub struct Connection<T> {
  receiver: BoxStream<'static, Result<Incoming, connection::Error>>,
- sender: Pin<Box<dyn Sink<OutgoingData, Error = connection::Error> + Send>>,
+ sender: Pin<Box<dyn Sink<OutgoingData, Error = quicksink::Error<connection::Error>> + Send>>,
  _marker: std::marker::PhantomData<T>,
 }
 

transports/websocket/src/quicksink.rs

@@ -30,14 +30,6 @@
 // Ok::<_, io::Error>(stdout)
 // });
 // ```
-//
-// # Panics
-//
-// - If any of the [`Sink`] methods produce an error, the sink transitions
-// to a failure state and none of its methods must be called afterwards or
-// else a panic will occur.
-// - If [`Sink::poll_close`] has been called, no other sink method must be
-// called afterwards or else a panic will be caused.
 
 use futures::{ready, sink::Sink};
 use pin_project_lite::pin_project;
@@ -102,6 +94,15 @@ enum State {
  Failed,
 }
 
+/// Errors the `Sink` may return.
+#[derive(Debug, thiserror::Error)]
+pub(crate) enum Error<E> {
+ #[error("Error while sending over the sink, {0}")]
+ Send(E),
+ #[error("The Sink has closed")]
+ Closed,
+}
+
 pin_project! {
  /// `SinkImpl` implements the `Sink` trait.
  #[derive(Debug)]
@@ -119,7 +120,7 @@ where
  F: FnMut(S, Action<A>) -> T,
  T: Future<Output = Result<S, E>>,
 {
- type Error = E;
+ type Error = Error<E>;
 
  fn poll_ready(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
  let mut this = self.project();
@@ -135,28 +136,27 @@ where
  Err(e) => {
  this.future.set(None);
  *this.state = State::Failed;
- Poll::Ready(Err(e))
+ Poll::Ready(Err(Error::Send(e)))
  }
  }
  }
  State::Closing => match ready!(this.future.as_mut().as_pin_mut().unwrap().poll(cx)) {
  Ok(_) => {
  this.future.set(None);
  *this.state = State::Closed;
- panic!("SinkImpl::poll_ready called on a closing sink.")
+ Poll::Ready(Err(Error::Closed))
  }
  Err(e) => {
  this.future.set(None);
  *this.state = State::Failed;
- Poll::Ready(Err(e))
+ Poll::Ready(Err(Error::Send(e)))
  }
  },
  State::Empty => {
  assert!(this.param.is_some());
  Poll::Ready(Ok(()))
  }
- State::Closed => panic!("SinkImpl::poll_ready called on a closed sink."),
- State::Failed => panic!("SinkImpl::poll_ready called after error."),
+ State::Closed | State::Failed => Poll::Ready(Err(Error::Closed)),
  }
  }
 
@@ -193,7 +193,7 @@ where
  Err(e) => {
  this.future.set(None);
  *this.state = State::Failed;
- return Poll::Ready(Err(e));
+ return Poll::Ready(Err(Error::Send(e)));
  }
  },
  State::Flushing => {
@@ -207,7 +207,7 @@ where
  Err(e) => {
  this.future.set(None);
  *this.state = State::Failed;
- return Poll::Ready(Err(e));
+ return Poll::Ready(Err(Error::Send(e)));
  }
  }
  }
@@ -221,11 +221,10 @@ where
  Err(e) => {
  this.future.set(None);
  *this.state = State::Failed;
- return Poll::Ready(Err(e));
+ return Poll::Ready(Err(Error::Send(e)));
  }
  },
- State::Closed => return Poll::Ready(Ok(())),
- State::Failed => panic!("SinkImpl::poll_flush called after error."),
+ State::Closed | State::Failed => return Poll::Ready(Err(Error::Closed)),
  }
  }
  }
@@ -253,7 +252,7 @@ where
  Err(e) => {
  this.future.set(None);
  *this.state = State::Failed;
- return Poll::Ready(Err(e));
+ return Poll::Ready(Err(Error::Send(e)));
  }
  },
  State::Flushing => {
@@ -266,7 +265,7 @@ where
  Err(e) => {
  this.future.set(None);
  *this.state = State::Failed;
- return Poll::Ready(Err(e));
+ return Poll::Ready(Err(Error::Send(e)));
  }
  }
  }
@@ -280,11 +279,11 @@ where
  Err(e) => {
  this.future.set(None);
  *this.state = State::Failed;
- return Poll::Ready(Err(e));
+ return Poll::Ready(Err(Error::Send(e)));
  }
  },
  State::Closed => return Poll::Ready(Ok(())),
- State::Failed => panic!("SinkImpl::poll_closed called after error."),
+ State::Failed => return Poll::Ready(Err(Error::Closed)),
  }
  }
  }
@@ -347,4 +346,31 @@ mod tests {
  assert_eq!(&expected[..], &actual[..])
  });
  }
+
+ #[test]
+ fn error_does_not_panic() {
+ task::block_on(async {
+ let sink = make_sink(io::stdout(), |mut _stdout, _action| async move {
+ Err(io::Error::new(io::ErrorKind::Other, "oh no"))
+ });
+
+ futures::pin_mut!(sink);
+
+ let result = sink.send("hello").await;
+ match result {
+ Err(crate::quicksink::Error::Send(e)) => {
+ assert_eq!(e.kind(), io::ErrorKind::Other);
+ assert_eq!(e.to_string(), "oh no")
+ }
+ _ => panic!("unexpected result: {:?}", result),
+ };
+
+ // Call send again, expect not to panic.
+ let result = sink.send("hello").await;
+ match result {
+ Err(crate::quicksink::Error::Closed) => {}
+ _ => panic!("unexpected result: {:?}", result),
+ };
+ })
+ }
 }