Add Vec<u8> websocket wrapper.

src/network/mod.rs

@@ -1,3 +1,4 @@
 pub mod client;
 pub mod packet;
 pub mod server;
+pub mod websocket;

src/network/websocket.rs

@@ -0,0 +1,73 @@
+use async_tungstenite::tungstenite::{Error as WsError, Message};
+use async_tungstenite::WebSocketStream;
+use futures::prelude::*;
+use futures::sink::Sink;
+use futures::stream::Stream;
+use futures::task::Context;
+use futures::task::Poll;
+use std::pin::Pin;
+
+pub struct WebSocketConnection<S> {
+    stream: Pin<Box<WebSocketStream<S>>>,
+}
+
+impl<S> Sink<Vec<u8>> for WebSocketConnection<S>
+where
+    S: AsyncRead + AsyncWrite + Unpin,
+{
+    type Error = WsError;
+
+    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
+        // Safe, as we will not move self_.
+        let self_ = unsafe { self.get_unchecked_mut() };
+
+        Pin::as_mut(&mut self_.stream).poll_ready(cx)
+    }
+
+    fn start_send(self: Pin<&mut Self>, item: Vec<u8>) -> Result<(), Self::Error> {
+        // Safe, as we will not move self_.
+        let self_ = unsafe { self.get_unchecked_mut() };
+
+        Pin::as_mut(&mut self_.stream).start_send(Message::Binary(item))
+    }
+
+    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
+        // Safe, as we will not move self_.
+        let self_ = unsafe { self.get_unchecked_mut() };
+
+        Pin::as_mut(&mut self_.stream).poll_flush(cx)
+    }
+
+    fn poll_close(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
+        // Safe, as we will not move self_.
+        let self_ = unsafe { self.get_unchecked_mut() };
+
+        Pin::as_mut(&mut self_.stream).poll_close(cx)
+    }
+}
+
+impl<S> Stream for WebSocketConnection<S>
+where
+    S: AsyncRead + AsyncWrite + Unpin,
+{
+    type Item = Result<Vec<u8>, WsError>;
+
+    fn poll_next(
+        self: Pin<&mut Self>,
+        cx: &mut Context,
+    ) -> Poll<Option<<WebSocketConnection<S> as Stream>::Item>> {
+        // Safe, as we will not move self_.
+        let self_ = unsafe { self.get_unchecked_mut() };
+
+        match Pin::as_mut(&mut self_.stream).poll_next(cx) {
+            Poll::Ready(Some(Ok(Message::Binary(data)))) => Poll::Ready(Some(Ok(data))),
+            Poll::Ready(Some(Ok(_))) => unsafe {
+                // Ignore other message types, try again.
+                Pin::new_unchecked(self_).poll_next(cx)
+            },
+            Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e))),
+            Poll::Ready(None) => Poll::Ready(None),
+            Poll::Pending => Poll::Pending,
+        }
+    }
+}