Index deserialization and traversal. Tests.
This commit is contained in:
parent
5283830c3c
commit
9db07297ac
GPG key ID:
08F7911F0E650C67
2 changed files with 206 additions and 57 deletions
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue