mdebray/tvl-depot
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
Show stats 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::{
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();
}