about summary refs log tree commit diff
path: root/tvix/nix-compat/src/wire/bytes/reader/mod.rs
diff options
context:
space:
mode:
authoredef <edef@edef.eu>2024-04-29T14·56+0000
committeredef <edef@edef.eu>2024-04-29T15·33+0000
commit84b27760d0e3260434656e15e82362dcea39319f (patch)
tree78470cc81d743ae9296a82c20ba13aa4114717af /tvix/nix-compat/src/wire/bytes/reader/mod.rs
parentfdecf52a52026c47d393df3316c97d54109f58c4 (diff)
refactor(tvix/nix-compat/wire/bytes): use RangeInclusive for limits r/8031
The (min, max) pair is already a RangeInclusive in essence, so we might
as well represent it that way.

Change-Id: I2f67f3c47dc36b87e866ff5dc2e0cd28f01fbb04
Reviewed-on: https://cl.tvl.fyi/c/depot/+/11540
Tested-by: BuildkiteCI
Reviewed-by: flokli <flokli@flokli.de>
Diffstat (limited to '')
-rw-r--r--tvix/nix-compat/src/wire/bytes/reader/mod.rs21
1 files changed, 7 insertions, 14 deletions
diff --git a/tvix/nix-compat/src/wire/bytes/reader/mod.rs b/tvix/nix-compat/src/wire/bytes/reader/mod.rs
index 5f6081b404..b46b0b5339 100644
--- a/tvix/nix-compat/src/wire/bytes/reader/mod.rs
+++ b/tvix/nix-compat/src/wire/bytes/reader/mod.rs
@@ -1,7 +1,7 @@
 use std::{
     future::Future,
     io,
-    ops::{Bound, RangeBounds},
+    ops::{Bound, RangeBounds, RangeInclusive},
     pin::Pin,
     task::{self, ready, Poll},
 };
@@ -37,10 +37,7 @@ enum State<R> {
     /// The data size is being read.
     Size {
         reader: Option<R>,
-        /// Minimum length (inclusive)
-        user_len_min: u64,
-        /// Maximum length (inclusive)
-        user_len_max: u64,
+        allowed_size: RangeInclusive<u64>,
         filled: u8,
         buf: [u8; 8],
     },
@@ -64,13 +61,11 @@ where
 {
     /// Constructs a new BytesReader, using the underlying passed reader.
     pub fn new<S: RangeBounds<u64>>(reader: R, allowed_size: S) -> Self {
-        let user_len_min = match allowed_size.start_bound() {
+        let allowed_size = match allowed_size.start_bound() {
             Bound::Included(&n) => n,
             Bound::Excluded(&n) => n.saturating_add(1),
             Bound::Unbounded => 0,
-        };
-
-        let user_len_max = match allowed_size.end_bound() {
+        }..=match allowed_size.end_bound() {
             Bound::Included(&n) => n,
             Bound::Excluded(&n) => n.checked_sub(1).unwrap(),
             Bound::Unbounded => u64::MAX,
@@ -79,8 +74,7 @@ where
         Self {
             state: State::Size {
                 reader: Some(reader),
-                user_len_min,
-                user_len_max,
+                allowed_size,
                 filled: 0,
                 buf: [0; 8],
             },
@@ -128,15 +122,14 @@ impl<R: AsyncRead + Unpin> AsyncRead for BytesReader<R> {
             match this {
                 State::Size {
                     reader,
-                    user_len_min,
-                    user_len_max,
+                    allowed_size,
                     filled: 8,
                     buf,
                 } => {
                     let reader = reader.take().unwrap();
 
                     let data_len = u64::from_le_bytes(*buf);
-                    if data_len < *user_len_min || data_len > *user_len_max {
+                    if !allowed_size.contains(&data_len) {
                         return Err(io::Error::new(io::ErrorKind::InvalidData, "invalid size"))
                             .into();
                     }