1 files changed, 38 insertions, 18 deletions

diff --git a/azalea-nbt/src/decode.rs b/azalea-nbt/src/decode.rs
index 2ab337fc..8a1dfab5 100755
--- a/azalea-nbt/src/decode.rs
+++ b/azalea-nbt/src/decode.rs
@@ -4,20 +4,31 @@ use ahash::AHashMap;
 use azalea_buf::{BufReadError, McBufReadable};
 use byteorder::{ReadBytesExt, BE};
 use flate2::read::{GzDecoder, ZlibDecoder};
+use std::io::Cursor;
 use std::io::{BufRead, Read};
 
 #[inline]
-fn read_string(stream: &mut impl Read) -> Result<String, Error> {
-    let length = stream.read_u16::<BE>()?;
+fn read_bytes<'a>(buf: &'a mut Cursor<&[u8]>, length: usize) -> Result<&'a [u8], Error> {
+    if length > buf.get_ref().len() {
+        return Err(Error::UnexpectedEof);
+    }
+    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)
+}
 
-    let mut buf = vec![0; length as usize];
-    stream.read_exact(&mut buf)?;
-    Ok(String::from_utf8(buf)?)
+#[inline]
+fn read_string(stream: &mut Cursor<&[u8]>) -> Result<String, Error> {
+    let length = stream.read_u16::<BE>()? as usize;
+
+    let buf = read_bytes(stream, length)?;
+    Ok(std::str::from_utf8(buf)?.to_string())
 }
 
 impl Tag {
     #[inline]
-    fn read_known(stream: &mut impl Read, id: u8) -> Result<Tag, Error> {
+    fn read_known(stream: &mut Cursor<&[u8]>, id: u8) -> Result<Tag, Error> {
         Ok(match id {
             // Signifies the end of a TAG_Compound. It is only ever used inside
             // a TAG_Compound, and is not named despite being in a TAG_Compound
@@ -39,9 +50,8 @@ impl Tag {
             // A length-prefixed array of signed bytes. The prefix is a signed
             // integer (thus 4 bytes)
             7 => {
-                let length = stream.read_u32::<BE>()?;
-                let mut bytes = vec![0; length as usize];
-                stream.read_exact(&mut bytes)?;
+                let length = stream.read_u32::<BE>()? as usize;
+                let bytes = read_bytes(stream, length)?.to_vec();
                 Tag::ByteArray(bytes)
             }
             // A length-prefixed modified UTF-8 string. The prefix is an
@@ -58,8 +68,8 @@ impl Tag {
             // parsers should accept any type if the length is <= 0).
             9 => {
                 let type_id = stream.read_u8()?;
-                let length = stream.read_i32::<BE>()?;
-                let mut list = Vec::with_capacity(length as usize);
+                let length = stream.read_u32::<BE>()?;
+                let mut list = Vec::new();
                 for _ in 0..length {
                     list.push(Tag::read_known(stream, type_id)?);
                 }
@@ -84,7 +94,10 @@ impl Tag {
             // signed integer (thus 4 bytes) and indicates the number of 4 byte
             // integers.
             11 => {
-                let length = stream.read_u32::<BE>()?;
+                let length = stream.read_u32::<BE>()? as usize;
+                if length * 4 > stream.get_ref().len() {
+                    return Err(Error::UnexpectedEof);
+                }
                 let mut ints = Vec::with_capacity(length as usize);
                 for _ in 0..length {
                     ints.push(stream.read_i32::<BE>()?);
@@ -94,7 +107,10 @@ impl Tag {
             // A length-prefixed array of signed longs. The prefix is a signed
             // integer (thus 4 bytes) and indicates the number of 8 byte longs.
             12 => {
-                let length = stream.read_u32::<BE>()?;
+                let length = stream.read_u32::<BE>()? as usize;
+                if length * 8 > stream.get_ref().len() {
+                    return Err(Error::UnexpectedEof);
+                }
                 let mut longs = Vec::with_capacity(length as usize);
                 for _ in 0..length {
                     longs.push(stream.read_i64::<BE>()?);
@@ -105,7 +121,7 @@ impl Tag {
         })
     }
 
-    pub fn read(stream: &mut impl Read) -> Result<Tag, Error> {
+    pub fn read(stream: &mut Cursor<&[u8]>) -> Result<Tag, Error> {
         // default to compound tag
 
         // the parent compound only ever has one item
@@ -123,17 +139,21 @@ impl Tag {
 
     pub fn read_zlib(stream: &mut impl BufRead) -> Result<Tag, Error> {
         let mut gz = ZlibDecoder::new(stream);
-        Tag::read(&mut gz)
+        let mut buf = Vec::new();
+        gz.read_to_end(&mut buf)?;
+        Tag::read(&mut Cursor::new(&buf))
     }
 
-    pub fn read_gzip(stream: &mut impl Read) -> Result<Tag, Error> {
+    pub fn read_gzip(stream: &mut Cursor<Vec<u8>>) -> Result<Tag, Error> {
         let mut gz = GzDecoder::new(stream);
-        Tag::read(&mut gz)
+        let mut buf = Vec::new();
+        gz.read_to_end(&mut buf)?;
+        Tag::read(&mut Cursor::new(&buf))
     }
 }
 
 impl McBufReadable for Tag {
-    fn read_from(buf: &mut impl Read) -> Result<Self, BufReadError> {
+    fn read_from(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
         Ok(Tag::read(buf)?)
     }
 }