From 9db07297acd30c4c9921ef3292d4f5cb679915c0 Mon Sep 17 00:00:00 2001 From: femsci Date: Thu, 2 Feb 2023 00:42:55 +0100 Subject: [PATCH] Index deserialization and traversal. Tests. --- src/store/local.rs | 207 ++++++++++++++++++++++++++++++++------------- src/store/mod.rs | 56 ++++++++++++ 2 files changed, 206 insertions(+), 57 deletions(-) diff --git a/src/store/local.rs b/src/store/local.rs index 209582f..e730212 100644 --- a/src/store/local.rs +++ b/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) -> bool { header.chksum_crc = header.get_checksum(); + let mut index_block: Vec = 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) -> Result for IndexNodeType { } } +struct IndexNodeEntry { + key: u64, + pointer: u64, +} + struct IndexNode { node_type: IndexNodeType, - children: [u64], + children: Vec, +} + +impl IndexNode { + fn children(&self) -> &Vec { + &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 { let mut file: BufReader = BufReader::new(File::open(&self.path).expect("Cannot open file")); @@ -551,33 +606,36 @@ 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 subnode_pos > self.meta.index_node_arity { - return None; + if node.node_type == IndexNodeType::Invalid { + 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); + } } - subnode_pos += 1; - continue; + return None; + } + + 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, offset: usize) -> Vec { 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) { - nonce_copy[len - (i as usize + 1)] += ((offset & (0xff << (8 * i))) >> (8 * i)) as u8; - } + let mut num = u64::from_be_bytes(nonce_copy[start..len].try_into().unwrap()); + 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, offset: usize) -> Result, &'static str> { + fn decrypt_block( + &self, + ciphertext: &Vec, + offset: usize, + ) -> Result { 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, offset: usize) -> Result, &'static str> { + fn encrypt_block( + &self, + plaintext: &Vec, + offset: usize, + ) -> Result { 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, - ) -> Result, &'static str> { + ) -> Result { let mut buf: Vec = vec![0; len + 16]; - if file.read_exact(&mut buf).is_err() { - return Err("Read error."); + match file.read_exact(&mut buf) { + Err(err) => Err(format!("Read error: {}", err.to_string())), + Ok(_) => self.decrypt_block(&buf, offset), } - - self.decrypt_block(&buf, offset) } } @@ -680,21 +752,42 @@ mod tests { #[test] fn nonce_offset_increment() { - let nonce: Vec = vec![0; 12]; + { + let nonce: Vec = vec![0; 12]; - let nonce_off = nonce_offset(&nonce, 0xfeeddeadbeef); + let nonce_off = nonce_offset(&nonce, 0xfeeddeadbeef); - assert!( - nonce_off[11] == 0xef - && nonce_off[10] == 0xbe - && nonce_off[9] == 0xad - && nonce_off[8] == 0xde - && nonce_off[7] == 0xed - && nonce_off[6] == 0xfe, - "\nn1 = {:x?}\nn2 = {:x?}", - nonce, - nonce_off - ); + assert!( + nonce_off[11] == 0xef + && nonce_off[10] == 0xbe + && nonce_off[9] == 0xad + && nonce_off[8] == 0xde + && nonce_off[7] == 0xed + && nonce_off[6] == 0xfe, + "\ninitial = {:x?}\nnew = {:x?}", + nonce, + nonce_off + ); + } + + //Carry test + { + let nonce: Vec = 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] diff --git a/src/store/mod.rs b/src/store/mod.rs index 99a01e8..06e8464 100644 --- a/src/store/mod.rs +++ b/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); + } +}