From bc3aa9467ae1e2d0ea1727093af9b0af14965e69 Mon Sep 17 00:00:00 2001 From: mat <27899617+mat-1@users.noreply.github.com> Date: Fri, 7 Oct 2022 20:12:36 -0500 Subject: Replace impl Read with Cursor<&[u8]> (#26) * Start getting rid of Cursor * try to make the tests pass and fail * make the tests pass * remove unused uses * fix clippy warnings * fix potential OOM exploits * fix OOM in az-nbt * fix nbt benchmark * fix a test * start replacing it with Cursor> * wip * fix all the issues * fix all tests * fix nbt benchmark * fix warnings --- azalea-buf/src/definitions.rs | 4 +- azalea-buf/src/lib.rs | 73 +++++++++--------- azalea-buf/src/read.rs | 145 ++++++++++++++++++++---------------- azalea-buf/src/serializable_uuid.rs | 6 +- 4 files changed, 122 insertions(+), 106 deletions(-) (limited to 'azalea-buf/src') diff --git a/azalea-buf/src/definitions.rs b/azalea-buf/src/definitions.rs index ae7a7407..77309c46 100644 --- a/azalea-buf/src/definitions.rs +++ b/azalea-buf/src/definitions.rs @@ -2,10 +2,10 @@ use std::ops::Deref; /// A Vec that isn't prefixed by a VarInt with the size. #[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub struct UnsizedByteArray(Vec); +pub struct UnsizedByteArray(pub Vec); impl Deref for UnsizedByteArray { - type Target = Vec; + type Target = [u8]; fn deref(&self) -> &Self::Target { &self.0 diff --git a/azalea-buf/src/lib.rs b/azalea-buf/src/lib.rs index b8daa325..2c42f2ca 100644 --- a/azalea-buf/src/lib.rs +++ b/azalea-buf/src/lib.rs @@ -12,7 +12,7 @@ mod write; pub use azalea_buf_macros::*; pub use definitions::*; -pub use read::{read_varint_async, BufReadError, McBufReadable, McBufVarReadable}; +pub use read::{BufReadError, McBufReadable, McBufVarReadable}; pub use serializable_uuid::*; pub use write::{McBufVarWritable, McBufWritable}; @@ -23,7 +23,7 @@ const MAX_STRING_LENGTH: u16 = 32767; #[cfg(test)] mod tests { use super::*; - use std::{collections::HashMap, io::Cursor}; + use std::collections::HashMap; #[test] fn test_write_varint() { @@ -74,44 +74,44 @@ mod tests { #[test] fn test_read_varint() { - let mut buf = Cursor::new(vec![0]); - assert_eq!(i32::var_read_from(&mut buf).unwrap(), 0); + let buf = &mut &vec![0][..]; + assert_eq!(i32::var_read_from(buf).unwrap(), 0); - let mut buf = Cursor::new(vec![1]); - assert_eq!(i32::var_read_from(&mut buf).unwrap(), 1); + let buf = &mut &vec![1][..]; + assert_eq!(i32::var_read_from(buf).unwrap(), 1); - let mut buf = Cursor::new(vec![2]); - assert_eq!(i32::var_read_from(&mut buf).unwrap(), 2); + let buf = &mut &vec![2][..]; + assert_eq!(i32::var_read_from(buf).unwrap(), 2); - let mut buf = Cursor::new(vec![127]); - assert_eq!(i32::var_read_from(&mut buf).unwrap(), 127); + let buf = &mut &vec![127][..]; + assert_eq!(i32::var_read_from(buf).unwrap(), 127); - let mut buf = Cursor::new(vec![128, 1]); - assert_eq!(i32::var_read_from(&mut buf).unwrap(), 128); + let buf = &mut &vec![128, 1][..]; + assert_eq!(i32::var_read_from(buf).unwrap(), 128); - let mut buf = Cursor::new(vec![255, 1]); - assert_eq!(i32::var_read_from(&mut buf).unwrap(), 255); + let buf = &mut &vec![255, 1][..]; + assert_eq!(i32::var_read_from(buf).unwrap(), 255); - let mut buf = Cursor::new(vec![221, 199, 1]); - assert_eq!(i32::var_read_from(&mut buf).unwrap(), 25565); + let buf = &mut &vec![221, 199, 1][..]; + assert_eq!(i32::var_read_from(buf).unwrap(), 25565); - let mut buf = Cursor::new(vec![255, 255, 127]); - assert_eq!(i32::var_read_from(&mut buf).unwrap(), 2097151); + let buf = &mut &vec![255, 255, 127][..]; + assert_eq!(i32::var_read_from(buf).unwrap(), 2097151); - let mut buf = Cursor::new(vec![255, 255, 255, 255, 7]); - assert_eq!(i32::var_read_from(&mut buf).unwrap(), 2147483647); + let buf = &mut &vec![255, 255, 255, 255, 7][..]; + assert_eq!(i32::var_read_from(buf).unwrap(), 2147483647); - let mut buf = Cursor::new(vec![255, 255, 255, 255, 15]); - assert_eq!(i32::var_read_from(&mut buf).unwrap(), -1); + let buf = &mut &vec![255, 255, 255, 255, 15][..]; + assert_eq!(i32::var_read_from(buf).unwrap(), -1); - let mut buf = Cursor::new(vec![128, 128, 128, 128, 8]); - assert_eq!(i32::var_read_from(&mut buf).unwrap(), -2147483648); + let buf = &mut &vec![128, 128, 128, 128, 8][..]; + assert_eq!(i32::var_read_from(buf).unwrap(), -2147483648); } #[test] fn test_read_varint_longer() { - let mut buf = Cursor::new(vec![138, 56, 0, 135, 56, 123]); - assert_eq!(i32::var_read_from(&mut buf).unwrap(), 7178); + let buf = &mut &vec![138, 56, 0, 135, 56, 123][..]; + assert_eq!(i32::var_read_from(buf).unwrap(), 7178); } #[test] @@ -123,8 +123,8 @@ mod tests { dbg!(&buf); - let mut buf = Cursor::new(buf); - let result = Vec::::read_from(&mut buf).unwrap(); + let buf = &mut &buf[..]; + let result = Vec::::read_from(buf).unwrap(); assert_eq!(result, original_vec); } @@ -133,9 +133,9 @@ mod tests { let mut buf = Vec::new(); vec![1, 2, 3].var_write_into(&mut buf).unwrap(); - let mut buf = Cursor::new(buf); + let buf = &mut &buf[..]; - let result = Vec::::var_read_from(&mut buf).unwrap(); + let result = Vec::::var_read_from(buf).unwrap(); assert_eq!(result, vec![1, 2, 3]); } @@ -149,20 +149,19 @@ mod tests { let mut buf = Vec::new(); original_map.var_write_into(&mut buf).unwrap(); - let mut buf = Cursor::new(buf); + let buf = &mut &buf[..]; - let result = HashMap::::var_read_from(&mut buf).unwrap(); + let result = HashMap::::var_read_from(buf).unwrap(); assert_eq!(result, original_map); } #[test] fn test_long() { - let mut buf = Vec::new(); - 123456u64.write_into(&mut buf).unwrap(); - - let mut buf = Cursor::new(buf); + let buf: &mut Vec = &mut Vec::new(); + 123456u64.write_into(buf).unwrap(); - assert_eq!(u64::read_from(&mut buf).unwrap(), 123456); + let buf = &mut &buf[..]; + assert_eq!(u64::read_from(buf).unwrap(), 123456); } } diff --git a/azalea-buf/src/read.rs b/azalea-buf/src/read.rs index 78c858e4..29f351c6 100644 --- a/azalea-buf/src/read.rs +++ b/azalea-buf/src/read.rs @@ -1,8 +1,11 @@ use super::{UnsizedByteArray, MAX_STRING_LENGTH}; use byteorder::{ReadBytesExt, BE}; -use std::{collections::HashMap, hash::Hash, io::Read}; +use std::{ + collections::HashMap, + hash::Hash, + io::{Cursor, Read}, +}; use thiserror::Error; -use tokio::io::{AsyncRead, AsyncReadExt}; #[derive(Error, Debug)] pub enum BufReadError { @@ -26,6 +29,11 @@ pub enum BufReadError { UnexpectedEnumVariant { id: i32 }, #[error("Unexpected enum variant {id}")] UnexpectedStringEnumVariant { id: String }, + #[error("Tried to read {attempted_read} bytes but there were only {actual_read}")] + UnexpectedEof { + attempted_read: usize, + actual_read: usize, + }, #[error("{0}")] Custom(String), #[cfg(feature = "serde_json")] @@ -33,7 +41,20 @@ pub enum BufReadError { Deserialization(#[from] serde_json::Error), } -fn read_utf_with_len(buf: &mut impl Read, max_length: u32) -> Result { +fn read_bytes<'a>(buf: &'a mut Cursor<&[u8]>, length: usize) -> Result<&'a [u8], BufReadError> { + if length > buf.get_ref().len() { + return Err(BufReadError::UnexpectedEof { + attempted_read: length, + actual_read: buf.get_ref().len(), + }); + } + let initial_position = buf.position() as usize; + buf.set_position(buf.position() + length as u64); + let data = &buf.get_ref()[initial_position..initial_position + length]; + Ok(data) +} + +fn read_utf_with_len(buf: &mut Cursor<&[u8]>, max_length: u32) -> Result { let length = u32::var_read_from(buf)?; // i don't know why it's multiplied by 4 but it's like that in mojang's code so if length as u32 > max_length * 4 { @@ -43,12 +64,10 @@ fn read_utf_with_len(buf: &mut impl Read, max_length: u32) -> Result length as usize { return Err(BufReadError::StringLengthTooLong { length, max_length }); } @@ -58,37 +77,37 @@ fn read_utf_with_len(buf: &mut impl Read, max_length: u32) -> Result Result { - let mut buffer = [0]; - let mut ans = 0; - for i in 0..5 { - reader.read_exact(&mut buffer).await?; - ans |= ((buffer[0] & 0b0111_1111) as i32) << (7 * i); - if buffer[0] & 0b1000_0000 == 0 { - break; - } - } - Ok(ans) -} +// pub async fn read_varint_async( +// reader: &mut (dyn AsyncRead + Unpin + Send), +// ) -> Result { +// let mut buffer = [0]; +// let mut ans = 0; +// for i in 0..5 { +// reader.read_exact(&mut buffer).await?; +// ans |= ((buffer[0] & 0b0111_1111) as i32) << (7 * i); +// if buffer[0] & 0b1000_0000 == 0 { +// break; +// } +// } +// Ok(ans) +// } pub trait McBufReadable where Self: Sized, { - fn read_from(buf: &mut impl Read) -> Result; + fn read_from(buf: &mut Cursor<&[u8]>) -> Result; } pub trait McBufVarReadable where Self: Sized, { - fn var_read_from(buf: &mut impl Read) -> Result; + fn var_read_from(buf: &mut Cursor<&[u8]>) -> Result; } impl McBufReadable for i32 { - fn read_from(buf: &mut impl Read) -> Result { + fn read_from(buf: &mut Cursor<&[u8]>) -> Result { Ok(buf.read_i32::()?) } } @@ -96,7 +115,7 @@ impl McBufReadable for i32 { impl McBufVarReadable for i32 { // fast varints modified from https://github.com/luojia65/mc-varint/blob/master/src/lib.rs#L67 /// Read a single varint from the reader and return the value - fn var_read_from(buf: &mut impl Read) -> Result { + fn var_read_from(buf: &mut Cursor<&[u8]>) -> Result { let mut buffer = [0]; let mut ans = 0; for i in 0..5 { @@ -112,7 +131,7 @@ impl McBufVarReadable for i32 { impl McBufVarReadable for i64 { // fast varints modified from https://github.com/luojia65/mc-varint/blob/master/src/lib.rs#L54 - fn var_read_from(buf: &mut impl Read) -> Result { + fn var_read_from(buf: &mut Cursor<&[u8]>) -> Result { let mut buffer = [0]; let mut ans = 0; for i in 0..8 { @@ -127,25 +146,26 @@ impl McBufVarReadable for i64 { } } impl McBufVarReadable for u64 { - fn var_read_from(buf: &mut impl Read) -> Result { + fn var_read_from(buf: &mut Cursor<&[u8]>) -> Result { i64::var_read_from(buf).map(|i| i as u64) } } impl McBufReadable for UnsizedByteArray { - fn read_from(buf: &mut impl Read) -> Result { + fn read_from(buf: &mut Cursor<&[u8]>) -> Result { // read to end of the buffer - let mut bytes = vec![]; - buf.read_to_end(&mut bytes) - .map_err(|_| BufReadError::CouldNotReadBytes)?; - Ok(bytes.into()) + let data = buf.get_ref()[buf.position() as usize..].to_vec(); + buf.set_position((buf.position()) + data.len() as u64); + Ok(UnsizedByteArray(data)) } } impl McBufReadable for Vec { - default fn read_from(buf: &mut impl Read) -> Result { + default fn read_from(buf: &mut Cursor<&[u8]>) -> Result { let length = u32::var_read_from(buf)? as usize; - let mut contents = Vec::with_capacity(length); + // we don't set the capacity here so we can't get exploited into + // allocating a bunch + let mut contents = vec![]; for _ in 0..length { contents.push(T::read_from(buf)?); } @@ -154,9 +174,9 @@ impl McBufReadable for Vec { } impl McBufReadable for HashMap { - default fn read_from(buf: &mut impl Read) -> Result { + default fn read_from(buf: &mut Cursor<&[u8]>) -> Result { let length = i32::var_read_from(buf)? as usize; - let mut contents = HashMap::with_capacity(length); + let mut contents = HashMap::new(); for _ in 0..length { contents.insert(K::read_from(buf)?, V::read_from(buf)?); } @@ -167,9 +187,9 @@ impl McBufReadable impl McBufVarReadable for HashMap { - default fn var_read_from(buf: &mut impl Read) -> Result { + default fn var_read_from(buf: &mut Cursor<&[u8]>) -> Result { let length = i32::var_read_from(buf)? as usize; - let mut contents = HashMap::with_capacity(length); + let mut contents = HashMap::new(); for _ in 0..length { contents.insert(K::read_from(buf)?, V::var_read_from(buf)?); } @@ -178,55 +198,52 @@ impl McBufVarRe } impl McBufReadable for Vec { - fn read_from(buf: &mut impl Read) -> Result { + fn read_from(buf: &mut Cursor<&[u8]>) -> Result { let length = i32::var_read_from(buf)? as usize; - let mut buffer = vec![0; length]; - buf.read_exact(&mut buffer) - .map_err(|_| BufReadError::CouldNotReadBytes)?; - Ok(buffer) + read_bytes(buf, length).map(|b| b.to_vec()) } } impl McBufReadable for String { - fn read_from(buf: &mut impl Read) -> Result { + fn read_from(buf: &mut Cursor<&[u8]>) -> Result { read_utf_with_len(buf, MAX_STRING_LENGTH.into()) } } impl McBufReadable for u32 { - fn read_from(buf: &mut impl Read) -> Result { + fn read_from(buf: &mut Cursor<&[u8]>) -> Result { Ok(i32::read_from(buf)? as u32) } } impl McBufVarReadable for u32 { - fn var_read_from(buf: &mut impl Read) -> Result { + fn var_read_from(buf: &mut Cursor<&[u8]>) -> Result { Ok(i32::var_read_from(buf)? as u32) } } impl McBufReadable for u16 { - fn read_from(buf: &mut impl Read) -> Result { + fn read_from(buf: &mut Cursor<&[u8]>) -> Result { i16::read_from(buf).map(|i| i as u16) } } impl McBufReadable for i16 { - fn read_from(buf: &mut impl Read) -> Result { + fn read_from(buf: &mut Cursor<&[u8]>) -> Result { Ok(buf.read_i16::()?) } } impl McBufVarReadable for u16 { - fn var_read_from(buf: &mut impl Read) -> Result { + fn var_read_from(buf: &mut Cursor<&[u8]>) -> Result { Ok(i32::var_read_from(buf)? as u16) } } impl McBufVarReadable for Vec { - fn var_read_from(buf: &mut impl Read) -> Result { + fn var_read_from(buf: &mut Cursor<&[u8]>) -> Result { let length = i32::var_read_from(buf)? as usize; - let mut contents = Vec::with_capacity(length); + let mut contents = Vec::new(); for _ in 0..length { contents.push(T::var_read_from(buf)?); } @@ -235,49 +252,49 @@ impl McBufVarReadable for Vec { } impl McBufReadable for i64 { - fn read_from(buf: &mut impl Read) -> Result { + fn read_from(buf: &mut Cursor<&[u8]>) -> Result { Ok(buf.read_i64::()?) } } impl McBufReadable for u64 { - fn read_from(buf: &mut impl Read) -> Result { + fn read_from(buf: &mut Cursor<&[u8]>) -> Result { i64::read_from(buf).map(|i| i as u64) } } impl McBufReadable for bool { - fn read_from(buf: &mut impl Read) -> Result { + fn read_from(buf: &mut Cursor<&[u8]>) -> Result { Ok(u8::read_from(buf)? != 0) } } impl McBufReadable for u8 { - fn read_from(buf: &mut impl Read) -> Result { + fn read_from(buf: &mut Cursor<&[u8]>) -> Result { Ok(buf.read_u8()?) } } impl McBufReadable for i8 { - fn read_from(buf: &mut impl Read) -> Result { + fn read_from(buf: &mut Cursor<&[u8]>) -> Result { u8::read_from(buf).map(|i| i as i8) } } impl McBufReadable for f32 { - fn read_from(buf: &mut impl Read) -> Result { + fn read_from(buf: &mut Cursor<&[u8]>) -> Result { Ok(buf.read_f32::()?) } } impl McBufReadable for f64 { - fn read_from(buf: &mut impl Read) -> Result { + fn read_from(buf: &mut Cursor<&[u8]>) -> Result { Ok(buf.read_f64::()?) } } impl McBufReadable for Option { - default fn read_from(buf: &mut impl Read) -> Result { + default fn read_from(buf: &mut Cursor<&[u8]>) -> Result { let present = bool::read_from(buf)?; Ok(if present { Some(T::read_from(buf)?) @@ -288,7 +305,7 @@ impl McBufReadable for Option { } impl McBufVarReadable for Option { - default fn var_read_from(buf: &mut impl Read) -> Result { + default fn var_read_from(buf: &mut Cursor<&[u8]>) -> Result { let present = bool::read_from(buf)?; Ok(if present { Some(T::var_read_from(buf)?) @@ -300,13 +317,13 @@ impl McBufVarReadable for Option { // [String; 4] impl McBufReadable for [T; N] { - default fn read_from(buf: &mut impl Read) -> Result { + default fn read_from(buf: &mut Cursor<&[u8]>) -> Result { let mut contents = Vec::with_capacity(N); for _ in 0..N { contents.push(T::read_from(buf)?); } contents.try_into().map_err(|_| { - panic!("Panic is not possible since the Vec is the same size as the array") + unreachable!("Panic is not possible since the Vec is the same size as the array") }) } } diff --git a/azalea-buf/src/serializable_uuid.rs b/azalea-buf/src/serializable_uuid.rs index 66cdda73..dc251269 100644 --- a/azalea-buf/src/serializable_uuid.rs +++ b/azalea-buf/src/serializable_uuid.rs @@ -1,5 +1,5 @@ use crate::{read::BufReadError, McBufReadable, McBufWritable}; -use std::io::{Read, Write}; +use std::io::{Cursor, Write}; use uuid::Uuid; pub trait SerializableUuid { @@ -33,7 +33,7 @@ impl SerializableUuid for Uuid { } impl McBufReadable for Uuid { - fn read_from(buf: &mut impl Read) -> Result { + fn read_from(buf: &mut Cursor<&[u8]>) -> Result { Ok(Uuid::from_int_array([ u32::read_from(buf)?, u32::read_from(buf)?, @@ -80,7 +80,7 @@ mod tests { u.write_into(&mut buf).unwrap(); println!("{:?}", buf); assert_eq!(buf.len(), 16); - let u2 = Uuid::read_from(&mut buf.as_slice()).unwrap(); + let u2 = Uuid::read_from(&mut Cursor::new(&buf)).unwrap(); assert_eq!(u, u2); } } -- cgit v1.2.3