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<Vec<u8>>

* wip

* fix all the issues

* fix all tests

* fix nbt benchmark

* fix warnings
---
 azalea-buf/src/read.rs | 145 +++++++++++++++++++++++++++----------------------
 1 file changed, 81 insertions(+), 64 deletions(-)

(limited to 'azalea-buf/src/read.rs')

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<String, BufReadError> {
+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<String, BufReadError> {
     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<String, Buf
         });
     }
 
-    // this is probably quite inefficient, idk how to do it better
-    let mut string = String::new();
-    let mut buffer = vec![0; length as usize];
-    buf.read_exact(&mut buffer)
-        .map_err(|_| BufReadError::InvalidUtf8)?;
-    string.push_str(std::str::from_utf8(&buffer).map_err(|_| BufReadError::InvalidUtf8)?);
+    let buffer = read_bytes(buf, length as usize)?;
+    let string = std::str::from_utf8(buffer)
+        .map_err(|_| BufReadError::InvalidUtf8)?
+        .to_string();
     if string.len() > 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<String, Buf
 
 // 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, along with the number of bytes read
-pub async fn read_varint_async(
-    reader: &mut (dyn AsyncRead + Unpin + Send),
-) -> Result<i32, BufReadError> {
-    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<i32, BufReadError> {
+//     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<Self, BufReadError>;
+    fn read_from(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError>;
 }
 
 pub trait McBufVarReadable
 where
     Self: Sized,
 {
-    fn var_read_from(buf: &mut impl Read) -> Result<Self, BufReadError>;
+    fn var_read_from(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError>;
 }
 
 impl McBufReadable for i32 {
-    fn read_from(buf: &mut impl Read) -> Result<Self, BufReadError> {
+    fn read_from(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
         Ok(buf.read_i32::<BE>()?)
     }
 }
@@ -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<Self, BufReadError> {
+    fn var_read_from(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
         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<Self, BufReadError> {
+    fn var_read_from(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
         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<Self, BufReadError> {
+    fn var_read_from(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
         i64::var_read_from(buf).map(|i| i as u64)
     }
 }
 
 impl McBufReadable for UnsizedByteArray {
-    fn read_from(buf: &mut impl Read) -> Result<Self, BufReadError> {
+    fn read_from(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
         // 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<T: McBufReadable + Send> McBufReadable for Vec<T> {
-    default fn read_from(buf: &mut impl Read) -> Result<Self, BufReadError> {
+    default fn read_from(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
         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<T: McBufReadable + Send> McBufReadable for Vec<T> {
 }
 
 impl<K: McBufReadable + Send + Eq + Hash, V: McBufReadable + Send> McBufReadable for HashMap<K, V> {
-    default fn read_from(buf: &mut impl Read) -> Result<Self, BufReadError> {
+    default fn read_from(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
         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<K: McBufReadable + Send + Eq + Hash, V: McBufReadable + Send> McBufReadable
 impl<K: McBufReadable + Send + Eq + Hash, V: McBufVarReadable + Send> McBufVarReadable
     for HashMap<K, V>
 {
-    default fn var_read_from(buf: &mut impl Read) -> Result<Self, BufReadError> {
+    default fn var_read_from(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
         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<K: McBufReadable + Send + Eq + Hash, V: McBufVarReadable + Send> McBufVarRe
 }
 
 impl McBufReadable for Vec<u8> {
-    fn read_from(buf: &mut impl Read) -> Result<Self, BufReadError> {
+    fn read_from(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
         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<Self, BufReadError> {
+    fn read_from(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
         read_utf_with_len(buf, MAX_STRING_LENGTH.into())
     }
 }
 
 impl McBufReadable for u32 {
-    fn read_from(buf: &mut impl Read) -> Result<Self, BufReadError> {
+    fn read_from(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
         Ok(i32::read_from(buf)? as u32)
     }
 }
 
 impl McBufVarReadable for u32 {
-    fn var_read_from(buf: &mut impl Read) -> Result<Self, BufReadError> {
+    fn var_read_from(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
         Ok(i32::var_read_from(buf)? as u32)
     }
 }
 
 impl McBufReadable for u16 {
-    fn read_from(buf: &mut impl Read) -> Result<Self, BufReadError> {
+    fn read_from(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
         i16::read_from(buf).map(|i| i as u16)
     }
 }
 
 impl McBufReadable for i16 {
-    fn read_from(buf: &mut impl Read) -> Result<Self, BufReadError> {
+    fn read_from(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
         Ok(buf.read_i16::<BE>()?)
     }
 }
 
 impl McBufVarReadable for u16 {
-    fn var_read_from(buf: &mut impl Read) -> Result<Self, BufReadError> {
+    fn var_read_from(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
         Ok(i32::var_read_from(buf)? as u16)
     }
 }
 
 impl<T: McBufVarReadable> McBufVarReadable for Vec<T> {
-    fn var_read_from(buf: &mut impl Read) -> Result<Self, BufReadError> {
+    fn var_read_from(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
         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<T: McBufVarReadable> McBufVarReadable for Vec<T> {
 }
 
 impl McBufReadable for i64 {
-    fn read_from(buf: &mut impl Read) -> Result<Self, BufReadError> {
+    fn read_from(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
         Ok(buf.read_i64::<BE>()?)
     }
 }
 
 impl McBufReadable for u64 {
-    fn read_from(buf: &mut impl Read) -> Result<Self, BufReadError> {
+    fn read_from(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
         i64::read_from(buf).map(|i| i as u64)
     }
 }
 
 impl McBufReadable for bool {
-    fn read_from(buf: &mut impl Read) -> Result<Self, BufReadError> {
+    fn read_from(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
         Ok(u8::read_from(buf)? != 0)
     }
 }
 
 impl McBufReadable for u8 {
-    fn read_from(buf: &mut impl Read) -> Result<Self, BufReadError> {
+    fn read_from(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
         Ok(buf.read_u8()?)
     }
 }
 
 impl McBufReadable for i8 {
-    fn read_from(buf: &mut impl Read) -> Result<Self, BufReadError> {
+    fn read_from(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
         u8::read_from(buf).map(|i| i as i8)
     }
 }
 
 impl McBufReadable for f32 {
-    fn read_from(buf: &mut impl Read) -> Result<Self, BufReadError> {
+    fn read_from(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
         Ok(buf.read_f32::<BE>()?)
     }
 }
 
 impl McBufReadable for f64 {
-    fn read_from(buf: &mut impl Read) -> Result<Self, BufReadError> {
+    fn read_from(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
         Ok(buf.read_f64::<BE>()?)
     }
 }
 
 impl<T: McBufReadable> McBufReadable for Option<T> {
-    default fn read_from(buf: &mut impl Read) -> Result<Self, BufReadError> {
+    default fn read_from(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
         let present = bool::read_from(buf)?;
         Ok(if present {
             Some(T::read_from(buf)?)
@@ -288,7 +305,7 @@ impl<T: McBufReadable> McBufReadable for Option<T> {
 }
 
 impl<T: McBufVarReadable> McBufVarReadable for Option<T> {
-    default fn var_read_from(buf: &mut impl Read) -> Result<Self, BufReadError> {
+    default fn var_read_from(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
         let present = bool::read_from(buf)?;
         Ok(if present {
             Some(T::var_read_from(buf)?)
@@ -300,13 +317,13 @@ impl<T: McBufVarReadable> McBufVarReadable for Option<T> {
 
 // [String; 4]
 impl<T: McBufReadable, const N: usize> McBufReadable for [T; N] {
-    default fn read_from(buf: &mut impl Read) -> Result<Self, BufReadError> {
+    default fn read_from(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
         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")
         })
     }
 }
-- 
cgit v1.2.3