twfss/src/network/websocket.rs

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> WebSocketConnection<S> {
    pub fn new(s: WebSocketStream<S>) -> Self {
        Self {
            stream: Box::pin(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,
        }
    }
}