Fix memory leak and wrong nbt scan results

Fix a memory leak error caused by incorrect parsign of ByteArrays.
ByteArrays are now skipped completely when parsing nbt data since they
take up a lot of memory. The NBT Data is now parsed dynamically from
the BufReader of the file rather than using a preallocated slice.

Signed-off-by: trivernis <trivernis@protonmail.com>

Cargo.lock

@@ -307,7 +307,7 @@ dependencies = [
 
 [[package]]
 name = "minecraft-regions-tool"
-version = "0.4.0"
+version = "0.4.1"
 dependencies = [
  "byteorder",
  "colored",

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "minecraft-regions-tool"
-version = "0.4.0"
+version = "0.4.1"
 authors = ["trivernis <trivernis@protonmail.com>"]
 edition = "2018"
 license = "GPL-3.0"

src/chunk.rs

@@ -1,11 +1,10 @@
 use crate::nbt::{NBTError, NBTReader, NBTValue};
-use crate::region_file::BLOCK_SIZE;
-use byteorder::{BigEndian, ByteOrder, ReadBytesExt};
+use byteorder::{BigEndian, ReadBytesExt};
 
-use flate2::bufread::ZlibDecoder;
+use crate::region_file::BLOCK_SIZE;
+use flate2::read::ZlibDecoder;
 use std::fmt::{Display, Formatter};
-use std::fs::File;
-use std::io::{self, BufReader, Error, Read, Seek, SeekFrom};
+use std::io::{self, BufReader, Error};
 
 type IOResult<T> = io::Result<T>;
 
@@ -17,48 +16,33 @@ const TAG_Z_POS: &str = "zPos";
 pub struct Chunk {
     pub length: u32,
     pub compression_type: u8,
-    nbt_raw: Vec<u8>,
 }
 
 impl Chunk {
-    pub fn from_buf_reader(reader: &mut BufReader<File>, include_nbt: bool) -> IOResult<Self> {
-        let mut length_raw = [0u8; 4];
-        reader.read_exact(&mut length_raw)?;
-        let length = BigEndian::read_u32(&length_raw);
-        let compression_type = reader.read_u8()?;
-
-        let mut nbt_raw = Vec::new();
-        if include_nbt {
-            for _ in 0..((length - 1) as f32 / BLOCK_SIZE as f32).ceil() as u8 {
-                let mut buffer = [0u8; BLOCK_SIZE];
-                reader.read(&mut buffer)?;
-                nbt_raw.append(&mut buffer.to_vec());
-            }
-            nbt_raw.truncate((length - 1) as usize);
-        }
-
-        if length > 0 {
-            reader.seek(SeekFrom::Current((length - 1) as i64))?;
-        } else {
-            reader.seek(SeekFrom::Current((length) as i64))?;
+    pub fn from_buf_reader<R: io::Read + io::Seek>(reader: &mut R) -> IOResult<Self> {
+        let length = reader.read_u32::<BigEndian>()?;
+        if length > 128 * BLOCK_SIZE as u32 {
+            return Err(io::Error::from(io::ErrorKind::InvalidData));
         }
+        let compression_type = reader.read_u8()?;
 
         Ok(Self {
             compression_type,
             length,
-            nbt_raw,
         })
     }
 
-    pub fn validate_nbt_data(&mut self) -> Result<(), ChunkScanError> {
-        if self.compression_type == 2 {
-            let mut decoder = ZlibDecoder::new(&self.nbt_raw[..]);
-            let mut data = Vec::new();
-            decoder.read_to_end(&mut data)?;
-            self.nbt_raw = data;
-        }
-        let mut reader = NBTReader::new(&self.nbt_raw[..]);
-        let data = reader.parse()?;
+    pub fn validate_nbt_data<R: io::Read + io::Seek>(
+        &mut self,
+        reader: &mut R,
+    ) -> Result<(), ChunkScanError> {
+        let data = if self.compression_type == 2 {
+            let mut nbt_reader = NBTReader::new(BufReader::new(ZlibDecoder::new(reader)));
+            nbt_reader.parse()?
+        } else {
+            let mut nbt_reader = NBTReader::new(reader);
+            nbt_reader.parse()?
+        };
 
         if !data.contains_key(TAG_LEVEL) {
             Err(ChunkScanError::MissingTag(TAG_LEVEL))
@@ -87,6 +71,7 @@ pub enum ChunkScanError {
     NBTError(NBTError),
     MissingTag(&'static str),
     InvalidFormat(&'static str),
+    InvalidLength(u32),
 }
 
 impl Display for ChunkScanError {
@@ -97,6 +82,7 @@ impl Display for ChunkScanError {
             Self::NBTError(nbt) => write!(f, "NBT Error: {}", nbt),
             Self::MissingTag(tag) => write!(f, "Missing Tag in NBT Data: {}", tag),
             Self::InvalidFormat(tag) => write!(f, "Unexpected data format for NBT Tag {}", tag),
+            Self::InvalidLength(length) => write!(f, "Invalid chunk data length: {}", length),
         }
     }
 }

src/nbt.rs

@@ -4,8 +4,11 @@ use std::error::Error;
 use std::fmt::{self, Display, Formatter};
 use std::io::{self, Read};
 
+const MAX_RECURSION: u64 = 100;
+
 pub struct NBTReader<R> {
     inner: Box<R>,
+    recursion: u64,
 }
 
 type NBTResult<T> = Result<T, NBTError>;
@@ -17,6 +20,7 @@ where
     pub fn new(inner: R) -> Self {
         Self {
             inner: Box::new(inner),
+            recursion: 0,
         }
     }
 
@@ -27,13 +31,18 @@ where
         if tag != 10 {
             return Err(NBTError::MissingRootTag);
         }
-        let _ = self.parse_string()?;
+        let mut buf = [0u8; 2];
+        self.inner.read(&mut buf)?;
 
         self.parse_compound()
     }
 
     /// Parses a compound tag
     fn parse_compound(&mut self) -> NBTResult<HashMap<String, NBTValue>> {
+        self.recursion += 1;
+        if self.recursion > MAX_RECURSION {
+            return Err(NBTError::RecursionError);
+        }
         let mut root_value = HashMap::new();
         loop {
             let tag = self.inner.read_u8()?;
@@ -59,17 +68,18 @@ where
             };
             root_value.insert(name, value);
         }
-
+        self.recursion -= 1;
         Ok(root_value)
     }
 
     /// Parses an array of bytes
     fn parse_byte_array(&mut self) -> NBTResult<Vec<u8>> {
         let length = self.inner.read_u32::<BigEndian>()?;
-        let mut buf = Vec::with_capacity(length as usize);
-        self.inner.read_exact(&mut buf)?;
+        for _ in 0..length {
+            self.inner.read_u8()?;
+        }
 
-        Ok(buf)
+        Ok(Vec::with_capacity(0))
     }
 
     /// Parses a string value
@@ -78,10 +88,8 @@ where
         if length == 0 {
             return Ok(String::new());
         }
-        let mut buf = Vec::with_capacity(length as usize);
-        for _ in 0..length {
-            buf.push(self.inner.read_u8()?);
-        }
+        let mut buf = vec![0u8; length as usize];
+        self.inner.read_exact(&mut buf)?;
 
         String::from_utf8(buf).map_err(|_| NBTError::InvalidName)
     }
@@ -161,6 +169,7 @@ pub enum NBTError {
     MissingRootTag,
     InvalidTag(u8),
     InvalidName,
+    RecursionError,
 }
 
 impl Display for NBTError {
@@ -170,6 +179,7 @@ impl Display for NBTError {
             Self::InvalidTag(tag) => write!(f, "Invalid Tag: 0x{:x}", tag),
             Self::MissingRootTag => write!(f, "Missing root tag!"),
             Self::InvalidName => write!(f, "Encountered invalid tag name"),
+            Self::RecursionError => write!(f, "Reached recursion limit"),
         }
     }
 }

src/region_file.rs

@@ -52,7 +52,7 @@ impl RegionFile {
             let reader_offset = *offset as u64 * BLOCK_SIZE as u64;
             self.reader.seek(SeekFrom::Start(reader_offset))?;
 
-            match Chunk::from_buf_reader(&mut self.reader, true) {
+            match Chunk::from_buf_reader(&mut self.reader) {
                 Ok(mut chunk) => {
                     let chunk_sections = ((chunk.length + 4) as f64 / BLOCK_SIZE as f64).ceil();
 
@@ -63,7 +63,8 @@ impl RegionFile {
                             self.writer.write_u8(1)?;
                         }
                     } else {
-                        if let Err(e) = chunk.validate_nbt_data() {
+                        self.reader.seek(SeekFrom::Start(reader_offset + 5))?;
+                        if let Err(e) = chunk.validate_nbt_data(&mut self.reader) {
                             match e {
                                 ChunkScanError::IO(e) => {
                                     log::debug!(
@@ -73,8 +74,12 @@ impl RegionFile {
                                     );
                                     statistic.corrupted_compression += 1;
                                 }
+                                ChunkScanError::NBTError(e) => {
+                                    log::debug!("Corrupted nbt data for chunk {}: {}", offset, e);
+                                    statistic.corrupted_nbt += 1;
+                                }
                                 _ => {
-                                    log::debug!("Missing nbt for chunk {}: {}", offset, e);
+                                    log::debug!("Missing nbt data for chunk {}: {}", offset, e);
                                     statistic.missing_nbt += 1;
                                 }
                             }

src/scan.rs

@@ -7,6 +7,7 @@ pub struct ScanStatistics {
     pub invalid_length: u64,
     pub invalid_compression_method: u64,
     pub missing_nbt: u64,
+    pub corrupted_nbt: u64,
     pub failed_to_read: u64,
     pub corrupted_compression: u64,
 }
@@ -18,6 +19,7 @@ impl ScanStatistics {
             invalid_length: 0,
             invalid_compression_method: 0,
             missing_nbt: 0,
+            corrupted_nbt: 0,
             corrupted_compression: 0,
             failed_to_read: 0,
         }
@@ -34,6 +36,7 @@ impl Add for ScanStatistics {
         self.failed_to_read += rhs.failed_to_read;
         self.missing_nbt += rhs.missing_nbt;
         self.corrupted_compression += rhs.corrupted_compression;
+        self.corrupted_nbt += rhs.corrupted_nbt;
 
         self
     }
@@ -43,17 +46,20 @@ impl Display for ScanStatistics {
     fn fmt(&self, f: &mut Formatter<'_>) -> Result {
         write!(
             f,
-            "Total Chunks: {}
+            "
+            Total Chunks: {}
             Failed to Read: {}
             Chunks with invalid length: {}
             Chunks with invalid compression method: {}
             Chunks with missing nbt data: {}
+            Chunks with corrupted nbt data: {}
             Chunks with corrupted compressed data {}",
             self.total_chunks,
             self.failed_to_read,
             self.invalid_length,
             self.invalid_compression_method,
             self.missing_nbt,
+            self.corrupted_nbt,
             self.corrupted_compression
         )
     }