diff options
Diffstat (limited to 'tvix/nix-compat/src/wire')
-rw-r--r-- | tvix/nix-compat/src/wire/bytes/reader/mod.rs | 14 | ||||
-rw-r--r-- | tvix/nix-compat/src/wire/bytes/reader/trailer.rs | 17 |
2 files changed, 31 insertions, 0 deletions
diff --git a/tvix/nix-compat/src/wire/bytes/reader/mod.rs b/tvix/nix-compat/src/wire/bytes/reader/mod.rs index 78615faf0f4c..9f7013a4e9b1 100644 --- a/tvix/nix-compat/src/wire/bytes/reader/mod.rs +++ b/tvix/nix-compat/src/wire/bytes/reader/mod.rs @@ -89,6 +89,20 @@ where }, } } + + /// Remaining data length, ie not including data already read. + /// + /// If the size has not been read yet, this is [None]. + #[allow(clippy::len_without_is_empty)] // if size is unknown, we can't answer that + pub fn len(&self) -> Option<u64> { + match self.state { + State::Size { .. } => None, + State::Body { + consumed, user_len, .. + } => Some(user_len - consumed), + State::Trailer(ref r) => Some(r.len() as u64), + } + } } impl<R: AsyncRead + Unpin> AsyncRead for BytesReader<R> { diff --git a/tvix/nix-compat/src/wire/bytes/reader/trailer.rs b/tvix/nix-compat/src/wire/bytes/reader/trailer.rs index 958cead42d8f..61a77678080a 100644 --- a/tvix/nix-compat/src/wire/bytes/reader/trailer.rs +++ b/tvix/nix-compat/src/wire/bytes/reader/trailer.rs @@ -81,6 +81,12 @@ pub(crate) fn read_trailer<R: AsyncRead + Unpin, T: Tag>( } } +impl<R, T: Tag> ReadTrailer<R, T> { + pub fn len(&self) -> u8 { + self.data_len + } +} + impl<R: AsyncRead + Unpin, T: Tag> Future for ReadTrailer<R, T> { type Output = io::Result<Trailer>; @@ -141,6 +147,17 @@ impl<R: AsyncRead + Unpin> TrailerReader<R> { pub fn new(reader: R, data_len: u8) -> Self { Self::Reading(read_trailer(reader, data_len)) } + + pub fn len(&self) -> u8 { + match self { + TrailerReader::Reading(fut) => fut.len(), + &TrailerReader::Releasing { + off, + data: Trailer { data_len, .. }, + } => data_len - off, + TrailerReader::Done => 0, + } + } } impl<R: AsyncRead + Unpin> AsyncRead for TrailerReader<R> { |