borer-core 0.5.9

#[allow(dead_code)]
const MAX_READ_SIZE: usize = 1024;

#[allow(unused_imports)]
use std::{pin::Pin, task::Poll};

#[cfg(feature = "websocket")]
/// `AsyncRead`/`AsyncWrite` adapter over a tungstenite websocket stream.
pub struct WebSocketCopyStream<T> {
    inner: tokio_tungstenite::WebSocketStream<T>,
}

#[cfg(feature = "websocket")]
impl<T> WebSocketCopyStream<T> {
    /// Wrap a websocket stream so it can participate in bidirectional I/O copying.
    pub fn new(inner: tokio_tungstenite::WebSocketStream<T>) -> Self {
        Self { inner }
    }
}

#[cfg(feature = "websocket")]
impl<T> tokio::io::AsyncRead for WebSocketCopyStream<T>
where
    T: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
{
    fn poll_read(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
        buf: &mut tokio::io::ReadBuf<'_>,
    ) -> Poll<std::io::Result<()>> {
        use futures::StreamExt;
        use std::io::Error;
        use std::io::ErrorKind;
        use tokio_tungstenite::tungstenite::Message;

        let me = self.get_mut();
        let next = me.inner.poll_next_unpin(cx);
        match next {
            Poll::Ready(t) => match t {
                Some(Ok(Message::Binary(b))) => {
                    //todo flow_control. buf remaining painc
                    buf.put_slice(b.as_ref());
                    Poll::Ready(Ok(()))
                }
                Some(Ok(Message::Close(_))) => {
                    Poll::Ready(Err(Error::new(ErrorKind::Other, "websocket close message")))
                }
                Some(Ok(_)) => Poll::Ready(Ok(())),
                Some(Err(e)) => Poll::Ready(Err(Error::new(ErrorKind::Other, e))),
                None => Poll::Ready(Err(Error::new(ErrorKind::Other, "websocket poll none"))),
            },
            Poll::Pending => Poll::Pending,
        }
    }
}

