snix/tvix/nix-compat/src/wire/ser/writer.rs

use std::fmt::{self, Write as _};
use std::future::poll_fn;
use std::io::{self, Cursor};
use std::pin::Pin;
use std::task::{ready, Context, Poll};

use bytes::{Buf, BufMut, BytesMut};
use pin_project_lite::pin_project;
use tokio::io::{AsyncWrite, AsyncWriteExt};

use crate::wire::{padding_len, ProtocolVersion, EMPTY_BYTES};

use super::{Error, NixWrite};

pub struct NixWriterBuilder {
    buf: Option<BytesMut>,
    reserved_buf_size: usize,
    max_buf_size: usize,
    version: ProtocolVersion,
}

impl Default for NixWriterBuilder {
    fn default() -> Self {
        Self {
            buf: Default::default(),
            reserved_buf_size: 8192,
            max_buf_size: 8192,
            version: Default::default(),
        }
    }
}

impl NixWriterBuilder {
    pub fn set_buffer(mut self, buf: BytesMut) -> Self {
        self.buf = Some(buf);
        self
    }

    pub fn set_reserved_buf_size(mut self, size: usize) -> Self {
        self.reserved_buf_size = size;
        self
    }

    pub fn set_max_buf_size(mut self, size: usize) -> Self {
        self.max_buf_size = size;
        self
    }

    pub fn set_version(mut self, version: ProtocolVersion) -> Self {
        self.version = version;
        self
    }

    pub fn build<W>(self, writer: W) -> NixWriter<W> {
        let buf = self
            .buf
            .unwrap_or_else(|| BytesMut::with_capacity(self.max_buf_size));
        NixWriter {
            buf,
            inner: writer,
            reserved_buf_size: self.reserved_buf_size,
            max_buf_size: self.max_buf_size,
            version: self.version,
        }
    }
}

pin_project! {
    pub struct NixWriter<W> {
        #[pin]
        inner: W,
        buf: BytesMut,
        reserved_buf_size: usize,
        max_buf_size: usize,
        version: ProtocolVersion,
    }
}

impl NixWriter<Cursor<Vec<u8>>> {
    pub fn builder() -> NixWriterBuilder {
        NixWriterBuilder::default()
    }
}

impl<W> NixWriter<W>
where
    W: AsyncWriteExt,
{
    pub fn new(writer: W) -> NixWriter<W> {
        NixWriter::builder().build(writer)
    }

    pub fn buffer(&self) -> &[u8] {
        &self.buf[..]
    }

    pub fn set_version(&mut self, version: ProtocolVersion) {
        self.version = version;
    }

    /// Remaining capacity in internal buffer
    pub fn remaining_mut(&self) -> usize {
        self.buf.capacity() - self.buf.len()
    }

    fn poll_flush_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
        let mut this = self.project();
        while !this.buf.is_empty() {
            let n = ready!(this.inner.as_mut().poll_write(cx, &this.buf[..]))?;
            if n == 0 {
                return Poll::Ready(Err(io::Error::new(
                    io::ErrorKind::WriteZero,
                    "failed to write the buffer",
                )));
            }
            this.buf.advance(n);
        }
        Poll::Ready(Ok(()))
    }
}

impl<W> NixWriter<W>
where
    W: AsyncWriteExt + Unpin,
{
    async fn flush_buf(&mut self) -> Result<(), io::Error> {
        let mut s = Pin::new(self);
        poll_fn(move |cx| s.as_mut().poll_flush_buf(cx)).await
    }
}

