fix CustomModelData and WrittenBookContent datacomponents

1 files changed, 30 insertions, 1 deletions

diff --git a/azalea-buf/src/read.rs b/azalea-buf/src/read.rs
index d7a55716..71c96c34 100755
--- a/azalea-buf/src/read.rs
+++ b/azalea-buf/src/read.rs
@@ -21,6 +21,8 @@ pub enum BufReadError {
     CouldNotReadBytes,
     #[error("The received encoded string buffer length is longer than maximum allowed ({length} > {max_length})")]
     StringLengthTooLong { length: u32, max_length: u32 },
+    #[error("The received Vec length is longer than maximum allowed ({length} > {max_length})")]
+    VecLengthTooLong { length: u32, max_length: u32 },
     #[error("{source}")]
     Io {
         #[from]
@@ -183,7 +185,7 @@ impl AzaleaRead for UnsizedByteArray {
     }
 }
 
-impl<T: AzaleaRead + Send> AzaleaRead for Vec<T> {
+impl<T: AzaleaRead> AzaleaRead for Vec<T> {
     default fn azalea_read(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
         let length = u32::azalea_read_var(buf)? as usize;
         // we limit the capacity to not get exploited into allocating a bunch
@@ -194,6 +196,23 @@ impl<T: AzaleaRead + Send> AzaleaRead for Vec<T> {
         Ok(contents)
     }
 }
+impl<T: AzaleaRead> AzaleaReadLimited for Vec<T> {
+    fn azalea_read_limited(buf: &mut Cursor<&[u8]>, limit: usize) -> Result<Self, BufReadError> {
+        let length = u32::azalea_read_var(buf)? as usize;
+        if length > limit {
+            return Err(BufReadError::VecLengthTooLong {
+                length: length as u32,
+                max_length: limit as u32,
+            });
+        }
+
+        let mut contents = Vec::with_capacity(usize::min(length, 65536));
+        for _ in 0..length {
+            contents.push(T::azalea_read(buf)?);
+        }
+        Ok(contents)
+    }
+}
 
 impl<K: AzaleaRead + Send + Eq + Hash, V: AzaleaRead + Send> AzaleaRead for HashMap<K, V> {
     fn azalea_read(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
@@ -343,6 +362,16 @@ impl<T: AzaleaReadVar> AzaleaReadVar for Option<T> {
         })
     }
 }
+impl<T: AzaleaReadLimited> AzaleaReadLimited for Option<T> {
+    fn azalea_read_limited(buf: &mut Cursor<&[u8]>, limit: usize) -> Result<Self, BufReadError> {
+        let present = bool::azalea_read(buf)?;
+        Ok(if present {
+            Some(T::azalea_read_limited(buf, limit)?)
+        } else {
+            None
+        })
+    }
+}
 
 // [String; 4]
 impl<T: AzaleaRead, const N: usize> AzaleaRead for [T; N] {