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

5 files changed, 67 insertions, 42 deletions

diff --git a/azalea-nbt/benches/my_benchmark.rs b/azalea-nbt/benches/my_benchmark.rs
index ed963fb1..2e4efa97 100755
--- a/azalea-nbt/benches/my_benchmark.rs
+++ b/azalea-nbt/benches/my_benchmark.rs
@@ -1,25 +1,26 @@
 use azalea_nbt::Tag;
-use criterion::{criterion_group, criterion_main, Criterion, Throughput};
+use criterion::{black_box, criterion_group, criterion_main, Criterion, Throughput};
 use flate2::read::GzDecoder;
 use std::{
     fs::File,
-    io::{self, Read, Seek, SeekFrom},
+    io::{self, Cursor, Read},
 };
 
 fn bench_serialize(filename: &str, c: &mut Criterion) {
     let mut file = File::open(filename).unwrap();
     let mut contents = Vec::new();
     file.read_to_end(&mut contents).unwrap();
-    let mut src = std::io::Cursor::new(&contents[..]);
+    let mut src = &contents[..];
 
     // decode the original src so most of the time isn't spent on unzipping
     let mut decoded_src_decoder = GzDecoder::new(&mut src);
     let mut decoded_src = Vec::new();
     decoded_src_decoder.read_to_end(&mut decoded_src).unwrap();
-    let mut decoded_src_stream = std::io::Cursor::new(decoded_src.clone());
 
-    file.seek(SeekFrom::Start(0)).unwrap();
+    let mut decoded_src_stream = Cursor::new(&decoded_src[..]);
+
     let nbt = Tag::read(&mut decoded_src_stream).unwrap();
+    decoded_src_stream.set_position(0);
 
     let mut group = c.benchmark_group(filename);
 
@@ -27,9 +28,8 @@ fn bench_serialize(filename: &str, c: &mut Criterion) {
 
     group.bench_function("Decode", |b| {
         b.iter(|| {
-            let mut owned_decoded_src_stream = decoded_src_stream.clone();
-            owned_decoded_src_stream.seek(SeekFrom::Start(0)).unwrap();
-            Tag::read(&mut owned_decoded_src_stream).unwrap();
+            black_box(Tag::read(&mut decoded_src_stream).unwrap());
+            decoded_src_stream.set_position(0);
         })
     });
 
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)?)
     }
 }
diff --git a/azalea-nbt/src/error.rs b/azalea-nbt/src/error.rs
index ef4a9e9f..308c74c8 100755
--- a/azalea-nbt/src/error.rs
+++ b/azalea-nbt/src/error.rs
@@ -3,7 +3,8 @@ pub enum Error {
     InvalidTagType(u8),
     InvalidTag,
     WriteError(std::io::Error),
-    Utf8Error(std::string::FromUtf8Error),
+    Utf8Error(std::str::Utf8Error),
+    UnexpectedEof,
 }
 
 impl std::fmt::Display for Error {
@@ -13,6 +14,7 @@ impl std::fmt::Display for Error {
             Error::InvalidTag => write!(f, "Invalid tag"),
             Error::WriteError(e) => write!(f, "Write error: {}", e),
             Error::Utf8Error(e) => write!(f, "Utf8 error: {}", e),
+            Error::UnexpectedEof => write!(f, "Unexpected EOF"),
         }
     }
 }
@@ -22,8 +24,8 @@ impl From<std::io::Error> for Error {
         Error::WriteError(e)
     }
 }
