From a6f3a3626dce59c0f5a61e5d36042898d58417d6 Mon Sep 17 00:00:00 2001 From: Matt Bilker Date: Thu, 6 Sep 2018 03:21:31 +0000 Subject: [PATCH] sixbit: read from borrowed byte array - lib: supports uncompressed node names --- src/compression.rs | 2 +- src/lib.rs | 17 ++++++++++--- src/node/definition.rs | 7 +----- src/reader.rs | 10 +++++--- src/sixbit.rs | 57 +++++++++++++++++++++++++++--------------- 5 files changed, 59 insertions(+), 34 deletions(-) diff --git a/src/compression.rs b/src/compression.rs index be5d319..65cd9f8 100644 --- a/src/compression.rs +++ b/src/compression.rs @@ -16,7 +16,7 @@ impl Compression { } } - pub fn _to_byte(&self) -> u8 { + pub fn to_byte(&self) -> u8 { match *self { Compression::Compressed => SIG_COMPRESSED, Compression::Uncompressed => SIG_UNCOMPRESSED, diff --git a/src/lib.rs b/src/lib.rs index f78cfee..2e45085 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -46,6 +46,7 @@ use reader::Reader; use sixbit::Sixbit; // Public exports +pub use compression::Compression; pub use encoding_type::EncodingType; pub use printer::Printer; pub use error::{KbinError, KbinErrorKind, Result}; @@ -195,9 +196,17 @@ impl KbinXml { array_mask, count); - // TODO: support uncompressed node_buf.write_u8(node_type.id | array_mask).context(KbinErrorKind::DataWrite(node_type.name))?; - Sixbit::pack(&mut **node_buf, input.name())?; + + match self.options.compression { + Compression::Compressed => Sixbit::pack(&mut **node_buf, input.name())?, + Compression::Uncompressed => { + let data = self.options.encoding.encode_bytes(input.name())?; + let len = (data.len() - 1) as u8; + node_buf.write_u8(len | ARRAY_MASK).context(KbinErrorKind::DataWrite("node name length"))?; + node_buf.write_all(&data).context(KbinErrorKind::DataWrite("node name bytes"))?; + }, + }; match node_type { StandardType::NodeStart => {}, @@ -258,7 +267,9 @@ impl KbinXml { fn to_binary_internal(&mut self, input: &Element) -> Result> { let mut header = Cursor::new(Vec::with_capacity(8)); header.write_u8(SIGNATURE).context(KbinErrorKind::HeaderWrite("signature"))?; - header.write_u8(SIG_COMPRESSED).context(KbinErrorKind::HeaderWrite("compression"))?; + + let compression = self.options.compression.to_byte(); + header.write_u8(compression).context(KbinErrorKind::HeaderWrite("compression"))?; let encoding = self.options.encoding.to_byte(); header.write_u8(encoding).context(KbinErrorKind::HeaderWrite("encoding"))?; diff --git a/src/node/definition.rs b/src/node/definition.rs index 5b5ff77..c147e2f 100644 --- a/src/node/definition.rs +++ b/src/node/definition.rs @@ -1,5 +1,3 @@ -use std::io::Cursor; - use byte_buffer::strip_trailing_null_bytes; use encoding_type::EncodingType; use error::{KbinError, KbinErrorKind}; @@ -42,8 +40,7 @@ impl<'buf> Key<'buf> { fn to_string(&self) -> Result { match self { Key::Compressed { ref size, ref data } => { - let mut data = Cursor::new(data); - Ok(Sixbit::unpack(&mut data, *size)?) + Ok(Sixbit::unpack(data, *size)?) }, Key::Uncompressed { encoding, ref data } => { Ok(encoding.decode_bytes(data)?) @@ -96,7 +93,6 @@ impl<'buf> NodeDefinition<'buf> { }, (node_type, NodeData::Some { ref value_data, .. }) => { let value = Value::from_standard_type(node_type, self.is_array, value_data)?; - debug!("value: {:?}", value); match value { Some(value) => Ok(value), None => Err(KbinErrorKind::InvalidNodeType(node_type).into()), @@ -122,7 +118,6 @@ impl<'buf> NodeDefinition<'buf> { (_, NodeData::Some { key, .. }) => { let key = key.to_string()?; let value = self.value()?; - debug!("value: {:?}", value); Ok(Node::with_value(key, value)) }, (node_type, NodeData::None) => { diff --git a/src/reader.rs b/src/reader.rs index 921e692..5e2eb0e 100644 --- a/src/reader.rs +++ b/src/reader.rs @@ -127,7 +127,8 @@ impl<'buf> Reader<'buf> { let value = match self.compression { Compression::Compressed => { let size = Sixbit::size(&mut *self.node_buf)?; - Sixbit::unpack(&mut *self.node_buf, size)? + let data = self.node_buf.get(size.real_len as u32)?; + Sixbit::unpack(data, size)? }, Compression::Uncompressed => { let length = (self.node_buf.read_u8().context(KbinErrorKind::DataRead(1))? & !ARRAY_MASK) + 1; @@ -156,7 +157,8 @@ impl<'buf> Reader<'buf> { let value = match self.compression { Compression::Compressed => { let size = Sixbit::size(&mut *self.node_buf)?; - Sixbit::unpack(&mut *self.node_buf, size)? + let data = self.node_buf.get(size.real_len as u32)?; + Sixbit::unpack(data, size)? }, Compression::Uncompressed => { let length = (self.node_buf.read_u8().context(KbinErrorKind::DataRead(1))? & !ARRAY_MASK) + 1; @@ -189,7 +191,7 @@ impl<'buf> Reader<'buf> { }, node_type => self.data_buf.get_aligned(*node_type)?, }; - debug!("Reader::read_node_data(node_type: {:?}, is_array: {}) => value: {:?}", node_type, is_array, value); + debug!("Reader::read_node_data(node_type: {:?}, is_array: {}) => value: 0x{:02x?}", node_type, is_array, value); Ok(value) } @@ -205,7 +207,7 @@ impl<'buf> Reader<'buf> { let key = match self.compression { Compression::Compressed => { let size = Sixbit::size(&mut *self.node_buf)?; - let data = self.node_buf.get(size.1 as u32)?; + let data = self.node_buf.get(size.real_len as u32)?; Key::Compressed { size, data } }, Compression::Uncompressed => { diff --git a/src/sixbit.rs b/src/sixbit.rs index eae3ab2..eb12db0 100644 --- a/src/sixbit.rs +++ b/src/sixbit.rs @@ -20,7 +20,11 @@ lazy_static! { }; } -pub type SixbitSize = (u8, usize); +#[derive(Clone, Copy, Debug)] +pub struct SixbitSize { + pub sixbit_len: u8, + pub real_len: usize, +} pub struct Sixbit; @@ -28,12 +32,12 @@ impl Sixbit { pub fn size(reader: &mut T) -> Result where T: Read { - let len = reader.read_u8().context(KbinErrorKind::SixbitLengthRead)?; - let real_len = (f32::from(len * 6) / 8f32).ceil(); + let sixbit_len = reader.read_u8().context(KbinErrorKind::SixbitLengthRead)?; + let real_len = (f32::from(sixbit_len * 6) / 8f32).ceil(); let real_len = (real_len as u32) as usize; - debug!("sixbit_len: {}, real_len: {}", len, real_len); + debug!("sixbit_len: {}, real_len: {}", sixbit_len, real_len); - Ok((len, real_len)) + Ok(SixbitSize { sixbit_len, real_len }) } pub fn pack(writer: &mut T, input: &str) -> Result<(), KbinError> @@ -44,6 +48,7 @@ impl Sixbit { .map(|ch| { *BYTE_MAP.get(&ch).expect("Character must be a valid sixbit character") }); + let len = input.len(); let real_len = (f64::from(len as u32 * 6) / 8f64).ceil() as usize; debug!("sixbit_len: {}, real_len: {}", len, real_len); @@ -60,18 +65,17 @@ impl Sixbit { } writer.write_u8(len as u8).context(KbinErrorKind::SixbitLengthWrite)?; - writer.write(&bytes).context(KbinErrorKind::SixbitWrite)?; + writer.write_all(&bytes).context(KbinErrorKind::SixbitWrite)?; Ok(()) } - pub fn unpack(reader: &mut T, size: SixbitSize) -> Result - where T: Read - { - let (sixbit_len, len) = size; + pub fn unpack(buf: &[u8], size: SixbitSize) -> Result { + let SixbitSize { sixbit_len, real_len } = size; - let mut buf = vec![0; len]; - reader.read_exact(&mut buf).context(KbinErrorKind::SixbitRead)?; + if buf.len() < real_len { + return Err(KbinErrorKind::SixbitRead.into()); + } let sixbit_len = sixbit_len as usize; let mut result = String::with_capacity(sixbit_len); @@ -100,23 +104,25 @@ mod tests { use super::Sixbit; + const TEST1_STR: &str = "hello"; + const TEST1_BYTES: &[u8] = &[5,182,172,113,208]; + #[test] fn test_pack() { let _ = pretty_env_logger::try_init(); let mut data: Cursor> = Cursor::new(Vec::new()); - Sixbit::pack(&mut data, "hello").expect("Failed to pack 'hello' as sixbit"); - assert_eq!(data.into_inner(), &[5,182,172,113,208]); + Sixbit::pack(&mut data, TEST1_STR).expect("Failed to pack sixbit"); + assert_eq!(data.into_inner(), TEST1_BYTES); } #[test] fn test_unpack() { let _ = pretty_env_logger::try_init(); - let mut data = Cursor::new(&[5,182,172,113,208]); - let size = Sixbit::size(&mut data).expect("Failed to get size of 'hello' sixbit string"); - let result = Sixbit::unpack(&mut data, size).expect("Failed to unpack 'hello' sixbit string"); - assert_eq!(result, "hello"); + let size = Sixbit::size(&mut Cursor::new(TEST1_BYTES)).expect("Failed to get size of sixbit string"); + let result = Sixbit::unpack(&TEST1_BYTES[1..], size).expect("Failed to unpack sixbit string"); + assert_eq!(result, TEST1_STR); } #[bench] @@ -126,10 +132,21 @@ mod tests { b.iter(|| { for _ in 0..100 { data.seek(SeekFrom::Start(0)).unwrap(); - black_box(Sixbit::pack(&mut data, "hello").unwrap()); + black_box(Sixbit::pack(&mut data, TEST1_STR).unwrap()); } }); - assert_eq!(data.into_inner(), &[5,182,172,113,208]); + assert_eq!(data.into_inner(), TEST1_BYTES); + } + + #[bench] + fn bench_unpack(b: &mut Bencher) { + b.iter(|| { + for _ in 0..100 { + let size = Sixbit::size(&mut Cursor::new(TEST1_BYTES)).expect("Failed to get size of sixbit string"); + let result = Sixbit::unpack(&TEST1_BYTES[1..], size).expect("Failed to unpack sixbit string"); + black_box(result); + } + }); } }