nxtquic-api 0.1.2

//! QUIC stream types implementing Tokio async I/O traits.

use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender};

pub(crate) enum WriteCommand {
    Data { stream_id: u64, offset: u64, data: Vec<u8>, fin: bool },
}

/// A stream that can be written to.
pub struct SendStream {
    tx: Option<UnboundedSender<Option<Vec<u8>>>>,
    network_tx: Option<UnboundedSender<WriteCommand>>,
    stream_id: u64,
    offset: u64,
}

impl SendStream {
    pub(crate) fn pair() -> (Self, RecvStream) {
        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
        (
            Self { tx: Some(tx), network_tx: None, stream_id: 0, offset: 0 },
            RecvStream {
                rx,
                pending: None,
                offset: 0,
            },
        )
    }

    pub(crate) fn network(tx: UnboundedSender<WriteCommand>, stream_id: u64) -> Self {
        Self { tx: None, network_tx: Some(tx), stream_id, offset: 0 }
    }
}

impl AsyncWrite for SendStream {
    fn poll_write(
        self: Pin<&mut Self>,
        _cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<std::io::Result<usize>> {
        let this = self.get_mut();
        if let Some(tx) = this.network_tx.as_ref() {
            let offset = this.offset;
            tx.send(WriteCommand::Data { stream_id: this.stream_id, offset, data: buf.to_vec(), fin: false })
                .map_err(|_| std::io::Error::new(std::io::ErrorKind::BrokenPipe, "connection closed"))?;
            this.offset += buf.len() as u64;
            return Poll::Ready(Ok(buf.len()));
        }
        let tx = this.tx.as_ref().ok_or_else(|| {
            std::io::Error::new(std::io::ErrorKind::BrokenPipe, "stream is finished")
        });
        let tx = match tx {
            Ok(tx) => tx,
            Err(err) => return Poll::Ready(Err(err)),
        };
        tx.send(Some(buf.to_vec()))
            .map_err(|_| std::io::Error::new(std::io::ErrorKind::BrokenPipe, "peer closed"))?;
        Poll::Ready(Ok(buf.len()))
    }

    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
        Poll::Ready(Ok(()))
    }

    fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
        let this = self.get_mut();
        if let Some(tx) = this.network_tx.take() {
            let _ = tx.send(WriteCommand::Data { stream_id: this.stream_id, offset: this.offset, data: Vec::new(), fin: true });
            return Poll::Ready(Ok(()));
        }
        if let Some(tx) = this.tx.take() {
            let _ = tx.send(None);
        }
        Poll::Ready(Ok(()))
    }
}

/// A stream that can be read from.
pub struct RecvStream {
    rx: UnboundedReceiver<Option<Vec<u8>>>,
    pending: Option<Vec<u8>>,
    offset: usize,
}

impl RecvStream {
    pub(crate) fn from_receiver(rx: UnboundedReceiver<Option<Vec<u8>>>) -> Self {
        Self { rx, pending: None, offset: 0 }
    }
}

impl AsyncRead for RecvStream {
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<std::io::Result<()>> {
        loop {
            if let Some(data) = self.pending.as_ref() {
                let remaining = &data[self.offset..];
                if remaining.is_empty() {
                    self.pending = None;
                    self.offset = 0;
                    continue;
                }
                let n = remaining.len().min(buf.remaining());
                buf.put_slice(&remaining[..n]);
                self.offset += n;
                return Poll::Ready(Ok(()));
            }
            match Pin::new(&mut self.rx).poll_recv(cx) {
                Poll::Ready(Some(Some(data))) => {
                    self.pending = Some(data);
                }
                Poll::Ready(Some(None)) | Poll::Ready(None) => return Poll::Ready(Ok(())),
                Poll::Pending => return Poll::Pending,
            }
        }
    }
}