node_definition: working data reading and parsing

2026-09-09 19:15:27 -05:00 · 2018-09-04 21:01:57 -04:00
parent 9bf833751e
commit 5acd855cf6
6 changed files with 303 additions and 51 deletions
--- a/src/error.rs
+++ b/src/error.rs
@@ -6,7 +6,7 @@ use std::string::FromUtf8Error;
 use failure::{Backtrace, Compat, Context, Fail};
 use serde::{de, ser};

-use node_types::KbinType;
+use node_types::{KbinType, StandardType};

 pub type Result<T> = StdResult<T, KbinError>;

@@ -107,6 +107,12 @@ pub enum KbinErrorKind {
  #[fail(display = "Type mismatch, expected: {}, found: {}", _0, _1)]
  TypeMismatch(KbinType, KbinType),

+  #[fail(display = "Invalid input for boolean: {}", _0)]
+  InvalidBooleanInput(u8),
+
+  #[fail(display = "Invalid node type to create `Node` object from, node_type: {:?}", _0)]
+  InvalidNodeType(StandardType),
+
  #[fail(display = "Invalid state")]
  InvalidState,
 }
--- a/src/node_definition.rs
+++ b/src/node_definition.rs
@@ -1,72 +1,119 @@
 use std::io::Cursor;

-use error::KbinError;
+use byte_buffer::strip_trailing_null_bytes;
+use encoding_type::EncodingType;
+use error::{KbinError, KbinErrorKind};
 use node::Node;
 use node_types::StandardType;
 use sixbit::{Sixbit, SixbitSize};
+use value::Value;

 #[derive(Debug)]
 pub enum Key<'buf> {
-  Some {
+  Compressed {
    size: SixbitSize,
    data: &'buf [u8],
  },
+  Uncompressed {
+    encoding: EncodingType,
+    data: &'buf [u8],
+  },
+}
+
+#[derive(Debug)]
+pub enum NodeData<'buf> {
+  Some {
+    key: Key<'buf>,
+    value_data: &'buf [u8],
+  },
  None,
 }

 #[derive(Debug)]
 pub struct NodeDefinition<'buf> {
+  encoding: EncodingType,
  pub node_type: StandardType,
  pub is_array: bool,

-  pub key: Key<'buf>,
+  data: NodeData<'buf>,
+}

