websocket: Fix connection stability on decrypt messages (#393)

This PR greatly improves the WebSocket connection stability by relying
on the interval buffers of tungstenite instead of buffering at a higher
level. The fix passes through the messages to the tungstenite socket
directly.

This is a long-lasting issue (reproducible on all older versions
silently with IO errors) that manifested as a decryption error after the
state fixes:
- #325
- #327

Issue context:
- node is under stress due to handling multiple substreams
- the issue affected only long running WebSocket substreams and
manifested as an IO error from crypto/noise decoding
- tungstenite `WebSocketStream` already has a 128KiB buffer for writing
- litep2p has a **redundant** 8 KiB buffer for writing 
- litep2p buffered internally multiple packets, tunstenite accepted the
batch. I expect this creates a wrongly framed packet that fails to
decode at the crypto/noise level


## Investigation

We have noted several errors that manifested as crypto/nosie decoding
failures on our Kusama validators:
- paritytech/polkadot-sdk#8525

```rust
litep2p::crypto::noise: failed to decrypt message error=Decrypt
```

Upon further investigation, the errors affected only WebSocket
connections. The issue could be reproduced by running a local node in
Kusama with more than 500 peers in and out. As well as running
subp2p-explorer with adjusted protocols:

```yaml
2025-05-15T14:58:08.095961Z ERROR {peer_id=peer_id=12D3KooWGsDvWrbApFTCpF8h7YCKHuvJbok6HAq5ZnPgE9LGWnsv}:
litep2p::crypto::noise: failed to decrypt message for bigger buffers error=Decrypt peer=PeerId("12D3KooWSa5SbCHGKpNeSs3Qak2TrM5gTkEBrPfvo6TyxhUpEHeu")

2025-05-15T14:58:08.096419Z DEBUG 
{peer_id=peer_id=12D3KooWGsDvWrbApFTCpF8h7YCKHuvJbok6HAq5ZnPgE9LGWnsv}:
litep2p::websocket::connection: connection closed with error peer=PeerId("12D3KooWSa5SbCHGKpNeSs3Qak2TrM5gTkEBrPfvo6TyxhUpEHeu") error=Decode(Io(Custom { kind: Other, error: "failed to decrypt message bigger buffers: decrypt error 12D3KooWSa5SbCHGKpNeSs3Qak2TrM5gTkEBrPfvo6TyxhUpEHeu" }))
```



The issue also reproduced on the zombinet PR, which uses litep2p:
- paritytech/polkadot-sdk#8461

```yaml
2025-05-14 09:37:30.805  INFO tokio-runtime-worker sync: Warp sync is complete, continuing with state sync.    

2025-05-14 09:37:33.189 ERROR tokio-runtime-worker litep2p::crypto::noise: failed to decrypt message error=Decrypt
2025-05-14 09:37:33.283 ERROR tokio-runtime-worker litep2p::crypto::noise: failed to decrypt message error=Decrypt
2025-05-14 09:37:34.764 ERROR tokio-runtime-worker litep2p::crypto::noise: failed to decrypt message error=Decrypt
	
2025-05-14 09:37:35.656  INFO tokio-runtime-worker substrate: ⚙️  State sync, Downloading state, 22%, 2.21 Mib (0 peers), best: #0 (0xc5e7…d059), finalized #0 (0xc5e7…d059), ⬇ 707.8kiB/s ⬆ 0.5kiB/s    
	
2025-05-14 09:37:40.657  INFO tokio-runtime-worker substrate: ⚙️  State sync, Downloading state, 22%, 2.21 Mib (3 peers), best: #0 (0xc5e7…d059), finalized #0 (0xc5e7…d059), ⬇ 1.0kiB/s ⬆ 1.0kiB/s    
```

## Testing Done

### Performance

Tested the performance with litep2p-perf using the following branch:
-
https://github.com/lexnv/litep2p-perf/compare/lexnv/websocket-tests?expand=1

| Status     | Data Size | Time (s) | Bandwidth (Mbit/s) |
|------------|-----------|----------|-------------------|
| **Before** |           |          |                   |
| Uploaded   | 256.00 MiB| 15.1152  | 135.49            |
| Downloaded | 256.00 MiB| 13.2296  | 154.80            |
| **After**  |           |          |                   |
| Uploaded   | 256.00 MiB| 15.7178  | 130.30            |
| Downloaded | 256.00 MiB| 13.2435  | 154.64            |

From the performance table, we are within 3% of the original buggy
implementation. I would lean towards a normal variation in our results.
Therefore, the performance remains unimpacted.

### Repro Case

Have added a custom user protocol as part of our testing to filter out
these errors.

- The protocol opens 16 outbound substreams on the connection
established event. Therefore, it will handle 16 outbound substreams and
16 inbound substreams
- The outbound substreams will push a configurable number of packets,
each of size 128 bytes, to the remote peer. While the inbound substreams
will read the same number of packets from the remote peer.
 
Before this PR, the TCP was unaffected and the websocket reproduces the
decrypt failure. After this PR, the test passes.


Closes: paritytech/polkadot-sdk#8525

---------

Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io>

Author: lexnv

src/lib.rs

@@ -20,6 +20,7 @@
 
 #![allow(clippy::single_match)]
 #![allow(clippy::result_large_err)]
+#![allow(clippy::large_enum_variant)]
 #![allow(clippy::redundant_pattern_matching)]
 #![allow(clippy::type_complexity)]
 #![allow(clippy::result_unit_err)]
@@ -80,7 +81,7 @@ pub mod yamux;
 
 mod bandwidth;
 mod multistream_select;
-mod utils;
+pub mod utils;
 
 #[cfg(test)]
 mod mock;

src/multistream_select/negotiated.rs

@@ -370,6 +370,6 @@ impl From<NegotiationError> for io::Error {
         if let NegotiationError::ProtocolError(e) = err {
             return e.into();
         }
-        io::Error::new(io::ErrorKind::Other, err)
+        io::Error::other(err)
     }
 }

