Index deserialization and traversal. Tests.

src/store/local.rs

@@ -4,7 +4,8 @@ use std::{
     fs::File,
     hash::{Hash, Hasher},
     io::{BufReader, Read, Seek, SeekFrom, Write},
-    path::Path,
+    iter::TakeWhile,
+    path::{Iter, Path},
 };
 
 use aes_gcm::{aead::Aead, Aes256Gcm, KeyInit, Nonce};
@@ -12,7 +13,7 @@ use argon2::PasswordHasher;
 use password_hash::SaltString;
 use rand::{rngs::StdRng, RngCore, SeedableRng};
 
-use super::{EncryptionType, HashType, RecordType, Store, StoreRecord};
+use super::{EncryptionStream, EncryptionType, HashType, RecordType, Store, StoreRecord};
 
 pub struct LocalStore {
     path: String,
@@ -158,11 +159,11 @@ fn try_deserialize_superblock(
         .decrypt_block_from_file(header.enc_sup_block_size as usize, 0, file)
         .unwrap();
 
-    if data[0..enc.key.len()].ne(&enc.key) {
+    if data.data[0..enc.key.len()].ne(&enc.key) {
         return Err("Key mismatch.");
     }
 
-    let meta = deserialize_store_meta(&data[31..]);
+    let meta = deserialize_store_meta(&&data.data[32..]);
 
     Ok(meta)
 }
@@ -252,12 +253,15 @@ fn try_create_store(path: &String, passphrase: &Vec<u8>) -> bool {
 
     header.chksum_crc = header.get_checksum();
 
+    let mut index_block: Vec<u8> = vec![0; meta.data_block_size as usize];
+
     let mut buf = serialize_header(&header);
 
     buf.extend_from_slice(&salt_buf);
     buf.extend_from_slice(&iv);
     buf.append([b'\0'].repeat(16 - iv.len()).as_mut());
-    buf.extend_from_slice(&enc_buf);
+    buf.extend_from_slice(&enc_buf.data);
+    buf.extend_from_slice(&enc.encrypt_block(&index_block, 0x90).unwrap().data);
 
     match std::fs::File::create(path).unwrap().write(&buf) {
         Ok(x) => return x == buf.len(),
@@ -384,7 +388,7 @@ pub fn load(path: String, passphrase: Vec<u8>) -> Result<LocalStore, &'static st
     };
 
     let store = LocalStore {
-        path: path,
+        path,
         header,
         enc_ctx,
         meta,
@@ -485,7 +489,7 @@ impl LocalRecord {
                 key: key.to_owned(),
                 r#type: RecordType::from(r#type),
                 value: value.to_vec(),
-                meta: meta,
+                meta,
             },
             position: 0,
             size: val_offset + val_len + meta_len,
@@ -517,6 +521,7 @@ impl Store for LocalStore {
     }
 }
 
+#[derive(PartialEq, Eq, PartialOrd, Ord)]
 enum IndexNodeType {
     Internal,
     Leaf,
@@ -534,12 +539,62 @@ impl From<u8> for IndexNodeType {
     }
 }
 
+struct IndexNodeEntry {
+    key: u64,
+    pointer: u64,
+}
+
 struct IndexNode {
     node_type: IndexNodeType,
-    children: [u64],
+    children: Vec<IndexNodeEntry>,
+}
+
+impl IndexNode {
+    fn children(&self) -> &Vec<IndexNodeEntry> {
+        &self.children
+    }
+
+    fn len(&self) -> usize {
+        self.children.len()
+    }
 }
 
 impl LocalStore {
+    fn read_node(&self, data: &mut EncryptionStream) -> IndexNode {
+        let node_type = {
+            let mut buf: [u8; 2] = [0; 2];
+            data.read_exact(&mut buf).expect("Cannot read node type.");
+            IndexNodeType::from(buf[0])
+        };
+
+        let mut children = Vec::with_capacity(self.meta.index_node_arity as usize);
+        for _ in 0..self.meta.index_node_arity {
+            let mut buf: [u8; 16] = [0; 16];
+
+            data.read_exact(&mut buf)
+                .expect("Cannot read child node content.");
+
+            let key = u64::from_be_bytes(buf[0..8].try_into().unwrap());
+
+            if key == 0 {
+                break;
+            }
+
+            let pointer = u64::from_be_bytes(buf[0..8].try_into().unwrap());
+
+            let entry = IndexNodeEntry { key, pointer };
+
+            children.push(entry);
+        }
+
+        let node = IndexNode {
+            node_type,
+            children,
+        };
+
+        node
+    }
+
     fn btree_find_key(&self, key: &String) -> Option<u64> {
         let mut file: BufReader<File> =
             BufReader::new(File::open(&self.path).expect("Cannot open file"));
@@ -551,34 +606,37 @@ impl LocalStore {
         };
 
         let pos = file.seek(SeekFrom::Start(0x90)).unwrap();
+        let mut block = self
+            .enc_ctx
+            .decrypt_block_from_file(self.meta.data_block_size as usize, 0x90, &mut file)
+            .unwrap();
 
-        let mut read_node_type = || {
-            let mut buf: [u8; 2] = [0; 2];
-            file.read_exact(&mut buf).expect("Cannot read node type.");
-            IndexNodeType::from(buf[0])
-        };
-
-        let mut read_key_pos = || {
-            let mut buf: [u8; 16] = [0; 16];
-            file.read_exact(&mut buf).expect("Read error");
-
-            (
-                u64::from_be_bytes(buf[0..8].try_into().unwrap()),
-                u64::from_be_bytes(buf[8..16].try_into().unwrap()),
-            )
-        };
-
-        let mut subnode_pos = 0;
         loop {
-            let (key, offset) = read_key_pos();
+            let node = self.read_node(&mut block);
 
-            if hash > key {
+            if node.node_type == IndexNodeType::Invalid {
-                if subnode_pos > self.meta.index_node_arity {
+                panic!(
+                    "Invalid node detected: {}",
+                    block.stream_position().unwrap()
+                );
+            }
+
+            if node.node_type == IndexNodeType::Leaf {
+                for child in node.children() {
+                    if child.key == hash {
+                        return Some(child.pointer);
+                    }
+                }
                 return None;
             }
-                subnode_pos += 1;
+
+            for child in node.children {
+                if child.key < hash {
                     continue;
                 }
+
+                block.seek(SeekFrom::Start(child.pointer)).unwrap();
+            }
         }
     }
 }
@@ -588,19 +646,24 @@ fn nonce_offset(nonce_root: &Vec<u8>, offset: usize) -> Vec<u8> {
     if len < (usize::BITS / 8) as usize {
         panic!("Nonce length less than {}", usize::BITS / 8);
     }
+    let start = len - (usize::BITS / 8) as usize;
 
     let mut nonce_copy = vec![0; len];
     nonce_copy.clone_from_slice(&nonce_root);
 
-    for i in 0..(usize::BITS / 8) {
+    let mut num = u64::from_be_bytes(nonce_copy[start..len].try_into().unwrap());
-        nonce_copy[len - (i as usize + 1)] += ((offset & (0xff << (8 * i))) >> (8 * i)) as u8;
+    num += offset as u64;
-    }
+    nonce_copy[start..len].copy_from_slice(&num.to_be_bytes());
 
     nonce_copy
 }
 
 impl EncryptionContext {
-    fn decrypt_block(&self, ciphertext: &Vec<u8>, offset: usize) -> Result<Vec<u8>, &'static str> {
+    fn decrypt_block(
+        &self,
+        ciphertext: &Vec<u8>,
+        offset: usize,
+    ) -> Result<EncryptionStream, String> {
         let plaintext = match self.enc_type {
             EncryptionType::Invalid => unreachable!(),
             EncryptionType::AesGcm => {
@@ -616,13 +679,20 @@ impl EncryptionContext {
         };
 
         if plaintext.is_err() {
-            return Err("Decryption error");
+            return Err(format!("Decryption error: {}", plaintext.unwrap_err()));
         }
 
-        Ok(plaintext.unwrap())
+        Ok(EncryptionStream {
+            data: plaintext.unwrap(),
+            pos: 0,
+        })
     }
 
-    fn encrypt_block(&self, plaintext: &Vec<u8>, offset: usize) -> Result<Vec<u8>, &'static str> {
+    fn encrypt_block(
+        &self,
+        plaintext: &Vec<u8>,
+        offset: usize,
+    ) -> Result<EncryptionStream, &'static str> {
         let ciphertext = match self.enc_type {
             EncryptionType::Invalid => unreachable!(),
             EncryptionType::AesGcm => {
@@ -641,7 +711,10 @@ impl EncryptionContext {
             return Err("Encryption error");
         }
 
-        Ok(ciphertext.unwrap())
+        Ok(EncryptionStream {
+            data: ciphertext.unwrap(),
+            pos: 0,
+        })
     }
 
     fn decrypt_block_from_file(
@@ -649,13 +722,12 @@ impl EncryptionContext {
         len: usize,
         offset: usize,
         file: &mut BufReader<File>,
-    ) -> Result<Vec<u8>, &'static str> {
+    ) -> Result<EncryptionStream, String> {
         let mut buf: Vec<u8> = vec![0; len + 16];
-        if file.read_exact(&mut buf).is_err() {
+        match file.read_exact(&mut buf) {
-            return Err("Read error.");
+            Err(err) => Err(format!("Read error: {}", err.to_string())),
+            Ok(_) => self.decrypt_block(&buf, offset),
         }
-
-        self.decrypt_block(&buf, offset)
     }
 }
 
@@ -680,6 +752,7 @@ mod tests {
 
     #[test]
     fn nonce_offset_increment() {
+        {
             let nonce: Vec<u8> = vec![0; 12];
 
             let nonce_off = nonce_offset(&nonce, 0xfeeddeadbeef);
@@ -691,12 +764,32 @@ mod tests {
                     && nonce_off[8] == 0xde
                     && nonce_off[7] == 0xed
                     && nonce_off[6] == 0xfe,
-            "\nn1 = {:x?}\nn2 = {:x?}",
+                "\ninitial = {:x?}\nnew     = {:x?}",
                 nonce,
                 nonce_off
             );
         }
 
+        //Carry test
+        {
+            let nonce: Vec<u8> = vec![0xfe; 12];
+
+            let nonce_off = nonce_offset(&nonce, 0x2);
+
+            assert!(
+                nonce_off[11] == 0x00
+                    && nonce_off[10] == 0xff
+                    && nonce_off[9] == 0xfe
+                    && nonce_off[8] == 0xfe
+                    && nonce_off[7] == 0xfe
+                    && nonce_off[6] == 0xfe,
+                "\ninitial = {:x?}\nnew     = {:x?}",
+                nonce,
+                nonce_off
+            );
+        }
+    }
+
     #[test]
     fn localstore_record_serialize() {
         let record = LocalRecord {

src/store/mod.rs

@@ -99,6 +99,7 @@ impl Read for EncryptionStream {
         let range = self.pos..(self.pos + buf.len()).clamp(0, self.data.len() - buf.len());
         let size = range.len();
         buf.copy_from_slice(&self.data[range]);
+        self.pos += size;
         Ok(size)
     }
 }
@@ -146,3 +147,58 @@ impl Seek for EncryptionStream {
         Ok(cur as u64)
     }
 }
+
+impl EncryptionStream {
+    pub fn len(&self) -> usize {
+        self.data.len()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::io::{Read, Seek, SeekFrom};
+
+    use super::EncryptionStream;
+
+    fn make_enc_stream() -> EncryptionStream {
+        EncryptionStream {
+            data: vec![0; 128],
+            pos: 0,
+        }
+    }
+
+    #[test]
+    fn encryption_stream_seek() {
+        let mut stream = make_enc_stream();
+
+        {
+            {
+                let pos = stream.seek(SeekFrom::Start(64));
+                assert!(pos.is_ok());
+                assert_eq!(pos.unwrap(), 0);
+            }
+
+            {
+                let pos = stream.seek(SeekFrom::Start(128));
+                assert!(pos.is_ok());
+                assert_eq!(pos.unwrap(), 64);
+            }
+
+            {
+                let pos = stream.seek(SeekFrom::Start(256));
+                assert!(pos.is_err());
+            }
+        }
+    }
+
+    #[test]
+    fn encryption_stream_read() {
+        let mut stream = make_enc_stream();
+        let mut buf: [u8; 8] = [1; 8];
+
+        assert!(stream.read_exact(&mut buf).is_ok());
+
+        assert_eq!(&buf, &[0; 8]);
+        assert_eq!(stream.stream_position().unwrap(), 8);
+    }
+}