-  pub value_data: Option<&'buf [u8]>,
+impl<'buf> Key<'buf> {
+  fn to_string(&self) -> Result<String, KbinError> {
+    match self {
+      Key::Compressed { ref size, ref data } => {
+        let mut data = Cursor::new(data);
+        Ok(Sixbit::unpack(&mut data, *size)?)
+      },
+      Key::Uncompressed { encoding, ref data } => {
+        Ok(encoding.decode_bytes(data)?)
+      },
+    }
+  }
 }

 impl<'buf> NodeDefinition<'buf> {
-  /*
-  pub fn new(
-    node_type: (StandardType, bool),
-    value_data: Option<&'buf [u8]>,
-  ) -> Self {
-    Self::with_key(node_type, Key::None, value_data)
-  }
-  */
-
-  pub fn with_key(
-    node_type: (StandardType, bool),
-    key: Key<'buf>,
-    value_data: Option<&'buf [u8]>,
-  ) -> Self {
+  pub fn new(encoding: EncodingType, node_type: (StandardType, bool)) -> Self {
    let (node_type, is_array) = node_type;

    Self {
+      encoding,
      node_type,
      is_array,
-      key,
-      value_data,
+      data: NodeData::None,
+    }
+  }
+
+  pub fn with_data(encoding: EncodingType, node_type: (StandardType, bool), data: NodeData<'buf>) -> Self {
+    let (node_type, is_array) = node_type;
+
+    Self {
+      encoding,
+      node_type,
+      is_array,
+      data,
    }
  }

  pub fn key(&self) -> Result<Option<String>, KbinError> {
-    match self.key {
-      Key::Some { ref size, ref data } => {
-        let mut data = Cursor::new(data);
-        Ok(Some(Sixbit::unpack(&mut data, *size)?))
-      },
-      Key::None => Ok(None),
+    match self.data {
+      NodeData::Some { ref key, .. } => key.to_string().map(Some),
+      NodeData::None => Ok(None),
    }
  }

  pub fn into_node(self) -> Result<Node, KbinError> {
-    let key = self.key()?.unwrap_or_else(|| String::new());
-
-    if let Some(_value_data) = self.value_data {
-      //Ok(Node::with_value(key, value))
-      unimplemented!();
-    } else {
-      Ok(Node::new(key))
+    trace!("parsing definition: {:?}", self);
+    match (self.node_type, self.data) {
+      (StandardType::NodeStart, _) |
+      (StandardType::NodeEnd, _) |
+      (StandardType::FileEnd, _) => {
+        return Err(KbinErrorKind::InvalidNodeType(self.node_type).into());
+      },
+      (StandardType::Attribute, NodeData::Some { key, value_data }) => {
+        let key = key.to_string()?;
+        let data = strip_trailing_null_bytes(value_data);
+        let value = self.encoding.decode_bytes(data)?;
+        Ok(Node::with_value(key, Value::Attribute(value)))
+      },
+      (StandardType::String, NodeData::Some { key, value_data }) => {
+        let key = key.to_string()?;
+        let data = strip_trailing_null_bytes(value_data);
+        let value = self.encoding.decode_bytes(data)?;
+        Ok(Node::with_value(key, Value::String(value)))
+      },
+      (node_type, NodeData::Some { key, value_data }) => {
+        let key = key.to_string()?;
+        let value = Value::from_standard_type(node_type, self.is_array, value_data)?;
+        debug!("value: {:?}", value);
+        match value {
+          Some(value) => Ok(Node::with_value(key, value)),
+          None => Ok(Node::new(key)),
+        }
+      },
+      (node_type, NodeData::None) => {
+        Err(KbinErrorKind::InvalidNodeType(node_type).into())
+      },
    }
  }
 }
--- a/src/node_types.rs
+++ b/src/node_types.rs
@@ -153,7 +153,8 @@ macro_rules! construct_types {
    }

    $(
-      pub const $upcase: KbinType = KbinType {
+      #[allow(non_upper_case_globals)]
+      pub const $konst: KbinType = KbinType {
        id: $id,
        konst: stringify!($konst),
        name: $name,
@@ -217,7 +218,7 @@ macro_rules! construct_types {
      fn deref(&self) -> &KbinType {
        match *self {
          $(
-            StandardType::$konst => &$upcase,
+            StandardType::$konst => &$konst,
          )+
        }
      }
--- a/src/printer.rs
+++ b/src/printer.rs
@@ -8,10 +8,12 @@ impl Printer {
  pub fn run(input: &[u8]) -> Result<()> {
    let mut reader = Reader::new(input)?;
    let mut nodes = Vec::new();
+    let mut definitions = Vec::new();

    while let Ok(def) = reader.read_node_definition() {
      eprintln!("definition: {:?}", def);

+      let node_type = def.node_type;
      let key = match def.key() {
        Ok(v) => v,
        Err(e) => {
@@ -19,9 +21,10 @@ impl Printer {
          None
        },
      };
-      nodes.push((def.node_type, def.is_array, key));
+      nodes.push((node_type, def.is_array, key));
+      definitions.push(def);

-      if def.node_type == StandardType::FileEnd {
+      if node_type == StandardType::FileEnd {
        break;
      }
    }
@@ -41,6 +44,18 @@ impl Printer {
      };
    }

+    for def in definitions.into_iter() {
+      match def.node_type {
+        StandardType::NodeStart |
+        StandardType::NodeEnd |
+        StandardType::FileEnd => {},
+        _ => {
+          let node = def.into_node();
+          eprintln!("node: {:?}", node);
+        },
+      }
+    }
+
    Ok(())
  }
 }
--- a/src/reader.rs
+++ b/src/reader.rs
@@ -7,9 +7,10 @@ use byte_buffer::ByteBufferRead;
 use compression::Compression;
 use encoding_type::EncodingType;
 use error::{KbinErrorKind, Result};
-use node_definition::{Key, NodeDefinition};
+use node_definition::{Key, NodeData, NodeDefinition};
 use node_types::StandardType;
 use sixbit::Sixbit;
+
 use super::{ARRAY_MASK, SIGNATURE};

 pub struct Reader<'buf> {
@@ -116,8 +117,17 @@ impl<'buf> Reader<'buf> {
  pub fn peek_node_identifier(&mut self) -> Result<String> {
    let old_pos = self.node_buf.position();
    let _raw_node_type = self.node_buf.read_u8().context(KbinErrorKind::NodeTypeRead)?;
-    let size = Sixbit::size(&mut *self.node_buf)?;
-    let value = Sixbit::unpack(&mut *self.node_buf, size)?;
+    let value = match self.compression {
+      Compression::Compressed => {
+        let size = Sixbit::size(&mut *self.node_buf)?;
+        Sixbit::unpack(&mut *self.node_buf, size)?
+      },
+      Compression::Uncompressed => {
+        let length = (self.node_buf.read_u8().context(KbinErrorKind::DataRead(1))? & !ARRAY_MASK) + 1;
+        let bytes = self.node_buf.get(length as u32)?;
+        self.encoding.decode_bytes(bytes)?
+      },
+    };

    let size = self.node_buf.position() - old_pos;
    self.node_buf.seek(SeekFrom::Start(old_pos)).context(KbinErrorKind::DataRead(size as usize))?;
@@ -141,8 +151,8 @@ impl<'buf> Reader<'buf> {
      },
      Compression::Uncompressed => {
        let length = (self.node_buf.read_u8().context(KbinErrorKind::DataRead(1))? & !ARRAY_MASK) + 1;
-        let bytes = self.node_buf.get(length as u32)?;
-        self.encoding.decode_bytes(bytes)?
+        let data = self.node_buf.get(length as u32)?;
+        self.encoding.decode_bytes(data)?
      },
    };
    debug!("Reader::read_node_identifier() => value: {:?}", value);
@@ -152,18 +162,55 @@ impl<'buf> Reader<'buf> {
    Ok(value)
  }

+  pub fn read_node_data(&mut self, node_type: (StandardType, bool)) -> Result<&'buf [u8]> {
+    let (node_type, is_array) = node_type;
+    trace!("Reader::read_node_data(node_type: {:?}, is_array: {})", node_type, is_array);
+
+    let value = match node_type {
+      StandardType::Attribute |
+      StandardType::String => self.data_buf.buf_read()?,
+      StandardType::Binary => self.read_bytes()?,
+
+      _ if is_array => {
+        let arr_size = self.read_u32().context(KbinErrorKind::ArrayLengthRead)?;
+        let data = self.data_buf.get(arr_size)?;
+        self.data_buf.realign_reads(None)?;
+
+        data
+      },
+      node_type => self.data_buf.get_aligned(*node_type)?,
+    };
+    debug!("Reader::read_node_data(node_type: {:?}, is_array: {}) => value: {:?}", node_type, is_array, value);
+
+    Ok(value)
+  }
+
  pub fn read_node_definition(&mut self) -> Result<NodeDefinition<'buf>> {
    let node_type = self.read_node_type()?;
-    let key = match node_type.0 {
+    match node_type.0 {
      StandardType::NodeEnd |
-      StandardType::FileEnd => Key::None,
+      StandardType::FileEnd => {
+        Ok(NodeDefinition::new(self.encoding, node_type))
+      }
      _ => {
-        let size = Sixbit::size(&mut *self.node_buf)?;
-        let data = self.node_buf.get(size.1 as u32)?;
-        Key::Some { size, data }
+        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)?;
+            Key::Compressed { size, data }
+          },
+          Compression::Uncompressed => {
+            let encoding = self.encoding;
+            let length = (self.node_buf.read_u8().context(KbinErrorKind::DataRead(1))? & !ARRAY_MASK) + 1;
+            let data = self.node_buf.get(length as u32)?;
+            Key::Uncompressed { encoding, data }
+          },
+        };
+        let value_data = self.read_node_data(node_type)?;
+        let node_data = NodeData::Some { key, value_data };
+        Ok(NodeDefinition::with_data(self.encoding, node_type, node_data))
      },
-    };
-    Ok(NodeDefinition::with_key(node_type, key, None))
+    }
  }

  pub fn read_string(&mut self) -> Result<String> {
--- a/src/value/mod.rs
+++ b/src/value/mod.rs
@@ -1,16 +1,134 @@
 use std::fmt;
 use std::net::Ipv4Addr;

+use byteorder::{BigEndian, ByteOrder};
+use error::{KbinError, KbinErrorKind};
 use serde::de::{Deserialize, Deserializer, DeserializeSeed};
 use serde_bytes::ByteBuf;

 use node::Node;
 use node::de::NodeSeed;
-use node_types::StandardType;
+use node_types::{self, StandardType};

 mod de;
 mod ser;

+macro_rules! tuple {
+  (
+    byte: [
+      s8: [$($s8_konst:ident),*],
+      u8: [$($u8_konst:ident),*],
+      bool: [$($bool_konst:ident),*]
+    ],
+    multi: [
+      $($read_method:ident => [$($multi_konst:ident),*]),*
+    ]
+  ) => {
+    pub fn from_standard_type(node_type: StandardType, is_array: bool, input: &[u8]) -> Result<Option<Value>, KbinError> {
+      let node_size = node_type.size * node_type.count;
+
+      if is_array {
+        let mut values = Vec::new();
+
+        for chunk in input.chunks(node_size) {
+          trace!("chunk: {:?}", chunk);
+          match Value::from_standard_type(node_type, false, chunk)? {
+            Some(value) => values.push(value),
+            None => return Err(KbinErrorKind::InvalidState.into()),
+          }
+        }
+        debug!("values: {:?}", values);
+        return Ok(Some(Value::Array(node_type, values)));
+      }
+
+      match node_type {
+        StandardType::String |
+        StandardType::Binary => {},
+        _ => {
+          if input.len() != node_size {
+            return Err(KbinErrorKind::SizeMismatch(*node_type, node_size, input.len()).into());
+          }
+        },
+      };
+
+      let value = match node_type {
+        StandardType::NodeStart |
+        StandardType::NodeEnd |
+        StandardType::FileEnd => return Ok(None),
+        StandardType::S8 => Value::S8(input[0] as i8),
+        StandardType::U8 => Value::U8(input[0]),
+        StandardType::S16 => Value::S16(BigEndian::read_i16(input)),
+        StandardType::U16 => Value::U16(BigEndian::read_u16(input)),
+        StandardType::S32 => Value::S32(BigEndian::read_i32(input)),
+        StandardType::U32 => Value::U32(BigEndian::read_u32(input)),
+        StandardType::S64 => Value::S64(BigEndian::read_i64(input)),
+        StandardType::U64 => Value::U64(BigEndian::read_u64(input)),
+        StandardType::Attribute |
+        StandardType::String => unimplemented!(),
+        StandardType::Binary => Value::Binary(input.to_vec()),
+        StandardType::Time => Value::Time(BigEndian::read_u32(input)),
+        StandardType::Ip4 => {
+          let mut octets = [0; 4];
+          octets[0..4].copy_from_slice(&input[0..4]);
+          Value::Ip4(Ipv4Addr::from(octets))
+        },
+        StandardType::Float => Value::Float(BigEndian::read_f32(input)),
+        StandardType::Double => Value::Double(BigEndian::read_f64(input)),
+        StandardType::Boolean => Value::Boolean(match input[0] {
+          0x00 => false,
+          0x01 => true,
+          input => return Err(KbinErrorKind::InvalidBooleanInput(input).into()),
+        }),
+        $(
+          StandardType::$s8_konst => {
+            const COUNT: usize = node_types::$s8_konst.count;
+            let mut value = [0; COUNT];
+            for i in 0..COUNT {
+              value[i] = input[i] as i8;
+            }
+            Value::$s8_konst(value)
+          },
+        )*
+        $(
+          StandardType::$u8_konst => {
+            const COUNT: usize = node_types::$u8_konst.count;
+            let mut value = [0; COUNT];
+            value[0..COUNT].copy_from_slice(&input[0..COUNT]);
+            Value::$u8_konst(value)
+          },
+        )*
+        $(
+          StandardType::$bool_konst => {
+            const COUNT: usize = node_types::$bool_konst.count;
+            let mut value: [_; COUNT] = Default::default();
+            for i in 0..COUNT {
+              value[i] = match input[i] {
+                0x00 => false,
+                0x01 => true,
+                input => return Err(KbinErrorKind::InvalidBooleanInput(input).into()),
+              };
+            }
+            Value::$bool_konst(value)
+          },
+        )*
+        $(
+          $(
+            StandardType::$multi_konst => {
+              const COUNT: usize = node_types::$multi_konst.count;
+              const SIZE: usize = node_types::$multi_konst.size * COUNT;
+              let mut value: [_; COUNT] = Default::default();
+              BigEndian::$read_method(&input[0..SIZE], &mut value);
+              Value::$multi_konst(value)
+            },
+          )*
+        )*
+      };
+
+      Ok(Some(value))
+    }
+  };
+}
+
 macro_rules! construct_types {
  (
    $(
@@ -39,6 +157,24 @@ macro_rules! construct_types {
    )+

    impl Value {
+      tuple! {
+        byte: [
+          s8: [S8_2, S8_3, S8_4, Vs8],
+          u8: [U8_2, U8_3, U8_4, Vu8],
+          bool: [Boolean2, Boolean3, Boolean4, Vb]
+        ],
+        multi: [
+          read_i16_into => [S16_2, S16_3, S16_4, Vs16],
+          read_i32_into => [S32_2, S32_3, S32_4],
+          read_i64_into => [S64_2, S64_3, S64_4],
+          read_u16_into => [U16_2, U16_3, U16_4, Vu16],
+          read_u32_into => [U32_2, U32_3, U32_4],
+          read_u64_into => [U64_2, U64_3, U64_4],
+          read_f32_into_unchecked => [Float2, Float3, Float4],
+          read_f64_into_unchecked => [Double2, Double3, Double4]
+        ]
+      }
+
      pub fn standard_type(&self) -> StandardType {
        match *self {
          $(