From 44bd9543a6d23afd9b735bd8149341ea071019c2 Mon Sep 17 00:00:00 2001
From: edef <edef@edef.eu>
Date: Mon, 29 Apr 2024 13:57:20 +0000
Subject: feat(nix-compat/wire/bytes/reader): expose the remaining data length

The API is a bit odd here, because we don't have a distinct type for a
known-length reader.

Change-Id: I4a1dd07fbed0a400004dbe4aa2095c51898ad3bd
Reviewed-on: https://cl.tvl.fyi/c/depot/+/11538
Tested-by: BuildkiteCI
Reviewed-by: flokli <flokli@flokli.de>
Reviewed-by: Brian Olsen <me@griff.name>
---
 tvix/nix-compat/src/wire/bytes/reader/mod.rs     | 14 ++++++++++++++
 tvix/nix-compat/src/wire/bytes/reader/trailer.rs | 17 +++++++++++++++++
 2 files changed, 31 insertions(+)

(limited to 'tvix')

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> {
-- 
cgit 1.4.1