From 343e176bec6251b53006331e752c285d50c7167f Mon Sep 17 00:00:00 2001 From: edef Date: Tue, 30 Apr 2024 08:52:14 +0000 Subject: feat(nix-compat/wire/bytes/reader): parametrise on trailer tag 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 Reviewed-by: flokli Tested-by: BuildkiteCI --- tvix/nix-compat/src/wire/bytes/mod.rs | 2 +- tvix/nix-compat/src/wire/bytes/reader/mod.rs | 32 +++++++++++++++++++----- tvix/nix-compat/src/wire/bytes/reader/trailer.rs | 7 +++--- 3 files changed, 31 insertions(+), 10 deletions(-) (limited to 'tvix') 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 { - state: State, +#[allow(private_bounds)] +pub struct BytesReader { + state: State, } #[derive(Debug)] -enum State { +enum State { /// Full 8-byte blocks are being read and released to the caller. Body { reader: Option, @@ -38,7 +43,7 @@ enum State { user_len: u64, }, /// The trailer is in the process of being read. - ReadTrailer(ReadTrailer), + ReadTrailer(ReadTrailer), /// 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>(mut reader: R, allowed_size: S) -> io::Result { + pub async fn new>(reader: R, allowed_size: S) -> io::Result { + BytesReader::new_internal(reader, allowed_size).await + } +} + +#[allow(private_bounds)] +impl BytesReader +where + R: AsyncRead + Unpin, +{ + /// Constructs a new BytesReader, using the underlying passed reader. + pub(crate) async fn new_internal>( + mut reader: R, + allowed_size: S, + ) -> io::Result { let size = reader.read_u64_le().await?; if !allowed_size.contains(&size) { @@ -84,7 +103,8 @@ where } } -impl AsyncRead for BytesReader { +#[allow(private_bounds)] +impl AsyncRead for BytesReader { 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 { +pub(crate) struct ReadTrailer { reader: R, data_len: u8, filled: u8, -- cgit 1.4.1