diff --git a/src/lib/rainbowutils.rs b/src/lib/rainbowutils.rs index 5b370b4..fcb9e84 100644 --- a/src/lib/rainbowutils.rs +++ b/src/lib/rainbowutils.rs @@ -1,17 +1,21 @@ -use byteorder::{BigEndian, ByteOrder, ReadBytesExt}; -use rand::AsByteSliceMut; +use byteorder::{BigEndian, ByteOrder}; use std::collections::HashMap; +use std::convert::{TryFrom, TryInto}; use std::fs::File; use std::io::{BufReader, Read}; use std::io::{Error, ErrorKind}; pub const BDF_HDR: &[u8; 11] = b"BDF\x01RAINBOW"; pub const NULL_BYTES: &[u8; 4] = &[0u8; 4]; +pub const META_CHUNK_NAME: &str = "META"; +pub const HTBL_CHUNK_NAME: &str = "HTBL"; +pub const DTBL_CHUNK_NAME: &str = "DTBL"; pub struct BinaryDictionaryFile { name: String, reader: BufReader, metadata: Option, + lookup_table: Option, } #[derive(Debug, Clone)] @@ -30,6 +34,11 @@ pub struct MetaChunk { compression_method: Option, } +#[derive(Debug, Clone)] +pub struct HashLookupTable { + entries: HashMap, +} + #[derive(Debug, Clone)] pub struct HashEntry { id: u32, @@ -44,22 +53,48 @@ pub struct DataEntry { } impl BinaryDictionaryFile { - fn new(reader: BufReader) -> Self { + pub fn new(reader: BufReader) -> Self { Self { name: "".to_string(), metadata: None, + lookup_table: None, reader, } } - fn read_metadata(&mut self) -> Result { + pub fn read_metadata(&mut self) -> Result<&MetaChunk, Error> { if !self.validate_header() { - return Err(Error::new(ErrorKind::InvalidData, "Invalid BDF Header!")); + return Err(Error::new(ErrorKind::InvalidData, "invalid BDF Header")); + } + let meta_chunk: MetaChunk = self.next_chunk()?.try_into()?; + self.metadata = Some(meta_chunk); + + if let Some(chunk) = &self.metadata { + Ok(&chunk) + } else { + Err(Error::new( + ErrorKind::Other, + "Failed to read self assigned metadata.", + )) } - let meta_chunk = self.next_chunk().as_meta_chunk(); - self.metadata = Some(meta_chunk.clone()); + } + + pub fn read_lookup_table(&mut self) -> Result<&HashLookupTable, Error> { + match &self.metadata { + None => self.read_metadata()?, + Some(t) => t, + }; + let lookup_table: HashLookupTable = self.next_chunk()?.try_into()?; + self.lookup_table = Some(lookup_table); - Ok(meta_chunk) + if let Some(chunk) = &self.lookup_table { + Ok(&chunk) + } else { + Err(Error::new( + ErrorKind::Other, + "failed to read self assigned chunk", + )) + } } fn validate_header(&mut self) -> bool { @@ -69,52 +104,109 @@ impl BinaryDictionaryFile { header == BDF_HDR.as_ref() } - fn next_chunk(&mut self) -> GenericChunk { + /// Returns the next chunk if one is available. + pub fn next_chunk(&mut self) -> Result { let mut length_raw = [0u8; 4]; - let _ = self.reader.read(&mut length_raw); + let _ = self.reader.read_exact(&mut length_raw)?; let length = BigEndian::read_u32(&mut length_raw); let mut name_raw = [0u8; 4]; - let _ = self.reader.read(&mut name_raw); - let name = - String::from_utf8(name_raw.to_vec()).expect("Failed to parse chunk name to string!"); + let _ = self.reader.read_exact(&mut name_raw)?; + let name = String::from_utf8(name_raw.to_vec()).expect("Failed to parse name string."); let mut data = vec![0u8; length as usize]; - let _ = self.reader.read(&mut data); + let _ = self.reader.read_exact(&mut data)?; let mut crc_raw = [0u8; 4]; - let _ = self.reader.read(&mut crc_raw); + let _ = self.reader.read_exact(&mut crc_raw)?; let crc = BigEndian::read_u32(&mut crc_raw); - GenericChunk { + Ok(GenericChunk { length, name, data, crc, - } + }) } } -impl GenericChunk { - fn as_meta_chunk(&self) -> MetaChunk { - let mut chunk_count_raw = self.data[0..4].to_vec(); - let mut entries_per_chunk = self.data[4..8].to_vec(); - let mut total_number_of_entries = self.data[8..12].to_vec(); - let mut compression_method_raw = self.data[12..16].to_vec(); - let chunk_count = BigEndian::read_u32(&mut chunk_count_raw); - let entries_per_chunk = BigEndian::read_u32(&mut entries_per_chunk); - let entry_count = BigEndian::read_u32(&mut total_number_of_entries); +impl TryFrom for MetaChunk { + type Error = Error; + + fn try_from(chunk: GenericChunk) -> Result { + if &chunk.name != HTBL_CHUNK_NAME { + return Err(Error::new( + ErrorKind::InvalidData, + "chunk name doesn't match", + )); + } + if chunk.data.len() < 16 { + return Err(Error::new(ErrorKind::InvalidData, "invalid chunk data")); + } + let chunk_count_raw = &chunk.data[0..4]; + let entries_per_chunk = &chunk.data[4..8]; + let total_number_of_entries = &chunk.data[8..12]; + let compression_method_raw = chunk.data[12..16].to_vec(); + let chunk_count = BigEndian::read_u32(chunk_count_raw); + let entries_per_chunk = BigEndian::read_u32(entries_per_chunk); + let entry_count = BigEndian::read_u32(total_number_of_entries); let compression_method = if &compression_method_raw != NULL_BYTES { Some( - String::from_utf8(compression_method_raw.to_vec()) - .expect("Failed to parse compression method from meta string"), + String::from_utf8(compression_method_raw) + .expect("Failed to parse compression method name!"), ) } else { None }; - MetaChunk { + Ok(MetaChunk { chunk_count, entries_per_chunk, entry_count, compression_method, + }) + } +} + +impl HashLookupTable { + pub fn get_entry(&self, id: u32) -> Option<&HashEntry> { + self.entries.get(&id) + } +} + +impl TryFrom for HashLookupTable { + type Error = Error; + + fn try_from(chunk: GenericChunk) -> Result { + if &chunk.name != HTBL_CHUNK_NAME { + return Err(Error::new( + ErrorKind::InvalidData, + "chunk name doesn't match", + )); + } + let mut hash_entries: HashMap = HashMap::new(); + let mut position = 0; + while chunk.data.len() > (position + 12) { + let id_raw = &chunk.data[position..position + 4]; + position += 4; + let output_length_raw = &chunk.data[position..position + 4]; + position += 4; + let name_length_raw = &chunk.data[position..position + 4]; + position += 4; + let id = BigEndian::read_u32(id_raw); + let output_length = BigEndian::read_u32(output_length_raw); + let name_length = BigEndian::read_u32(name_length_raw); + let name_raw = &chunk.data[position..position + name_length as usize]; + let name = + String::from_utf8(name_raw.to_vec()).expect("Failed to parse hash function name!"); + hash_entries.insert( + id, + HashEntry { + id, + output_length, + name, + }, + ); } + Ok(HashLookupTable { + entries: hash_entries, + }) } }