impl<W> AsyncWrite for NixWriter<W>
where
    W: AsyncWrite,
{
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<Result<usize, io::Error>> {
        // Flush
        if self.remaining_mut() < buf.len() {
            ready!(self.as_mut().poll_flush_buf(cx))?;
        }
        let this = self.project();
        if buf.len() > this.buf.capacity() {
            this.inner.poll_write(cx, buf)
        } else {
            this.buf.put_slice(buf);
            Poll::Ready(Ok(buf.len()))
        }
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
        ready!(self.as_mut().poll_flush_buf(cx))?;
        self.project().inner.poll_flush(cx)
    }

    fn poll_shutdown(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Result<(), io::Error>> {
        ready!(self.as_mut().poll_flush_buf(cx))?;
        self.project().inner.poll_shutdown(cx)
    }
}

impl<W> NixWrite for NixWriter<W>
where
    W: AsyncWrite + Send + Unpin,
{
    type Error = io::Error;

    fn version(&self) -> ProtocolVersion {
        self.version
    }

    async fn write_number(&mut self, value: u64) -> Result<(), Self::Error> {
        let mut buf = [0u8; 8];
        BufMut::put_u64_le(&mut &mut buf[..], value);
        self.write_all(&buf).await
    }

    async fn write_slice(&mut self, buf: &[u8]) -> Result<(), Self::Error> {
        let padding = padding_len(buf.len() as u64) as usize;
        self.write_value(&buf.len()).await?;
        self.write_all(buf).await?;
        if padding > 0 {
            self.write_all(&EMPTY_BYTES[..padding]).await
        } else {
            Ok(())
        }
    }

    async fn write_display<D>(&mut self, msg: D) -> Result<(), Self::Error>
    where
        D: fmt::Display + Send,
        Self: Sized,
    {
        // Ensure that buffer has space for at least reserved_buf_size bytes
        if self.remaining_mut() < self.reserved_buf_size && !self.buf.is_empty() {
            self.flush_buf().await?;
        }
        let offset = self.buf.len();
        self.buf.put_u64_le(0);
        if let Err(err) = write!(self.buf, "{}", msg) {
            self.buf.truncate(offset);
            return Err(Self::Error::unsupported_data(err));
        }
        let len = self.buf.len() - offset - 8;
        BufMut::put_u64_le(&mut &mut self.buf[offset..(offset + 8)], len as u64);
        let padding = padding_len(len as u64) as usize;
        self.write_all(&EMPTY_BYTES[..padding]).await
    }
}

#[cfg(test)]
mod test {
    use std::time::Duration;

    use hex_literal::hex;
    use rstest::rstest;
    use tokio::io::AsyncWriteExt as _;
    use tokio_test::io::Builder;

    use crate::wire::ser::NixWrite;

    use super::NixWriter;

    #[rstest]
    #[case(1, &hex!("0100 0000 0000 0000"))]
    #[case::evil(666, &hex!("9A02 0000 0000 0000"))]
    #[case::max(u64::MAX, &hex!("FFFF FFFF FFFF FFFF"))]
    #[tokio::test]
    async fn test_write_number(#[case] number: u64, #[case] buf: &[u8]) {
        let mock = Builder::new().write(buf).build();
        let mut writer = NixWriter::new(mock);

        writer.write_number(number).await.unwrap();
        assert_eq!(writer.buffer(), buf);
        writer.flush().await.unwrap();
        assert_eq!(writer.buffer(), b"");
    }

    #[rstest]
    #[case::empty(b"", &hex!("0000 0000 0000 0000"))]
    #[case::one(b")", &hex!("0100 0000 0000 0000 2900 0000 0000 0000"))]
    #[case::two(b"it", &hex!("0200 0000 0000 0000 6974 0000 0000 0000"))]
    #[case::three(b"tea", &hex!("0300 0000 0000 0000 7465 6100 0000 0000"))]
    #[case::four(b"were", &hex!("0400 0000 0000 0000 7765 7265 0000 0000"))]
    #[case::five(b"where", &hex!("0500 0000 0000 0000 7768 6572 6500 0000"))]
    #[case::six(b"unwrap", &hex!("0600 0000 0000 0000 756E 7772 6170 0000"))]
    #[case::seven(b"where's", &hex!("0700 0000 0000 0000 7768 6572 6527 7300"))]
    #[case::aligned(b"read_tea", &hex!("0800 0000 0000 0000 7265 6164 5F74 6561"))]
    #[case::more_bytes(b"read_tess", &hex!("0900 0000 0000 0000 7265 6164 5F74 6573 7300 0000 0000 0000"))]
    #[tokio::test]
    async fn test_write_slice(
        #[case] value: &[u8],
        #[case] buf: &[u8],
        #[values(1, 2, 3, 4, 5, 6, 7, 8, 9, 1024)] chunks_size: usize,
        #[values(1, 2, 3, 4, 5, 6, 7, 8, 9, 1024)] buf_size: usize,
    ) {
        let mut builder = Builder::new();
        for chunk in buf.chunks(chunks_size) {
            builder.write(chunk);
            builder.wait(Duration::ZERO);
        }
        let mock = builder.build();
        let mut writer = NixWriter::builder().set_max_buf_size(buf_size).build(mock);

        writer.write_slice(value).await.unwrap();
        writer.flush().await.unwrap();
        assert_eq!(writer.buffer(), b"");
    }

    #[rstest]
    #[case::empty("", &hex!("0000 0000 0000 0000"))]
    #[case::one(")", &hex!("0100 0000 0000 0000 2900 0000 0000 0000"))]
    #[case::two("it", &hex!("0200 0000 0000 0000 6974 0000 0000 0000"))]
    #[case::three("tea", &hex!("0300 0000 0000 0000 7465 6100 0000 0000"))]
    #[case::four("were", &hex!("0400 0000 0000 0000 7765 7265 0000 0000"))]
    #[case::five("where", &hex!("0500 0000 0000 0000 7768 6572 6500 0000"))]
    #[case::six("unwrap", &hex!("0600 0000 0000 0000 756E 7772 6170 0000"))]
    #[case::seven("where's", &hex!("0700 0000 0000 0000 7768 6572 6527 7300"))]
    #[case::aligned("read_tea", &hex!("0800 0000 0000 0000 7265 6164 5F74 6561"))]
    #[case::more_bytes("read_tess", &hex!("0900 0000 0000 0000 7265 6164 5F74 6573 7300 0000 0000 0000"))]
    #[tokio::test]
    async fn test_write_display(
        #[case] value: &str,
        #[case] buf: &[u8],
        #[values(1, 2, 3, 4, 5, 6, 7, 8, 9, 1024)] chunks_size: usize,
    ) {
        let mut builder = Builder::new();
        for chunk in buf.chunks(chunks_size) {
            builder.write(chunk);
            builder.wait(Duration::ZERO);
        }
        let mock = builder.build();
        let mut writer = NixWriter::builder().build(mock);

        writer.write_display(value).await.unwrap();
        assert_eq!(writer.buffer(), buf);
        writer.flush().await.unwrap();
        assert_eq!(writer.buffer(), b"");
    }
}