Add fastwebsockets test and fix pre-handshake buffer issue

Cargo.toml

@@ -22,3 +22,4 @@ rustls-pemfile = "1.0"
 rustls = { version = "0.21", features = [ "dangerous_configuration" ] }
 ntest = "0.9"
 rstest = "0.18"
+fastwebsockets = "0.4.4"

src/lib.rs

@@ -5,6 +5,9 @@ mod connection_stream;
 mod handshake;
 mod stream;
 
+#[cfg(test)]
+mod system_test;
+
 pub use handshake::handshake_task;
 pub use stream::TlsHandshake;
 pub use stream::TlsStream;
@@ -21,22 +24,24 @@ struct TestOptions {
 }
 
 macro_rules! trace {
-    ($($args:expr),+) => {
-      #[cfg(feature="trace")]
-      {
-        println!($($args),+);
-      }
-      #[cfg(not(feature="trace"))]
-      {
-        format!($($args),+);
-      }
-    };
+  ($($args:expr),+) => {
+    #[cfg(feature="trace")]
+    {
+      println!($($args),+);
+    }
+    #[cfg(not(feature="trace"))]
+    {
+      format!($($args),+);
+    }
+  };
 }
 
 pub(crate) use trace;
 
 #[cfg(test)]
 mod tests {
+  pub use super::stream::tests::tls_pair;
+  pub use super::stream::tests::tls_pair_buffer_size;
   use rustls::client::ServerCertVerified;
   use rustls::client::ServerCertVerifier;
   use rustls::Certificate;

src/stream.rs

@@ -512,9 +512,13 @@ impl AsyncWrite for TlsStream {
               Poll::Pending
             } else {
               trace!("write buf");
-              let buf = &buf[0..remaining];
-              write_buf.extend_from_slice(buf);
-              Poll::Ready(Ok(buf.len()))
+              if buf.len() <= remaining {
+                write_buf.extend_from_slice(buf);
+                Poll::Ready(Ok(buf.len()))
+              } else {
+                write_buf.extend_from_slice(&buf[0..remaining]);
+                Poll::Ready(Ok(remaining))
+              }
             }
           } else {
             trace!("write buf");
@@ -594,7 +598,7 @@ impl Drop for TlsStream {
 }
 
 #[cfg(test)]
-mod tests {
+pub(super) mod tests {
   use super::*;
   use futures::stream::FuturesUnordered;
   use futures::FutureExt;
@@ -705,15 +709,21 @@ mod tests {
     (server, client)
   }
 
-  async fn tls_pair() -> (TlsStream, TlsStream) {
+  pub async fn tls_pair() -> (TlsStream, TlsStream) {
+    tls_pair_buffer_size(None).await
+  }
+
+  pub async fn tls_pair_buffer_size(
+    buffer_size: Option<NonZeroUsize>,
+  ) -> (TlsStream, TlsStream) {
     let (server, client) = tcp_pair().await;
     let server =
       TlsStream::new_server_side(server, server_config(&[]).into(), None);
     let client = TlsStream::new_client_side(
       client,
       client_config(&[]).into(),
       "example.com".try_into().unwrap(),
-      None,
+      buffer_size,
     );
 
     (server, client)

src/system_test/fastwebsockets.rs

@@ -0,0 +1,78 @@
+use std::num::NonZeroUsize;
+
+use fastwebsockets::Frame;
+use fastwebsockets::OpCode;
+use fastwebsockets::Payload;
+use fastwebsockets::Role;
+use fastwebsockets::WebSocket;
+use rstest::rstest;
+
+const LARGE_PAYLOAD: [u8; 48 * 1024] = [0xff; 48 * 1024];
+const SMALL_PAYLOAD: [u8; 16] = [0xff; 16];
+
+#[rstest]
+#[case(false, false, false)]
+#[case(false, false, true)]
+#[case(false, true, false)]
+#[case(false, true, true)]
+#[case(true, false, false)]
+#[case(true, false, true)]
+#[case(true, true, false)]
+#[case(true, true, true)]
+#[tokio::test]
+async fn test_fastwebsockets(
+  #[case] handshake: bool,
+  #[case] buffer_limit: bool,
+  #[case] large_payload: bool,
+) {
+  let payload = if large_payload {
+    LARGE_PAYLOAD.as_slice()
+  } else {
+    SMALL_PAYLOAD.as_slice()
+  };
+  let (mut client, mut server) = if buffer_limit {
+    crate::tests::tls_pair_buffer_size(Some(
+      NonZeroUsize::try_from(1024).unwrap(),
+    ))
+    .await
+  } else {
+    crate::tests::tls_pair().await
+  };
+  if handshake {
+    client.handshake().await.expect("failed handshake");
+    server.handshake().await.expect("failed handshake");
+  }
+
+  let a = tokio::spawn(async {
+    let mut ws = WebSocket::after_handshake(server, Role::Server);
+    ws.set_auto_close(true);
+    for _ in 0..1000 {
+      ws.write_frame(Frame::binary(Payload::Borrowed(payload)))
+        .await
+        .expect("failed to write");
+    }
+    let frame = ws.read_frame().await.expect("failed to read");
+    assert_eq!(frame.payload.len(), payload.len());
+    ws.write_frame(Frame::close(1000, &[]))
+      .await
+      .expect("failed to close");
+    let frame = ws.read_frame().await.expect("failed to read");
+    assert_eq!(frame.opcode, OpCode::Close);
+  });
+  let b = tokio::spawn(async {
+    let mut ws = WebSocket::after_handshake(client, Role::Client);
+    ws.set_auto_close(true);
+    for _ in 0..1000 {
+      let frame = ws.read_frame().await.expect("failed to read");
+      assert_eq!(frame.payload.len(), payload.len());
+    }
+    ws.write_frame(Frame::binary(Payload::Borrowed(payload)))
+      .await
+      .expect("failed to write");
+    let frame = ws.read_frame().await.expect("failed to read");
+    assert_eq!(frame.opcode, OpCode::Close);
+  });
+
+  a.await.expect("failed to join");
+  b.await.expect("failed to join");
+}

src/system_test/mod.rs

@@ -0,0 +1 @@
+mod fastwebsockets;