Index deserialization and traversal. Tests.

This commit is contained in:
femsci 2023-02-02 00:42:55 +01:00
parent 5283830c3c
commit 9db07297ac
Signed by: femsci
GPG key ID: 08F7911F0E650C67
2 changed files with 206 additions and 57 deletions

View file

@ -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 subnode_pos > self.meta.index_node_arity {
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);
}
}
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) {
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<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() {
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,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 {

View file

@ -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);
}
}