src/transport/websocket/stream.rs

@@ -21,7 +21,7 @@
 //! Stream implementation for `tokio_tungstenite::WebSocketStream` that implements
 //! `AsyncRead + AsyncWrite`
 
-use bytes::{Buf, Bytes, BytesMut};
+use bytes::{Buf, Bytes};
 use futures::{SinkExt, StreamExt};
 use tokio::io::{AsyncRead, AsyncWrite};
 use tokio_tungstenite::{tungstenite::Message, WebSocketStream};
@@ -31,121 +31,66 @@ use std::{
     task::{Context, Poll},
 };
 
-const DEFAULT_BUF_SIZE: usize = 8 * 1024;
-
-/// Send state.
-#[derive(Debug)]
-enum State {
-    /// State is poisoned.
-    Poisoned,
-
-    /// Sink is accepting input.
-    ReadyToSend,
-
-    /// Flush is pending for the sink.
-    FlushPending,
-}
+const LOG_TARGET: &str = "litep2p::transport::websocket::stream";
 
 /// Buffered stream which implements `AsyncRead + AsyncWrite`
 #[derive(Debug)]
 pub(super) struct BufferedStream<S: AsyncRead + AsyncWrite + Unpin> {
-    /// Write buffer.
-    write_buffer: BytesMut,
-
     /// Read buffer.
     ///
     /// The buffer is taken directly from the WebSocket stream.
     read_buffer: Bytes,
 
     /// Underlying WebSocket stream.
     stream: WebSocketStream<S>,
-
-    /// Read state.
-    state: State,
 }
 
 impl<S: AsyncRead + AsyncWrite + Unpin> BufferedStream<S> {
     /// Create new [`BufferedStream`].
     pub(super) fn new(stream: WebSocketStream<S>) -> Self {
         Self {
-            write_buffer: BytesMut::with_capacity(DEFAULT_BUF_SIZE),
             read_buffer: Bytes::new(),
             stream,
-            state: State::ReadyToSend,
         }
     }
 }
 
 impl<S: AsyncRead + AsyncWrite + Unpin> futures::AsyncWrite for BufferedStream<S> {
     fn poll_write(
         mut self: Pin<&mut Self>,
-        _cx: &mut Context<'_>,
+        cx: &mut Context<'_>,
         buf: &[u8],
     ) -> Poll<std::io::Result<usize>> {
-        self.write_buffer.extend_from_slice(buf);
-
-        Poll::Ready(Ok(buf.len()))
-    }
+        match futures::ready!(self.stream.poll_ready_unpin(cx)) {
+            Ok(()) => {
+                let message = Message::Binary(Bytes::copy_from_slice(buf));
 
-    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
-        if self.write_buffer.is_empty() {
-            return self
-                .stream
-                .poll_ready_unpin(cx)
-                .map_err(|_| std::io::ErrorKind::UnexpectedEof.into());
-        }
-
-        loop {
-            match std::mem::replace(&mut self.state, State::Poisoned) {
-                State::ReadyToSend => {
-                    match self.stream.poll_ready_unpin(cx) {
-                        Poll::Ready(Ok(())) => {}
-                        Poll::Ready(Err(_error)) =>
-                            return Poll::Ready(Err(std::io::ErrorKind::UnexpectedEof.into())),
-                        Poll::Pending => {
-                            self.state = State::ReadyToSend;
-                            return Poll::Pending;
-                        }
-                    }
-
-                    let message = std::mem::take(&mut self.write_buffer);
-                    match self.stream.start_send_unpin(Message::Binary(message.freeze())) {
-                        Ok(()) => {}
-                        Err(_error) =>
-                            return Poll::Ready(Err(std::io::ErrorKind::UnexpectedEof.into())),
-                    }
-
-                    // Transition to flush pending state.
-                    self.state = State::FlushPending;
-                    continue;
+                if let Err(err) = self.stream.start_send_unpin(message) {
+                    tracing::debug!(target: LOG_TARGET, "Error during start send: {:?}", err);
+                    return Poll::Ready(Err(std::io::ErrorKind::UnexpectedEof.into()));
                 }
 
-                State::FlushPending => {
-                    match self.stream.poll_flush_unpin(cx) {
-                        Poll::Ready(Ok(())) => {}
-                        Poll::Ready(Err(_error)) =>
-                            return Poll::Ready(Err(std::io::ErrorKind::UnexpectedEof.into())),
-                        Poll::Pending => {
-                            self.state = State::FlushPending;
-                            return Poll::Pending;
-                        }
-                    }
-
-                    self.state = State::ReadyToSend;
-                    self.write_buffer = BytesMut::with_capacity(DEFAULT_BUF_SIZE);
-                    return Poll::Ready(Ok(()));
-                }
-                State::Poisoned =>
-                    return Poll::Ready(Err(std::io::ErrorKind::UnexpectedEof.into())),
+                Poll::Ready(Ok(buf.len()))
+            }
+            Err(err) => {
+                tracing::debug!(target: LOG_TARGET, "Error during poll ready: {:?}", err);
+                Poll::Ready(Err(std::io::ErrorKind::UnexpectedEof.into()))
             }
         }
     }
 
+    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
+        self.stream.poll_flush_unpin(cx).map_err(|err| {
+            tracing::debug!(target: LOG_TARGET, "Error during poll flush: {:?}", err);
+            std::io::ErrorKind::UnexpectedEof.into()
+        })
+    }
+
     fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
-        match futures::ready!(self.stream.poll_close_unpin(cx)) {
-            Ok(_) => Poll::Ready(Ok(())),
-            Err(_) => Poll::Ready(Err(std::io::ErrorKind::PermissionDenied.into())),
-        }
+        self.stream.poll_close_unpin(cx).map_err(|err| {
+            tracing::debug!(target: LOG_TARGET, "Error during poll close: {:?}", err);
+            std::io::ErrorKind::PermissionDenied.into()
+        })
     }
 }
 
