diff options
author | edef <edef@edef.eu> | 2024-04-29T14·56+0000 |
---|---|---|
committer | edef <edef@edef.eu> | 2024-04-29T15·33+0000 |
commit | 84b27760d0e3260434656e15e82362dcea39319f (patch) | |
tree | 78470cc81d743ae9296a82c20ba13aa4114717af /tvix/nix-compat/src/wire/bytes | |
parent | fdecf52a52026c47d393df3316c97d54109f58c4 (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 'tvix/nix-compat/src/wire/bytes')
-rw-r--r-- | tvix/nix-compat/src/wire/bytes/reader/mod.rs | 21 |
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 5f6081b40451..b46b0b53396b 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(); } |