maurice/snix
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

refactor(tvix/nix-compat/wire/bytes): use RangeInclusive for limits

Browse source 
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>
This commit is contained in:
edef 2024-04-29 14:56:20 +00:00
parent fdecf52a52
commit 84b27760d0
1 changed files with 7 additions and 14 deletions
Download patch file Download diff file Expand all files Collapse all files

21
tvix/nix-compat/src/wire/bytes/reader/mod.rs
View file

@ -1,7 +1,7 @@
use std::{ use std::{
future::Future, future::Future,
io, io,
ops::{Bound, RangeBounds}, ops::{Bound, RangeBounds, RangeInclusive},
pin::Pin, pin::Pin,
task::{self, ready, Poll}, task::{self, ready, Poll},
}; };
@ -37,10 +37,7 @@ enum State<R> {
/// The data size is being read. /// The data size is being read.
Size { Size {
reader: Option<R>, reader: Option<R>,
/// Minimum length (inclusive) allowed_size: RangeInclusive<u64>,
user_len_min: u64,
/// Maximum length (inclusive)
user_len_max: u64,
filled: u8, filled: u8,
buf: [u8; 8], buf: [u8; 8],
}, },
@ -64,13 +61,11 @@ where
{ {
/// Constructs a new BytesReader, using the underlying passed reader. /// Constructs a new BytesReader, using the underlying passed reader.
pub fn new<S: RangeBounds<u64>>(reader: R, allowed_size: S) -> Self { 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::Included(&n) => n,
Bound::Excluded(&n) => n.saturating_add(1), Bound::Excluded(&n) => n.saturating_add(1),
Bound::Unbounded => 0, Bound::Unbounded => 0,
}; }..=match allowed_size.end_bound() {
let user_len_max = match allowed_size.end_bound() {
Bound::Included(&n) => n, Bound::Included(&n) => n,
Bound::Excluded(&n) => n.checked_sub(1).unwrap(), Bound::Excluded(&n) => n.checked_sub(1).unwrap(),
Bound::Unbounded => u64::MAX, Bound::Unbounded => u64::MAX,
@ -79,8 +74,7 @@ where
Self { Self {
state: State::Size { state: State::Size {
reader: Some(reader), reader: Some(reader),
user_len_min, allowed_size,
user_len_max,
filled: 0, filled: 0,
buf: [0; 8], buf: [0; 8],
}, },
@ -128,15 +122,14 @@ impl<R: AsyncRead + Unpin> AsyncRead for BytesReader<R> {
match this { match this {
State::Size { State::Size {
reader, reader,
user_len_min, allowed_size,
user_len_max,
filled: 8, filled: 8,
buf, buf,
} => { } => {
let reader = reader.take().unwrap(); let reader = reader.take().unwrap();
let data_len = u64::from_le_bytes(*buf); 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")) return Err(io::Error::new(io::ErrorKind::InvalidData, "invalid size"))
.into(); .into();
} }