diff --git a/src/byte_buffer.rs b/src/byte_buffer.rs index 4bbc38b..b401ec1 100644 --- a/src/byte_buffer.rs +++ b/src/byte_buffer.rs @@ -1,5 +1,5 @@ use std::cmp::max; -use std::io::{Cursor, Read, Seek, SeekFrom, Write}; +use std::io::{Cursor, Seek, SeekFrom, Write}; use std::ops::{Deref, DerefMut}; use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt}; @@ -10,8 +10,8 @@ use node_types::KbinType; pub use encoding_type::EncodingType; pub use error::{KbinError, KbinErrorKind, Result}; -pub struct ByteBufferRead> { - buffer: Cursor, +pub struct ByteBufferRead<'buf> { + buffer: Cursor<&'buf [u8]>, offset_1: u64, offset_2: u64, } @@ -22,8 +22,8 @@ pub struct ByteBufferWrite { offset_2: u64, } -impl> ByteBufferRead { - pub fn new(buffer: R) -> Self { +impl<'buf> ByteBufferRead<'buf> { + pub fn new(buffer: &'buf [u8]) -> Self { Self { buffer: Cursor::new(buffer), offset_1: 0, @@ -41,21 +41,34 @@ impl> ByteBufferRead { self.buffer.position() } - pub fn buf_read(&mut self) -> Result> { + fn buf_read_size(&mut self, size: usize) -> Result<&'buf [u8]> { + // To avoid an allocation of a `Vec` here, the raw input byte array is used + let start = self.buffer.position() as usize; + let end = start + size as usize; + if end > self.buffer.get_ref().len() { + return Err(KbinErrorKind::DataRead(size as usize).into()); + } + + let data = &self.buffer.get_ref()[start..end]; + trace!("buf_read_size => index: {}, size: {}, data: 0x{:02x?}", self.buffer.position(), data.len(), data); + + self.buffer.seek(SeekFrom::Current(size as i64)).context(KbinErrorKind::DataRead(size as usize))?; + + Ok(data) + } + + pub fn buf_read(&mut self) -> Result<&'buf [u8]> { let size = self.buffer.read_u32::().context(KbinErrorKind::DataReadSize)?; debug!("buf_read => index: {}, size: {}", self.buffer.position(), size); - let mut data = vec![0; size as usize]; - self.buffer.read_exact(&mut data).context(KbinErrorKind::DataRead(size as usize))?; - trace!("buf_read => index: {}, size: {}, data: 0x{:02x?}", self.buffer.position(), data.len(), data); - + let data = self.buf_read_size(size as usize)?; self.realign_reads(None)?; Ok(data) } pub fn read_str(&mut self, encoding: EncodingType) -> Result { - let mut data = self.buf_read()?; + let data = self.buf_read()?; // Remove trailing null bytes let mut index = data.len() - 1; @@ -63,20 +76,20 @@ impl> ByteBufferRead { while index > 0 && index < len && data[index] == 0x00 { index -= 1; } - data.truncate(index + 1); + let data = &data[..=index]; trace!("read_str => size: {}, data: 0x{:02x?}", data.len(), data); encoding.decode_bytes(data) } - pub fn get(&mut self, size: u32) -> Result> { - let mut data = vec![0; size as usize]; - self.buffer.read_exact(&mut data).context(KbinErrorKind::DataRead(size as usize))?; + pub fn get(&mut self, size: u32) -> Result<&'buf [u8]> { + let data = self.buf_read_size(size as usize)?; + trace!("get => size: {}, data: 0x{:02x?}", size, data); Ok(data) } - pub fn get_aligned(&mut self, data_type: KbinType) -> Result> { + pub fn get_aligned(&mut self, data_type: KbinType) -> Result<&'buf [u8]> { if self.offset_1 % 4 == 0 { self.offset_1 = self.data_buf_offset(); } @@ -92,23 +105,21 @@ impl> ByteBufferRead { 1 => { self.buffer.seek(SeekFrom::Start(self.offset_1)).context(KbinErrorKind::Seek)?; - let data = self.buffer.read_u8().context(KbinErrorKind::DataReadOneByte)?; + let data = self.buf_read_size(1).context(KbinErrorKind::DataReadOneByte)?; self.offset_1 += 1; - (true, vec![data]) + (true, data) }, 2 => { self.buffer.seek(SeekFrom::Start(self.offset_2)).context(KbinErrorKind::Seek)?; - let mut data = vec![0; 2]; - self.buffer.read_exact(&mut data).context(KbinErrorKind::DataReadTwoByte)?; + let data = self.buf_read_size(2).context(KbinErrorKind::DataReadTwoByte)?; self.offset_2 += 2; (true, data) }, size => { - let mut data = vec![0; size as usize]; - self.buffer.read_exact(&mut data).context(KbinErrorKind::DataReadAligned)?; + let data = self.buf_read_size(size as usize).context(KbinErrorKind::DataReadAligned)?; self.realign_reads(None)?; (false, data) @@ -264,15 +275,15 @@ impl ByteBufferWrite { } } -impl Deref for ByteBufferRead where R: AsRef<[u8]> { - type Target = Cursor; +impl<'buf> Deref for ByteBufferRead<'buf> { + type Target = Cursor<&'buf [u8]>; fn deref(&self) -> &Self::Target { &self.buffer } } -impl DerefMut for ByteBufferRead where R: AsRef<[u8]> { +impl<'buf> DerefMut for ByteBufferRead<'buf> { fn deref_mut(&mut self) -> &mut Self::Target { &mut self.buffer } diff --git a/src/encoding_type.rs b/src/encoding_type.rs index d0ba9ee..b845afa 100644 --- a/src/encoding_type.rs +++ b/src/encoding_type.rs @@ -50,19 +50,19 @@ impl EncodingType { /// /// A `Some` value indicates an encoding should be used from the `encoding` /// crate. A `None` value indicates Rust's own UTF-8 handling should be used. - pub fn decode_bytes(&self, input: Vec) -> Result { + pub fn decode_bytes(&self, input: &[u8]) -> Result { let decoder_fail = |e| { format_err!("{}", e).context(KbinErrorKind::Encoding) }; let result = match *self { EncodingType::None | - EncodingType::UTF_8 => String::from_utf8(input)?, + EncodingType::UTF_8 => String::from_utf8(input.to_vec())?, - EncodingType::ASCII => ASCII.decode(&input, DecoderTrap::Strict).map_err(decoder_fail)?, - EncodingType::ISO_8859_1 => ISO_8859_1.decode(&input, DecoderTrap::Strict).map_err(decoder_fail)?, - EncodingType::EUC_JP => EUC_JP.decode(&input, DecoderTrap::Strict).map_err(decoder_fail)?, - EncodingType::SHIFT_JIS => WINDOWS_31J.decode(&input, DecoderTrap::Strict).map_err(decoder_fail)?, + EncodingType::ASCII => ASCII.decode(input, DecoderTrap::Strict).map_err(decoder_fail)?, + EncodingType::ISO_8859_1 => ISO_8859_1.decode(input, DecoderTrap::Strict).map_err(decoder_fail)?, + EncodingType::EUC_JP => EUC_JP.decode(input, DecoderTrap::Strict).map_err(decoder_fail)?, + EncodingType::SHIFT_JIS => WINDOWS_31J.decode(input, DecoderTrap::Strict).map_err(decoder_fail)?, }; Ok(result) diff --git a/src/lib.rs b/src/lib.rs index 3696643..53d1b62 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -17,7 +17,7 @@ extern crate serde; use std::fmt::Write as FmtWrite; use std::io::{Cursor, Write}; -use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt}; +use byteorder::{BigEndian, WriteBytesExt}; use failure::ResultExt; use minidom::Element; use rustc_hex::FromHex; diff --git a/src/reader.rs b/src/reader.rs index 9f5a826..5ad7b0d 100644 --- a/src/reader.rs +++ b/src/reader.rs @@ -12,8 +12,8 @@ use super::{ARRAY_MASK, SIGNATURE, SIG_COMPRESSED}; pub struct Reader<'buf> { encoding: EncodingType, - pub(crate) node_buf: ByteBufferRead<&'buf [u8]>, - pub(crate) data_buf: ByteBufferRead<&'buf [u8]>, + pub(crate) node_buf: ByteBufferRead<'buf>, + pub(crate) data_buf: ByteBufferRead<'buf>, data_buf_start: u64, } @@ -81,13 +81,12 @@ impl<'buf> Reader<'buf> { self.data_buf_start } - pub fn read_node_type(&mut self) -> Result<(StandardType, bool)> { - let raw_node_type = self.node_buf.read_u8().context(KbinErrorKind::NodeTypeRead)?; + fn parse_node_type(&self, raw_node_type: u8) -> Result<(StandardType, bool)> { let is_array = raw_node_type & ARRAY_MASK == ARRAY_MASK; let node_type = raw_node_type & !ARRAY_MASK; let xml_type = StandardType::from_u8(node_type); - debug!("Reader::read_node_type() => raw_node_type: {}, node_type: {:?} ({}), is_array: {}", + debug!("Reader::parse_node_type() => raw_node_type: {}, node_type: {:?} ({}), is_array: {}", raw_node_type, xml_type, node_type, @@ -96,6 +95,17 @@ impl<'buf> Reader<'buf> { Ok((xml_type, is_array)) } + pub fn peek_node_type(&self) -> Result<(StandardType, bool)> { + let pos = self.node_buf.position(); + let raw_node_type = self.node_buf.get_ref()[pos as usize]; + self.parse_node_type(raw_node_type) + } + + pub fn read_node_type(&mut self) -> Result<(StandardType, bool)> { + let raw_node_type = self.node_buf.read_u8().context(KbinErrorKind::NodeTypeRead)?; + self.parse_node_type(raw_node_type) + } + pub fn read_node_identifier(&mut self) -> Result { let value = unpack_sixbit(&mut *self.node_buf)?; debug!("Reader::read_node_identifier() => value: {:?}", value); @@ -124,9 +134,8 @@ impl<'buf> Reader<'buf> { Ok(value) } - // TODO: make a more intelligent reader to avoid allocating the Vec #[inline] - pub fn read_bytes(&mut self) -> Result> { + pub fn read_bytes(&mut self) -> Result<&'buf [u8]> { self.data_buf.buf_read() } }