poll_* methods to support custom futures implementations

Cargo.toml

@@ -28,11 +28,15 @@ path = "examples/axum.rs"
 required-features = ["upgrade", "with_axum"]
 
 [dependencies]
-tokio = { version = "1.25.0",  default-features = false, features = ["io-util"] }
+tokio = { version = "1.25.0", default-features = false, features = ["io-util"] }
 simdutf8 = { version = "0.1.4", optional = true }
 hyper-util = { version = "0.1.0", features = ["tokio"], optional = true }
 http-body-util = { version = "0.1.0", optional = true }
-hyper = { version = "1", features = ["http1", "server", "client"], optional = true }
+hyper = { version = "1", features = [
+    "http1",
+    "server",
+    "client",
+], optional = true }
 pin-project = { version = "1.0.8", optional = true }
 base64 = { version = "0.21.0", optional = true }
 sha1 = { version = "0.10.5", optional = true }
@@ -45,11 +49,20 @@ bytes = "1.5.0"
 axum-core = { version = "0.4.3", optional = true }
 http = { version = "1", optional = true }
 async-trait = { version = "0.1", optional = true }
+tokio-util = { version = "0.7", features = ["codec", "io"] }
+futures = { version = "0.3", default-features = false, features = ["std"] }
 
 [features]
 default = ["simd"]
 simd = ["simdutf8/aarch64_neon"]
-upgrade = ["hyper", "pin-project", "base64", "sha1", "hyper-util", "http-body-util"]
+upgrade = [
+    "hyper",
+    "pin-project",
+    "base64",
+    "sha1",
+    "hyper-util",
+    "http-body-util",
+]
 unstable-split = []
 # Axum integration
 with_axum = ["axum-core", "http", "async-trait"]

src/frame.rs

@@ -14,7 +14,7 @@
 
 use tokio::io::AsyncWriteExt;
 
-use bytes::BytesMut;
+use bytes::{BufMut, BytesMut};
 use core::ops::Deref;
 
 use crate::WebSocketError;
@@ -259,26 +259,27 @@ impl<'f> Frame<'f> {
   /// # Panics
   ///
   /// This method panics if the head buffer is not at least n-bytes long, where n is the size of the length field (0, 2, 4, or 10)
-  pub fn fmt_head(&mut self, head: &mut [u8]) -> usize {
-    head[0] = (self.fin as u8) << 7 | (self.opcode as u8);
+  pub fn fmt_head(&mut self, mut head: impl BufMut) -> usize {
+    head.put_u8((self.fin as u8) << 7 | (self.opcode as u8));
+
+    let mask_bit = if self.mask.is_some() { 0x80 } else { 0x0 };
 
     let len = self.payload.len();
     let size = if len < 126 {
-      head[1] = len as u8;
+      head.put_u8(len as u8 | mask_bit);
       2
     } else if len < 65536 {
-      head[1] = 126;
-      head[2..4].copy_from_slice(&(len as u16).to_be_bytes());
+      head.put_u8(126u8 | mask_bit);
+      head.put_slice(&(len as u16).to_be_bytes());
       4
     } else {
-      head[1] = 127;
-      head[2..10].copy_from_slice(&(len as u64).to_be_bytes());
+      head.put_u8(127u8 | mask_bit);
+      head.put_slice(&(len as u64).to_be_bytes());
       10
     };
 
     if let Some(mask) = self.mask {
-      head[1] |= 0x80;
-      head[size..size + 4].copy_from_slice(&mask);
+      head.put_slice(&mask);
       size + 4
     } else {
       size
@@ -295,7 +296,7 @@ impl<'f> Frame<'f> {
     use std::io::IoSlice;
 
     let mut head = [0; MAX_HEAD_SIZE];
-    let size = self.fmt_head(&mut head);
+    let size = self.fmt_head(&mut head[..]);
 
     let total = size + self.payload.len();
 
@@ -330,7 +331,7 @@ impl<'f> Frame<'f> {
     let len = self.payload.len();
     reserve_enough(buf, len + MAX_HEAD_SIZE);
 
-    let size = self.fmt_head(buf);
+    let size = self.fmt_head(&mut *buf);
     buf[size..size + len].copy_from_slice(&self.payload);
     &buf[..size + len]
   }