From 8810290bafe53dece202c6ab40fbc29c1e5796b9 Mon Sep 17 00:00:00 2001 From: Matt Bilker Date: Fri, 14 Sep 2018 18:08:05 +0000 Subject: [PATCH] writer: new Writeable trait to support writing from Element, NodeCollection, and Node --- src/error.rs | 6 +- src/lib.rs | 10 ++- src/value/mod.rs | 54 +++++++----- src/writer.rs | 211 +++++++++++++++++++++++++++++++++++++++++------ 4 files changed, 229 insertions(+), 52 deletions(-) diff --git a/src/error.rs b/src/error.rs index ea087fc..e7e0862 100644 --- a/src/error.rs +++ b/src/error.rs @@ -5,6 +5,7 @@ use std::string::FromUtf8Error; use failure::{Backtrace, Compat, Context, Fail}; use serde::{de, ser}; +use value::Value; use node_types::{KbinType, StandardType}; @@ -107,10 +108,13 @@ pub enum KbinErrorKind { #[fail(display = "Type mismatch, expected: {}, found: {}", _0, _1)] TypeMismatch(StandardType, StandardType), + #[fail(display = "Value mismatch, expected {}, but found {:?}", _0, _1)] + ValueTypeMismatch(StandardType, Value), + #[fail(display = "Invalid input for boolean: {}", _0)] InvalidBooleanInput(u8), - #[fail(display = "Invalid node type for operation: {:?}", _0)] + #[fail(display = "Invalid node type {:?} for operation", _0)] InvalidNodeType(StandardType), #[fail(display = "Invalid state")] diff --git a/src/lib.rs b/src/lib.rs index 2583f9b..9c3de11 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -48,7 +48,7 @@ pub use options::Options; pub use de::from_bytes; pub use ser::to_bytes; pub use value::Value; -pub use writer::Writer; +pub use writer::{Writer, Writeable}; const SIGNATURE: u8 = 0xA0; @@ -161,12 +161,16 @@ pub fn node_from_binary(input: Bytes) -> Result<(Node, EncodingType)> { Ok((node, encoding)) } -pub fn to_binary(input: &Element) -> Result> { +pub fn to_binary(input: &T) -> Result> + where T: Writeable +{ let mut writer = Writer::new(); writer.to_binary(input) } -pub fn to_binary_with_options(options: Options, input: &Element) -> Result> { +pub fn to_binary_with_options(options: Options, input: &T) -> Result> + where T: Writeable +{ let mut writer = Writer::with_options(options); writer.to_binary(input) } diff --git a/src/value/mod.rs b/src/value/mod.rs index 55568ef..80db0e3 100644 --- a/src/value/mod.rs +++ b/src/value/mod.rs @@ -309,7 +309,7 @@ macro_rules! tuple { Ok(value) } - fn to_bytes_inner(self, output: &mut Vec) -> Result<(), KbinError> { + fn to_bytes_inner(&self, output: &mut Vec) -> Result<(), KbinError> { debug!("Value::to_bytes_inner(self: {:?})", self); macro_rules! gen_error { @@ -319,20 +319,20 @@ macro_rules! tuple { } match self { - Value::S8(n) => output.push(n as u8), - Value::U8(n) => output.push(n), - Value::S16(n) => output.write_i16::(n).context(gen_error!(S16))?, - Value::U16(n) => output.write_u16::(n).context(gen_error!(U16))?, - Value::S32(n) => output.write_i32::(n).context(gen_error!(S32))?, - Value::U32(n) => output.write_u32::(n).context(gen_error!(U32))?, - Value::S64(n) => output.write_i64::(n).context(gen_error!(S64))?, - Value::U64(n) => output.write_u64::(n).context(gen_error!(U64))?, - Value::Binary(data) => output.extend_from_slice(&data), - Value::Time(n) => output.write_u32::(n).context(gen_error!(Time))?, + Value::S8(ref n) => output.push(*n as u8), + Value::U8(ref n) => output.push(*n), + Value::S16(ref n) => output.write_i16::(*n).context(gen_error!(S16))?, + Value::U16(ref n) => output.write_u16::(*n).context(gen_error!(U16))?, + Value::S32(ref n) => output.write_i32::(*n).context(gen_error!(S32))?, + Value::U32(ref n) => output.write_u32::(*n).context(gen_error!(U32))?, + Value::S64(ref n) => output.write_i64::(*n).context(gen_error!(S64))?, + Value::U64(ref n) => output.write_u64::(*n).context(gen_error!(U64))?, + Value::Binary(ref data) => output.extend_from_slice(data), + Value::Time(ref n) => output.write_u32::(*n).context(gen_error!(Time))?, Value::Ip4(addr) => output.extend_from_slice(&addr.octets()), - Value::Float(n) => output.write_f32::(n).context(gen_error!(Float))?, - Value::Double(n) => output.write_f64::(n).context(gen_error!(Double))?, - Value::Boolean(v) => output.push(if v { 0x01 } else { 0x00 }), + Value::Float(ref n) => output.write_f32::(*n).context(gen_error!(Float))?, + Value::Double(ref n) => output.write_f64::(*n).context(gen_error!(Double))?, + Value::Boolean(ref v) => output.push(if *v { 0x01 } else { 0x00 }), Value::Array(_, values) => { for value in values { value.to_bytes_inner(output)?; @@ -428,13 +428,18 @@ macro_rules! construct_types { ] } - pub fn to_bytes(self) -> Result, KbinError> { + pub fn to_bytes(&self) -> Result, KbinError> { let mut output = Vec::new(); self.to_bytes_inner(&mut output)?; Ok(output) } + #[inline] + pub fn to_bytes_into(&self, output: &mut Vec) -> Result<(), KbinError> { + self.to_bytes_inner(output) + } + pub fn standard_type(&self) -> StandardType { match *self { $( @@ -448,17 +453,24 @@ macro_rules! construct_types { } } - pub fn as_i8(&self) -> Option { + pub fn as_i8(&self) -> Result { match self { - Value::S8(ref n) => Some(*n), - _ => None, + Value::S8(ref n) => Ok(*n), + value => Err(KbinErrorKind::ValueTypeMismatch(StandardType::S8, value.clone()).into()), } } - pub fn as_u8(&self) -> Option { + pub fn as_u8(&self) -> Result { match self { - Value::U8(ref n) => Some(*n), - _ => None, + Value::U8(ref n) => Ok(*n), + value => Err(KbinErrorKind::ValueTypeMismatch(StandardType::U8, value.clone()).into()), + } + } + + pub fn as_attribute(self) -> Result { + match self { + Value::Attribute(s) => Ok(s), + value => Err(KbinErrorKind::ValueTypeMismatch(StandardType::Attribute, value).into()), } } } diff --git a/src/writer.rs b/src/writer.rs index 18d8fa2..8994667 100644 --- a/src/writer.rs +++ b/src/writer.rs @@ -4,35 +4,67 @@ use byteorder::{BigEndian, WriteBytesExt}; use failure::ResultExt; use minidom::Element; use rustc_hex::FromHex; -use byte_buffer::ByteBufferWrite; -use node_types::StandardType; -use sixbit::Sixbit; +use byte_buffer::ByteBufferWrite; use compression::Compression; use error::{KbinErrorKind, Result}; +use node::{Node, NodeCollection}; +use node_types::StandardType; use options::Options; +use sixbit::Sixbit; use value::Value; use super::{ARRAY_MASK, SIGNATURE}; -pub struct Writer { - options: Options, +fn write_value(options: &Options, data_buf: &mut ByteBufferWrite, node_type: StandardType, is_array: bool, value: &Value) -> Result<()> { + match value { + Value::Binary(data) => { + trace!("data: 0x{:02x?}", data); + + let size = (data.len() as u32) * (node_type.size as u32); + data_buf.write_u32::(size).context(KbinErrorKind::DataWrite("binary node size"))?; + data_buf.write_all(&data).context(KbinErrorKind::DataWrite("binary"))?; + data_buf.realign_writes(None)?; + }, + Value::String(text) => { + data_buf.write_str(options.encoding, &text)?; + }, + Value::Array(node_type, values) => { + if !is_array { + return Err(KbinErrorKind::InvalidState.into()); + } + + let count = values.len(); + let total_size = count * node_type.count * node_type.size; + + let mut data = Vec::with_capacity(total_size); + for value in values { + value.to_bytes_into(&mut data)?; + } + + data_buf.write_u32::(total_size as u32).context(KbinErrorKind::DataWrite("node size"))?; + data_buf.write_all(&data).context(KbinErrorKind::DataWrite(node_type.name))?; + data_buf.realign_writes(None)?; + }, + value => { + if is_array { + return Err(KbinErrorKind::InvalidState.into()); + } else { + let data = value.to_bytes()?; + data_buf.write_aligned(*node_type, &data)?; + } + }, + }; + + Ok(()) } -impl Writer { - pub fn new() -> Self { - Self { - options: Options::default(), - } - } +pub trait Writeable { + fn write_node(options: &Options, node_buf: &mut ByteBufferWrite, data_buf: &mut ByteBufferWrite, input: &T) -> Result<()>; +} - pub fn with_options(options: Options) -> Self { - Self { - options, - } - } - - fn write_node(&mut self, node_buf: &mut ByteBufferWrite, data_buf: &mut ByteBufferWrite, input: &Element) -> Result<()> { +impl Writeable for Element { + fn write_node(options: &Options, node_buf: &mut ByteBufferWrite, data_buf: &mut ByteBufferWrite, input: &Element) -> Result<()> { let text = input.text(); let node_type = match input.attr("__type") { Some(name) => StandardType::from_name(name), @@ -66,10 +98,10 @@ impl Writer { count); node_buf.write_u8(node_type.id | array_mask).context(KbinErrorKind::DataWrite(node_type.name))?; - match self.options.compression { + match options.compression { Compression::Compressed => Sixbit::pack(&mut **node_buf, input.name())?, Compression::Uncompressed => { - let data = self.options.encoding.encode_bytes(input.name())?; + let data = 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"))?; @@ -89,7 +121,7 @@ impl Writer { data_buf.realign_writes(None)?; }, StandardType::String => { - data_buf.write_str(self.options.encoding, &text)?; + data_buf.write_str(options.encoding, &text)?; }, _ => { @@ -117,14 +149,14 @@ impl Writer { trace!("write_node => attr: {}, value: {}", key, value); - data_buf.write_str(self.options.encoding, value)?; + data_buf.write_str(options.encoding, value)?; let node_type = StandardType::Attribute; node_buf.write_u8(node_type.id).context(KbinErrorKind::DataWrite(node_type.name))?; - match self.options.compression { + match options.compression { Compression::Compressed => Sixbit::pack(&mut **node_buf, key)?, Compression::Uncompressed => { - let data = self.options.encoding.encode_bytes(key)?; + let data = options.encoding.encode_bytes(key)?; let len = (data.len() - 1) as u8; node_buf.write_u8(len | ARRAY_MASK).context(KbinErrorKind::DataWrite("attribute name length"))?; node_buf.write_all(&data).context(KbinErrorKind::DataWrite("node name bytes"))?; @@ -133,7 +165,7 @@ impl Writer { } for child in input.children() { - self.write_node(node_buf, data_buf, child)?; + Self::write_node(options, node_buf, data_buf, child)?; } // Always has the array bit set @@ -141,8 +173,133 @@ impl Writer { Ok(()) } +} - pub fn to_binary(&mut self, input: &Element) -> Result> { +impl Writeable for NodeCollection { + fn write_node(options: &Options, node_buf: &mut ByteBufferWrite, data_buf: &mut ByteBufferWrite, input: &NodeCollection) -> Result<()> { + let (node_type, is_array) = input.base().node_type_tuple(); + let array_mask = if is_array { ARRAY_MASK } else { 0 }; + let name = input.base().key()?.ok_or(KbinErrorKind::InvalidState)?; + + node_buf.write_u8(node_type.id | array_mask).context(KbinErrorKind::DataWrite(node_type.name))?; + match options.compression { + Compression::Compressed => Sixbit::pack(&mut **node_buf, &name)?, + Compression::Uncompressed => { + let data = options.encoding.encode_bytes(&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"))?; + }, + }; + + if node_type != StandardType::NodeStart { + let value = input.base().value()?; + write_value(options, data_buf, node_type, is_array, &value)?; + } + + for attr in input.attributes() { + let key = attr.key()?.ok_or(KbinErrorKind::InvalidState)?; + let value = attr.value_bytes().ok_or(KbinErrorKind::InvalidState)?; + + data_buf.buf_write(value)?; + + node_buf.write_u8(StandardType::Attribute.id).context(KbinErrorKind::DataWrite(StandardType::Attribute.name))?; + match options.compression { + Compression::Compressed => Sixbit::pack(&mut **node_buf, &key)?, + Compression::Uncompressed => { + let data = options.encoding.encode_bytes(&key)?; + let len = (data.len() - 1) as u8; + node_buf.write_u8(len | ARRAY_MASK).context(KbinErrorKind::DataWrite("attribute name length"))?; + node_buf.write_all(&data).context(KbinErrorKind::DataWrite("node name bytes"))?; + }, + }; + } + + for child in input.children() { + Self::write_node(options, node_buf, data_buf, child)?; + } + + // Always has the array bit set + node_buf.write_u8(StandardType::NodeEnd.id | ARRAY_MASK).context(KbinErrorKind::DataWrite("node end"))?; + + Ok(()) + } +} + +impl Writeable for Node { + fn write_node(options: &Options, node_buf: &mut ByteBufferWrite, data_buf: &mut ByteBufferWrite, input: &Node) -> Result<()> { + let (node_type, is_array) = match input.value() { + Some(Value::Array(node_type, _)) => (*node_type, true), + Some(ref value) => (value.standard_type(), false), + None => (StandardType::NodeStart, false), + }; + let array_mask = if is_array { ARRAY_MASK } else { 0 }; + + node_buf.write_u8(node_type.id | array_mask).context(KbinErrorKind::DataWrite(node_type.name))?; + match options.compression { + Compression::Compressed => Sixbit::pack(&mut **node_buf, &input.key())?, + Compression::Uncompressed => { + let data = options.encoding.encode_bytes(&input.key())?; + 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"))?; + }, + }; + + if let Some(value) = input.value() { + write_value(options, data_buf, node_type, is_array, value)?; + } + + if let Some(attributes) = input.attributes() { + for (key, value) in attributes { + data_buf.write_str(options.encoding, value)?; + + node_buf.write_u8(StandardType::Attribute.id).context(KbinErrorKind::DataWrite(StandardType::Attribute.name))?; + match options.compression { + Compression::Compressed => Sixbit::pack(&mut **node_buf, &key)?, + Compression::Uncompressed => { + let data = options.encoding.encode_bytes(&key)?; + let len = (data.len() - 1) as u8; + node_buf.write_u8(len | ARRAY_MASK).context(KbinErrorKind::DataWrite("attribute name length"))?; + node_buf.write_all(&data).context(KbinErrorKind::DataWrite("node name bytes"))?; + }, + }; + } + } + + if let Some(children) = input.children() { + for child in children { + Self::write_node(options, node_buf, data_buf, child)?; + } + } + + // Always has the array bit set + node_buf.write_u8(StandardType::NodeEnd.id | ARRAY_MASK).context(KbinErrorKind::DataWrite("node end"))?; + + Ok(()) + } +} + +pub struct Writer { + options: Options, +} + +impl Writer { + pub fn new() -> Self { + Self { + options: Options::default(), + } + } + + pub fn with_options(options: Options) -> Self { + Self { + options, + } + } + + pub fn to_binary(&mut self, input: &T) -> Result> + where T: Writeable + { let mut header = Cursor::new(Vec::with_capacity(8)); header.write_u8(SIGNATURE).context(KbinErrorKind::HeaderWrite("signature"))?; @@ -156,7 +313,7 @@ impl Writer { let mut node_buf = ByteBufferWrite::new(Vec::new()); let mut data_buf = ByteBufferWrite::new(Vec::new()); - self.write_node(&mut node_buf, &mut data_buf, input)?; + T::write_node(&self.options, &mut node_buf, &mut data_buf, input)?; node_buf.write_u8(StandardType::FileEnd.id | ARRAY_MASK).context(KbinErrorKind::DataWrite("file end"))?; node_buf.realign_writes(None)?;