nxtquic-api 0.1.3

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

use std::pin::Pin;
use std::sync::atomic::{AtomicI32, AtomicU64, Ordering};
use std::sync::Arc;
use std::task::{Context, Poll};
use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt, ReadBuf};
use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender};

pub(crate) enum WriteCommand {
    Data {
        stream_id: u64,
        offset: u64,
        data: Vec<u8>,
        fin: bool,
    },
    Reset {
        stream_id: u64,
        error_code: u64,
    },
    StopSending {
        stream_id: u64,
        error_code: u64,
    },
    Priority {
        stream_id: u64,
        priority: i32,
    },
}

/// 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,
    priority: Arc<AtomicI32>,
    stopped_reason: Arc<AtomicU64>,
}

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,
                priority: Arc::new(AtomicI32::new(0)),
                stopped_reason: Arc::new(AtomicU64::new(u64::MAX)),
            },
            RecvStream {
                rx,
                pending: None,
                offset: 0,
                stream_id: 0,
                priority: Arc::new(AtomicI32::new(0)),
                reset_reason: Arc::new(AtomicU64::new(u64::MAX)),
                network_tx: None,
            },
        )
    }

    pub(crate) fn network(tx: UnboundedSender<WriteCommand>, stream_id: u64) -> Self {
        Self {
            tx: None,
            network_tx: Some(tx),
            stream_id,
            offset: 0,
            priority: Arc::new(AtomicI32::new(0)),
            stopped_reason: Arc::new(AtomicU64::new(u64::MAX)),
        }
    }

    /// Returns the unique QUIC Stream ID (RFC 9000 §2.1).
    pub fn stream_id(&self) -> u64 {
        self.stream_id
    }

    /// Sets the stream priority for scheduling (RFC 9000 §2.3).
    pub fn set_priority(&mut self, priority: i32) {
        self.priority.store(priority, Ordering::Release);
        if let Some(tx) = &self.network_tx {
            let _ = tx.send(WriteCommand::Priority {
                stream_id: self.stream_id,
                priority,
            });
        }
    }

    /// Gets the current stream priority.
    pub fn priority(&self) -> i32 {
        self.priority.load(Ordering::Acquire)
    }

    /// Explicitly finishes the stream by sending a QUIC FIN (RFC 9000 §19.8).
    pub async fn finish(&mut self) -> std::io::Result<()> {
        self.shutdown().await
    }

    /// Abruptly terminates sending on this stream with an application error code (RFC 9000 §19.4).
    pub async fn reset(&mut self, error_code: u64) -> std::io::Result<()> {
        if let Some(tx) = self.network_tx.take() {
            let _ = tx.send(WriteCommand::Reset {
                stream_id: self.stream_id,
                error_code,
            });
        }
        if let Some(tx) = self.tx.take() {
            let _ = tx.send(None);
        }
        Ok(())
    }

    /// Returns the error code if the peer sent a `STOP_SENDING` frame (RFC 9000 §19.5).
    pub fn stopped(&self) -> Option<u64> {
        let code = self.stopped_reason.load(Ordering::Acquire);
        if code == u64::MAX {
            None
        } else {
            Some(code)
        }
    }

    /// Scatter-gather writes multiple contiguous byte chunks to the stream.
    pub async fn write_chunks(&mut self, bufs: &[bytes::Bytes]) -> std::io::Result<usize> {
        let mut total = 0;
        for buf in bufs {
            self.write_all(buf).await?;
            total += buf.len();
        }
        Ok(total)
    }
}

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,
    stream_id: u64,
    priority: Arc<AtomicI32>,
    reset_reason: Arc<AtomicU64>,
    network_tx: Option<UnboundedSender<WriteCommand>>,
}

impl RecvStream {
    pub(crate) fn from_receiver(rx: UnboundedReceiver<Option<Vec<u8>>>) -> Self {
        Self {
            rx,
            pending: None,
            offset: 0,
            stream_id: 0,
            priority: Arc::new(AtomicI32::new(0)),
            reset_reason: Arc::new(AtomicU64::new(u64::MAX)),
            network_tx: None,
        }
    }

    pub(crate) fn from_receiver_with_id(
        rx: UnboundedReceiver<Option<Vec<u8>>>,
        stream_id: u64,
        network_tx: Option<UnboundedSender<WriteCommand>>,
    ) -> Self {
        Self {
            rx,
            pending: None,
            offset: 0,
            stream_id,
            priority: Arc::new(AtomicI32::new(0)),
            reset_reason: Arc::new(AtomicU64::new(u64::MAX)),
            network_tx,
        }
    }

    /// Returns the unique QUIC Stream ID (RFC 9000 §2.1).
    pub fn stream_id(&self) -> u64 {
        self.stream_id
    }

    /// Sets the stream priority for scheduling (RFC 9000 §2.3).
    pub fn set_priority(&mut self, priority: i32) {
        self.priority.store(priority, Ordering::Release);
        if let Some(tx) = &self.network_tx {
            let _ = tx.send(WriteCommand::Priority {
                stream_id: self.stream_id,
                priority,
            });
        }
    }

    /// Gets the current stream priority.
    pub fn priority(&self) -> i32 {
        self.priority.load(Ordering::Acquire)
    }

    /// Signals the peer to stop sending on this stream with an application error code (RFC 9000 §19.5).
    pub async fn stop_sending(&mut self, error_code: u64) -> std::io::Result<()> {
        if let Some(tx) = &self.network_tx {
            let _ = tx.send(WriteCommand::StopSending {
                stream_id: self.stream_id,
                error_code,
            });
        }
        Ok(())
    }

    /// Returns the error code if the peer aborted this stream with a `RESET_STREAM` frame (RFC 9000 §19.4).
    pub fn received_reset(&self) -> Option<u64> {
        let code = self.reset_reason.load(Ordering::Acquire);
        if code == u64::MAX {
            None
        } else {
            Some(code)
        }
    }

    /// Reads an individual contiguous chunk of data without copying into an intermediate buffer.
    pub async fn read_chunk(&mut self, max: usize) -> std::io::Result<Option<bytes::Bytes>> {
        if let Some(data) = self.pending.take() {
            let remaining = &data[self.offset..];
            if !remaining.is_empty() {
                let chunk_size = remaining.len().min(max);
                let chunk = bytes::Bytes::copy_from_slice(&remaining[..chunk_size]);
                if chunk_size < remaining.len() {
                    self.pending = Some(data);
                    self.offset += chunk_size;
                } else {
                    self.offset = 0;
                }
                return Ok(Some(chunk));
            }
        }

        match self.rx.recv().await {
            Some(Some(data)) => {
                let chunk_size = data.len().min(max);
                let chunk = bytes::Bytes::copy_from_slice(&data[..chunk_size]);
                if chunk_size < data.len() {
                    self.pending = Some(data);
                    self.offset = chunk_size;
                }
                Ok(Some(chunk))
            }
            Some(None) | None => Ok(None),
        }
    }

    /// Reads into multiple buffers using vectored I/O.
    pub async fn read_chunks(&mut self, bufs: &mut [bytes::Bytes]) -> std::io::Result<usize> {
        let mut count = 0;
        for slot in bufs.iter_mut() {
            if let Some(chunk) = self.read_chunk(65536).await? {
                *slot = chunk;
                count += 1;
            } else {
                break;
            }
        }
        Ok(count)
    }
}

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,
            }
        }
    }
}