-impl From<std::string::FromUtf8Error> for Error {
-    fn from(e: std::string::FromUtf8Error) -> Self {
+impl From<std::str::Utf8Error> for Error {
+    fn from(e: std::str::Utf8Error) -> Self {
         Error::Utf8Error(e)
     }
 }
diff --git a/azalea-nbt/src/lib.rs b/azalea-nbt/src/lib.rs
index c644cfdc..e20049e0 100755
--- a/azalea-nbt/src/lib.rs
+++ b/azalea-nbt/src/lib.rs
@@ -8,10 +8,11 @@ pub use tag::Tag;
 
 #[cfg(test)]
 mod tests {
+    use std::io::Cursor;
+
     use super::*;
     use ahash::AHashMap;
     use azalea_buf::{McBufReadable, McBufWritable};
-    use std::io::Cursor;
 
     #[test]
     fn mcbuf_nbt() {
@@ -25,7 +26,7 @@ mod tests {
         )]));
         tag.write_into(&mut buf).unwrap();
 
-        let mut buf = Cursor::new(buf);
+        let mut buf = Cursor::new(&buf[..]);
 
         let result = Tag::read_from(&mut buf).unwrap();
         assert_eq!(
diff --git a/azalea-nbt/tests/tests.rs b/azalea-nbt/tests/tests.rs
index 41a14d1b..43c31590 100755
--- a/azalea-nbt/tests/tests.rs
+++ b/azalea-nbt/tests/tests.rs
@@ -9,7 +9,9 @@ use std::{
 fn test_decode_hello_world() {
     // read hello_world.nbt
     let mut file = File::open("tests/hello_world.nbt").unwrap();
-    let tag = Tag::read(&mut file).unwrap();
+    let mut buf = vec![];
+    file.read_to_end(&mut buf).unwrap();
+    let tag = Tag::read(&mut Cursor::new(&buf[..])).unwrap();
     assert_eq!(
         tag,
         Tag::Compound(AHashMap::from_iter(vec![(
@@ -28,14 +30,14 @@ fn test_roundtrip_hello_world() {
     let mut original = Vec::new();
     file.read_to_end(&mut original).unwrap();
 
-    let mut original_stream = Cursor::new(original.clone());
+    let mut original_stream = Cursor::new(&original[..]);
     let tag = Tag::read(&mut original_stream).unwrap();
 
     // write hello_world.nbt
-    let mut result = Cursor::new(Vec::new());
+    let mut result = Vec::new();
     tag.write(&mut result).unwrap();
 
-    assert_eq!(result.into_inner(), original);
+    assert_eq!(result, original);
 }
 
 #[test]
@@ -45,13 +47,13 @@ fn test_bigtest() {
     let mut original = Vec::new();
     file.read_to_end(&mut original).unwrap();
 
-    let mut original_stream = Cursor::new(original.clone());
+    let mut original_stream = Cursor::new(original);
     let original_tag = Tag::read_gzip(&mut original_stream).unwrap();
 
     let mut result = Vec::new();
     original_tag.write(&mut result).unwrap();
 
-    let decoded_tag = Tag::read(&mut Cursor::new(result)).unwrap();
+    let decoded_tag = Tag::read(&mut Cursor::new(&result)).unwrap();
 
     assert_eq!(decoded_tag, original_tag);
 }
@@ -83,7 +85,7 @@ fn test_stringtest() {
     let mut original = Vec::new();
     file.read_to_end(&mut original).unwrap();
 
-    let mut original_stream = Cursor::new(original.clone());
+    let mut original_stream = Cursor::new(original);
     let original_tag = Tag::read_gzip(&mut original_stream).unwrap();
 
     assert_eq!(original_tag, correct_tag);
@@ -95,13 +97,13 @@ fn test_complex_player() {
     let mut original = Vec::new();
     file.read_to_end(&mut original).unwrap();
 
-    let mut original_stream = Cursor::new(original.clone());
+    let mut original_stream = Cursor::new(original);
     let original_tag = Tag::read_gzip(&mut original_stream).unwrap();
 
     let mut result = Vec::new();
     original_tag.write(&mut result).unwrap();
 
-    let decoded_tag = Tag::read(&mut Cursor::new(result)).unwrap();
+    let decoded_tag = Tag::read(&mut Cursor::new(&result)).unwrap();
 
     assert_eq!(decoded_tag, original_tag);
 }
@@ -112,13 +114,13 @@ fn test_simple_player() {
     let mut original = Vec::new();
     file.read_to_end(&mut original).unwrap();
 
-    let mut original_stream = Cursor::new(original.clone());
+    let mut original_stream = Cursor::new(original);
     let original_tag = Tag::read_gzip(&mut original_stream).unwrap();
 
     let mut result = Vec::new();
     original_tag.write(&mut result).unwrap();
 
-    let decoded_tag = Tag::read(&mut Cursor::new(result)).unwrap();
+    let decoded_tag = Tag::read(&mut Cursor::new(&result)).unwrap();
 
     assert_eq!(decoded_tag, original_tag);
 }