1 files changed, 141 insertions, 0 deletions
diff --git a/tvix/nix-compat/src/nar/reader/read.rs b/tvix/nix-compat/src/nar/reader/read.rs
new file mode 100644
index 000000000000..9938581f2a2e
--- /dev/null
+++ b/tvix/nix-compat/src/nar/reader/read.rs
@@ -0,0 +1,141 @@
+//! Helpers for reading [crate::nar::wire] format.
+
+use std::io::{
+    self,
+    ErrorKind::{Interrupted, InvalidData, UnexpectedEof},
+};
+
+use super::Reader;
+use crate::nar::wire::Tag;
+
+/// Consume a little-endian [prim@u64] from the reader.
+pub fn u64(reader: &mut Reader) -> io::Result<u64> {
+    let mut buf = [0; 8];
+    reader.read_exact(&mut buf)?;
+    Ok(u64::from_le_bytes(buf))
+}
+
+/// Consume a byte string from the reader into a provided buffer,
+/// returning the data bytes.
+pub fn bytes_buf<'a, const N: usize>(
+    reader: &mut Reader,
+    buf: &'a mut [u8; N],
+    max_len: usize,
+) -> io::Result<&'a [u8]> {
+    assert_eq!(N % 8, 0);
+    assert!(max_len <= N);
+
+    // read the length, and reject excessively large values
+    let len = self::u64(reader)?;
+    if len > max_len as u64 {
+        return Err(InvalidData.into());
+    }
+    // we know the length fits in a usize now
+    let len = len as usize;
+
+    // read the data and padding into a buffer
+    let buf_len = (len + 7) & !7;
+    reader.read_exact(&mut buf[..buf_len])?;
+
+    // verify that the padding is all zeroes
+    for &b in &buf[len..buf_len] {
+        if b != 0 {
+            return Err(InvalidData.into());
+        }
+    }
+
+    Ok(&buf[..len])
+}
+
+/// Consume a byte string of up to `max_len` bytes from the reader.
+pub fn bytes(reader: &mut Reader, max_len: usize) -> io::Result<Vec<u8>> {
+    assert!(max_len <= isize::MAX as usize);
+
+    // read the length, and reject excessively large values
+    let len = self::u64(reader)?;
+    if len > max_len as u64 {
+        return Err(InvalidData.into());
+    }
+    // we know the length fits in a usize now
+    let len = len as usize;
+
+    // read the data and padding into a buffer
+    let buf_len = (len + 7) & !7;
+    let mut buf = vec![0; buf_len];
+    reader.read_exact(&mut buf)?;
+
+    // verify that the padding is all zeroes
+    for b in buf.drain(len..) {
+        if b != 0 {
+            return Err(InvalidData.into());
+        }
+    }
+
+    Ok(buf)
+}
+
+/// Consume a known token from the reader.
+pub fn token<const N: usize>(reader: &mut Reader, token: &[u8; N]) -> io::Result<()> {
+    let mut buf = [0u8; N];
+
+    // This implements something similar to [Read::read_exact], but verifies that
+    // the input data matches the token while we read it. These two slices respectively
+    // represent the remaining token to be verified, and the remaining input buffer.
+    let mut token = &token[..];
+    let mut buf = &mut buf[..];
+
+    while !token.is_empty() {
+        match reader.read(buf) {
+            Ok(0) => {
+                return Err(UnexpectedEof.into());
+            }
+            Ok(n) => {
+                let (t, b);
+                (t, token) = token.split_at(n);
+                (b, buf) = buf.split_at_mut(n);
+
+                if t != b {
+                    return Err(InvalidData.into());
+                }
+            }
+            Err(e) => {
+                if e.kind() != Interrupted {
+                    return Err(e);
+                }
+            }
+        }
+    }
+
+    Ok(())
+}
+
+/// Consume a [Tag] from the reader.
+pub fn tag<T: Tag>(reader: &mut Reader) -> io::Result<T> {
+    let mut buf = T::make_buf();
+    let buf = buf.as_mut();
+
+    // first read the known minimum length…
+    reader.read_exact(&mut buf[..T::MIN])?;
+
+    // then decide which tag we're expecting
+    let tag = T::from_u8(buf[T::OFF]).ok_or(InvalidData)?;
+    let (head, tail) = tag.as_bytes().split_at(T::MIN);
+
+    // make sure what we've read so far is valid
+    if buf[..T::MIN] != *head {
+        return Err(InvalidData.into());
+    }
+
+    // …then read the rest, if any
+    if !tail.is_empty() {
+        let rest = tail.len();
+        reader.read_exact(&mut buf[..rest])?;
+
+        // and make sure it's what we expect
+        if buf[..rest] != *tail {
+            return Err(InvalidData.into());
+        }
+    }
+
+    Ok(tag)
+}