1 files changed, 308 insertions, 0 deletions
diff --git a/tvix/nix-compat/src/nix_daemon/ser/writer.rs b/tvix/nix-compat/src/nix_daemon/ser/writer.rs
new file mode 100644
index 000000000000..87e30580af34
--- /dev/null
+++ b/tvix/nix-compat/src/nix_daemon/ser/writer.rs
@@ -0,0 +1,308 @@
+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::nix_daemon::ProtocolVersion;
+use crate::wire::padding_len;
+use crate::wire::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::nix_daemon::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"");
+    }
+}