feat(nix-compat/wire/bytes/reader): parametrise on trailer tag r/8074

This allows using BytesReader with a custom tag, eg the closing parens
for the NAR reader.

No public constructor is provided for custom-tagged readers, since this
feature isn't public API.

Change-Id: I82e73d064edc4b6783ead1d6fe46a5b35f45c844
Reviewed-on: https://cl.tvl.fyi/c/depot/+/11543
Reviewed-by: Brian Olsen <me@griff.name>
Reviewed-by: flokli <flokli@flokli.de>
Tested-by: BuildkiteCI

3 files changed, 31 insertions, 10 deletions

diff --git a/tvix/nix-compat/src/wire/bytes/mod.rs b/tvix/nix-compat/src/wire/bytes/mod.rs
index 740a7ebfd03e..ef0b59def8b9 100644
--- a/tvix/nix-compat/src/wire/bytes/mod.rs
+++ b/tvix/nix-compat/src/wire/bytes/mod.rs
@@ -4,7 +4,7 @@ use std::{
 };
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 
-mod reader;
+pub(crate) mod reader;
 pub use reader::BytesReader;
 mod writer;
 pub use writer::BytesWriter;
diff --git a/tvix/nix-compat/src/wire/bytes/reader/mod.rs b/tvix/nix-compat/src/wire/bytes/reader/mod.rs
index 50398d9b9e40..cd45f78a0c84 100644
--- a/tvix/nix-compat/src/wire/bytes/reader/mod.rs
+++ b/tvix/nix-compat/src/wire/bytes/reader/mod.rs
@@ -8,6 +8,10 @@ use std::{
 use tokio::io::{AsyncRead, AsyncReadExt, ReadBuf};
 
 use trailer::{read_trailer, ReadTrailer, Trailer};
+
+#[doc(hidden)]
+pub use self::trailer::Pad;
+pub(crate) use self::trailer::Tag;
 mod trailer;
 
 /// Reads a "bytes wire packet" from the underlying reader.
@@ -24,12 +28,13 @@ mod trailer;
 /// If the data is not read all the way to the end, or an error is encountered,
 /// the underlying reader is no longer usable and might return garbage.
 #[derive(Debug)]
-pub struct BytesReader<R> {
-    state: State<R>,
+#[allow(private_bounds)]
+pub struct BytesReader<R, T: Tag = Pad> {
+    state: State<R, T>,
 }
 
 #[derive(Debug)]
-enum State<R> {
+enum State<R, T: Tag> {
     /// Full 8-byte blocks are being read and released to the caller.
     Body {
         reader: Option<R>,
@@ -38,7 +43,7 @@ enum State<R> {
         user_len: u64,
     },
     /// The trailer is in the process of being read.
-    ReadTrailer(ReadTrailer<R>),
+    ReadTrailer(ReadTrailer<R, T>),
     /// The trailer has been fully read and validated,
     /// and data can now be released to the caller.
     ReleaseTrailer { consumed: u8, data: Trailer },
@@ -49,7 +54,21 @@ where
     R: AsyncRead + Unpin,
 {
     /// Constructs a new BytesReader, using the underlying passed reader.
-    pub async fn new<S: RangeBounds<u64>>(mut reader: R, allowed_size: S) -> io::Result<Self> {
+    pub async fn new<S: RangeBounds<u64>>(reader: R, allowed_size: S) -> io::Result<Self> {
+        BytesReader::new_internal(reader, allowed_size).await
+    }
+}
+
+#[allow(private_bounds)]
+impl<R, T: Tag> BytesReader<R, T>
+where
+    R: AsyncRead + Unpin,
+{
+    /// Constructs a new BytesReader, using the underlying passed reader.
+    pub(crate) async fn new_internal<S: RangeBounds<u64>>(
+        mut reader: R,
+        allowed_size: S,
+    ) -> io::Result<Self> {
         let size = reader.read_u64_le().await?;
 
         if !allowed_size.contains(&size) {
@@ -84,7 +103,8 @@ where
     }
 }
 
-impl<R: AsyncRead + Unpin> AsyncRead for BytesReader<R> {
+#[allow(private_bounds)]
+impl<R: AsyncRead + Unpin, T: Tag> AsyncRead for BytesReader<R, T> {
     fn poll_read(
         mut self: Pin<&mut Self>,
         cx: &mut task::Context,
diff --git a/tvix/nix-compat/src/wire/bytes/reader/trailer.rs b/tvix/nix-compat/src/wire/bytes/reader/trailer.rs
index 858026bf7135..0b0c7b13554d 100644
--- a/tvix/nix-compat/src/wire/bytes/reader/trailer.rs
+++ b/tvix/nix-compat/src/wire/bytes/reader/trailer.rs
@@ -1,4 +1,5 @@
 use std::{
+    fmt::Debug,
     future::Future,
     marker::PhantomData,
     ops::Deref,
@@ -33,14 +34,14 @@ pub(crate) trait Tag {
     /// Suitably sized buffer for reading [Self::PATTERN]
     ///
     /// HACK: This is a workaround for const generics limitations.
-    type Buf: AsRef<[u8]> + AsMut<[u8]> + Unpin;
+    type Buf: AsRef<[u8]> + AsMut<[u8]> + Debug + Unpin;
 
     /// Make an instance of [Self::Buf]
     fn make_buf() -> Self::Buf;
 }
 
 #[derive(Debug)]
-pub(crate) enum Pad {}
+pub enum Pad {}
 
 impl Tag for Pad {
     const PATTERN: &'static [u8] = &[0; 8];
@@ -53,7 +54,7 @@ impl Tag for Pad {
 }
 
 #[derive(Debug)]
-pub(crate) struct ReadTrailer<R, T: Tag = Pad> {
+pub(crate) struct ReadTrailer<R, T: Tag> {
     reader: R,
     data_len: u8,
     filled: u8,