diff --git a/Cargo.lock b/Cargo.lock index 6602e16..f7e6e15 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -307,7 +307,7 @@ dependencies = [ [[package]] name = "minecraft-regions-tool" -version = "0.4.0" +version = "0.4.1" dependencies = [ "byteorder", "colored", diff --git a/Cargo.toml b/Cargo.toml index bbf52a8..0a514d2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "minecraft-regions-tool" -version = "0.4.0" +version = "0.4.1" authors = ["trivernis "] edition = "2018" license = "GPL-3.0" diff --git a/src/chunk.rs b/src/chunk.rs index b1931ea..6d6970a 100644 --- a/src/chunk.rs +++ b/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 = io::Result; @@ -17,48 +16,33 @@ const TAG_Z_POS: &str = "zPos"; pub struct Chunk { pub length: u32, pub compression_type: u8, - nbt_raw: Vec, } impl Chunk { - pub fn from_buf_reader(reader: &mut BufReader, include_nbt: bool) -> IOResult { - 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(reader: &mut R) -> IOResult { + let length = reader.read_u32::()?; + 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( + &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), } } } diff --git a/src/nbt.rs b/src/nbt.rs index 724f0af..4b16f9a 100644 --- a/src/nbt.rs +++ b/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 { inner: Box, + recursion: u64, } type NBTResult = Result; @@ -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> { + 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> { let length = self.inner.read_u32::()?; - 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"), } } } diff --git a/src/region_file.rs b/src/region_file.rs index f4ae776..18967ad 100644 --- a/src/region_file.rs +++ b/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; } } diff --git a/src/scan.rs b/src/scan.rs index 9fbd6d4..3a4f8be 100644 --- a/src/scan.rs +++ b/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 ) }