@@ -183,7 +128,7 @@ impl<S: AsyncRead + AsyncWrite + Unpin> futures::AsyncRead for BufferedStream<S>
 #[cfg(test)]
 mod tests {
     use super::*;
-    use futures::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
+    use futures::{AsyncRead, AsyncReadExt, AsyncWriteExt};
     use tokio::io::DuplexStream;
     use tokio_tungstenite::{tungstenite::protocol::Role, WebSocketStream};
 
@@ -203,7 +148,6 @@ mod tests {
 
         let bytes_written = stream.write(data).await.unwrap();
         assert_eq!(bytes_written, data.len());
-        assert_eq!(&stream.write_buffer[..], data);
     }
 
     #[tokio::test]
@@ -253,38 +197,6 @@ mod tests {
         };
     }
 
-    #[tokio::test]
-    async fn test_poisoned_state() {
-        let (mut stream, server) = create_test_stream().await;
-        drop(server);
-
-        stream.state = State::Poisoned;
-
-        let mut buffer = [0u8; 10];
-        let result = stream.read(&mut buffer).await;
-        match result {
-            Err(error) => if error.kind() == std::io::ErrorKind::UnexpectedEof {},
-            state => panic!("Unexpected state {state:?}"),
-        };
-
-        let mut cx = std::task::Context::from_waker(futures::task::noop_waker_ref());
-        let mut pin_stream = Pin::new(&mut stream);
-
-        // Messages are buffered internally, the socket is not touched.
-        match pin_stream.as_mut().poll_write(&mut cx, &mut buffer) {
-            Poll::Ready(Ok(10)) => {}
-            state => panic!("Unexpected state {state:?}"),
-        }
-        // Socket is poisoned, the flush will fail.
-        match pin_stream.poll_flush(&mut cx) {
-            Poll::Ready(Err(error)) =>
-                if error.kind() == std::io::ErrorKind::UnexpectedEof {
-                    return;
-                },
-            state => panic!("Unexpected state {state:?}"),
-        }
-    }
-
     #[tokio::test]
     async fn test_read_poll_pending() {
         let (mut stream, mut _server) = create_test_stream().await;

tests/connection/mod.rs

@@ -44,6 +44,8 @@ use crate::common::{add_transport, Transport};
 
 #[cfg(test)]
 mod protocol_dial_invalid_address;
+#[cfg(test)]
+mod stability;
 
 #[tokio::test]
 async fn two_litep2ps_work_tcp() {