#[cfg(feature = "websocket")]
impl<T> tokio::io::AsyncWrite for WebSocketCopyStream<T>
where
    T: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
{
    fn poll_write(
        self: std::pin::Pin<&mut Self>,
        _cx: &mut std::task::Context<'_>,
        buf: &[u8],
    ) -> Poll<Result<usize, std::io::Error>> {
        use bytes::Bytes;
        use futures::SinkExt;
        use tokio_tungstenite::tungstenite::Message;
        let me = self.get_mut();
        let read_buf = &buf[0..std::cmp::min(buf.len(), MAX_READ_SIZE)];
        let bytes = Bytes::copy_from_slice(read_buf);
        let len = bytes.len();
        me.inner.start_send_unpin(Message::Binary(bytes)).ok();

        Poll::Ready(Ok(len))
    }

    fn poll_flush(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> Poll<Result<(), std::io::Error>> {
        use futures::Sink;
        use std::io::Error;
        use std::io::ErrorKind;

        let me = self.get_mut();
        let stream = &mut me.inner;
        match Pin::new(stream).poll_flush(cx) {
            Poll::Ready(Ok(_)) => Poll::Ready(Ok(())),
            Poll::Ready(Err(_)) => Poll::Ready(Err(Error::from(ErrorKind::Other))),
            Poll::Pending => Poll::Pending,
        }
    }

    fn poll_shutdown(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> Poll<Result<(), std::io::Error>> {
        use futures::Sink;
        use std::io::Error;
        use std::io::ErrorKind;

        let me = self.get_mut();
        let stream = &mut me.inner;
        match Pin::new(stream).poll_close(cx) {
            Poll::Ready(Ok(_)) => Poll::Ready(Ok(())),
            Poll::Ready(Err(_)) => Poll::Ready(Err(Error::from(ErrorKind::Other))),
            Poll::Pending => Poll::Pending,
        }
    }
}

#[cfg(feature = "axum")]
use axum::extract::ws::Message as AxumMessage;
#[cfg(feature = "axum")]
use axum::extract::ws::WebSocket;

#[cfg(feature = "axum")]
/// `AsyncRead`/`AsyncWrite` adapter over an axum websocket.
pub struct AxumWebSocketCopyStream {
    inner: WebSocket,
}

#[cfg(feature = "axum")]
impl AxumWebSocketCopyStream {
    /// Wrap an axum websocket so it can participate in bidirectional I/O copying.
    pub fn new(inner: WebSocket) -> Self {
        Self { inner }
    }
}

#[cfg(feature = "axum")]
impl tokio::io::AsyncRead for AxumWebSocketCopyStream {
    fn poll_read(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
        buf: &mut tokio::io::ReadBuf<'_>,
    ) -> Poll<std::io::Result<()>> {
        use futures::StreamExt;
        use std::io::Error;
        use std::io::ErrorKind;

        let me = self.get_mut();
        let next = me.inner.poll_next_unpin(cx);
        match next {
            Poll::Ready(t) => match t {
                Some(Ok(AxumMessage::Binary(b))) => {
                    //todo flow_control. buf remaining painc
                    buf.put_slice(b.as_ref());
                    Poll::Ready(Ok(()))
                }
                Some(Ok(AxumMessage::Close(_))) => {
                    Poll::Ready(Err(Error::new(ErrorKind::Other, "websocket close message")))
                }
                Some(Ok(_)) => Poll::Ready(Ok(())),
                Some(Err(e)) => Poll::Ready(Err(Error::new(ErrorKind::Other, e))),
                None => Poll::Ready(Err(Error::new(ErrorKind::Other, "websocket poll none"))),
            },
            Poll::Pending => Poll::Pending,
        }
    }
}

#[cfg(feature = "axum")]
impl tokio::io::AsyncWrite for AxumWebSocketCopyStream {
    fn poll_write(
        self: std::pin::Pin<&mut Self>,
        _cx: &mut std::task::Context<'_>,
        buf: &[u8],
    ) -> Poll<Result<usize, std::io::Error>> {
        use bytes::Bytes;
        use futures::SinkExt;

        let me = self.get_mut();
        let read_buf = &buf[0..std::cmp::min(buf.len(), MAX_READ_SIZE)];
        let bytes = Bytes::copy_from_slice(read_buf);
        let len = bytes.len();
        me.inner.start_send_unpin(AxumMessage::Binary(bytes)).ok();

        Poll::Ready(Ok(len))
    }

    fn poll_flush(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> Poll<Result<(), std::io::Error>> {
        use futures::Sink;
        use std::io::Error;
        use std::io::ErrorKind;

        let me = self.get_mut();
        let stream = &mut me.inner;
        match Pin::new(stream).poll_flush(cx) {
            Poll::Ready(Ok(_)) => Poll::Ready(Ok(())),
            Poll::Ready(Err(_)) => Poll::Ready(Err(Error::from(ErrorKind::Other))),
            Poll::Pending => Poll::Pending,
        }
    }

    fn poll_shutdown(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> Poll<Result<(), std::io::Error>> {
        use futures::Sink;
        use std::io::Error;
        use std::io::ErrorKind;

        let me = self.get_mut();
        let stream = &mut me.inner;
        match Pin::new(stream).poll_close(cx) {
            Poll::Ready(Ok(_)) => Poll::Ready(Ok(())),
            Poll::Ready(Err(_)) => Poll::Ready(Err(Error::from(ErrorKind::Other))),
            Poll::Pending => Poll::Pending,
        }
    }
}

#[cfg(all(test, feature = "websocket"))]
mod tests {
    use futures::{SinkExt, StreamExt};
    use tokio::io::{AsyncReadExt, AsyncWriteExt};
    use tokio_tungstenite::{
        WebSocketStream,
        tungstenite::{Message, protocol::Role},
    };

    use super::{MAX_READ_SIZE, WebSocketCopyStream};

    async fn websocket_pair() -> (
        WebSocketCopyStream<tokio::io::DuplexStream>,
        WebSocketStream<tokio::io::DuplexStream>,
    ) {
        let (server_io, client_io) = tokio::io::duplex(16 * 1024);
        let (server_ws, client_ws) = tokio::join!(
            WebSocketStream::from_raw_socket(server_io, Role::Server, None),
            WebSocketStream::from_raw_socket(client_io, Role::Client, None),
        );

        (WebSocketCopyStream::new(server_ws), client_ws)
    }

    #[tokio::test]
    async fn websocket_copy_stream_reads_binary_messages() {
        let (mut stream, mut peer) = websocket_pair().await;

        peer.send(Message::Binary(b"ping".to_vec().into()))
            .await
            .unwrap();

        let mut buf = [0u8; 4];
        stream.read_exact(&mut buf).await.unwrap();

        assert_eq!(&buf, b"ping");
    }

    #[tokio::test]
    async fn websocket_copy_stream_reports_close_as_error() {
        let (mut stream, mut peer) = websocket_pair().await;

        peer.close(None).await.unwrap();

        let err = stream.read_u8().await.unwrap_err();
        assert!(err.to_string().contains("websocket close message"));
    }

    #[tokio::test]
    async fn websocket_copy_stream_writes_binary_messages_capped_to_max_read_size() {
        let (mut stream, mut peer) = websocket_pair().await;
        let payload = vec![7u8; MAX_READ_SIZE + 128];

        stream.write_all(&payload).await.unwrap();
        stream.flush().await.unwrap();

        let message = peer.next().await.unwrap().unwrap();
        match message {
            Message::Binary(buf) => {
                assert_eq!(buf.len(), MAX_READ_SIZE);
                assert!(buf.iter().all(|b| *b == 7));
            }
            other => panic!("expected binary message, got {other:?}"),
        }
    }
}