Skip to content

Commit

Permalink
WIP use the aead crate and traits instead
Browse files Browse the repository at this point in the history
  • Loading branch information
@Kerollmops
Kerollmops committed Nov 12, 2022
1 parent 9c585f7 commit 7bf57bf
Show file tree
Hide file tree
Showing 4 changed files with 103 additions and 111 deletions.
3 changes: 2 additions & 1 deletion heed/Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ readme = "../README.md"
edition = "2021"

[dependencies]
aead = "0.5.1"
bytemuck = "1.12.1"
byteorder = { version = "1.4.3", default-features = false }
heed-traits = { version = "0.7.0", path = "../heed-traits" }
Expand All @@ -24,7 +25,7 @@ synchronoise = "1.0.1"

[dev-dependencies]
bytemuck = { version = "1.12.1", features = ["derive"] }
chacha20 = "0.9.0"
chacha20poly1305 = "0.10.1"
crc32fast = "1.3.2"
serde = { version = "1.0.144", features = ["derive"] }
tempfile = "3.3.0"
Expand Down
49 changes: 3 additions & 46 deletions heed/examples/encrypt.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,62 +2,19 @@ use std::error::Error;
use std::fs;
use std::path::Path;

use chacha20::cipher::{KeyIvInit, StreamCipher};
use chacha20::ChaCha20;
use chacha20poly1305::ChaCha20Poly1305;
use heed::types::*;
use heed::{Checksum, Database, Encrypt, EncryptDecrypt, EnvOpenOptions};

enum Crc32Checksum {}

impl Checksum for Crc32Checksum {
const SIZE: u32 = 32 / 8;

fn name() -> String {
String::from("crc32")
}

fn checksum(input: &[u8], output: &mut [u8], _key: Option<&[u8]>) {
let checksum = crc32fast::hash(input);
output.copy_from_slice(&checksum.to_le_bytes());
}
}

enum Chacha20Encrypt {}

impl Encrypt for Chacha20Encrypt {
fn name() -> String {
String::from("chacha20")
}

fn encrypt_decrypt(
_action: EncryptDecrypt,
input: &[u8],
output: &mut [u8],
key: &[u8],
iv: &[u8],
_auth: &[u8],
) -> Result<(), ()> {
Ok(ChaCha20::new_from_slices(key, &iv[..12])
.map_err(drop)?
.apply_keystream_b2b(input, output)
.map_err(drop)?)
}
}
use heed::{Database, EnvOpenOptions};

fn main() -> Result<(), Box<dyn Error>> {
let env_path = Path::new("target").join("encrypt.mdb");
let password: &[_; 32] = b"I told you this is my password!!";
let mac_size = 0;

let _ = fs::remove_dir_all(&env_path);
fs::create_dir_all(&env_path)?;

// We open the environment
let mut options = EnvOpenOptions::new()
.encrypt_with::<Chacha20Encrypt>(password.to_vec(), mac_size)
// By setting the checksum function we will have checksum errors if the decryption
// fail instead of random LMDB errors due to invalid data in the decrypted pages
.checksum_with::<Crc32Checksum>();
let mut options = EnvOpenOptions::new().encrypt_with::<ChaCha20Poly1305>(password.to_vec());
let env = options
.map_size(10 * 1024 * 1024) // 10MB
.max_dbs(3)
Expand Down
156 changes: 95 additions & 61 deletions heed/src/env.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,6 @@ use std::os::unix::{
};
use std::panic::catch_unwind;
use std::path::{Path, PathBuf};
use std::result::Result as StdResult;
use std::sync::{Arc, RwLock};
use std::time::Duration;
#[cfg(windows)]
Expand All @@ -20,6 +19,10 @@ use std::{
};
use std::{fmt, io, mem, ptr, sync};

use aead::consts::U0;
use aead::generic_array::typenum::Unsigned;
use aead::generic_array::GenericArray;
use aead::{AeadCore, AeadMutInPlace, Key, KeyInit, KeySizeUser, Nonce, Tag};
use lmdb_master3_sys::MDB_val;
use once_cell::sync::Lazy;
use synchronoise::event::SignalEvent;
Expand All @@ -46,8 +49,8 @@ struct EnvEntry {
pub struct SimplifiedOpenOptions {
/// The name of the checksum algorithm this [`Env`] has been opened with.
pub checksum_name: Option<String>,
/// The name of the encryption/decryption algorithm this [`Env`] has been opened with.
pub encrypt_name: Option<String>,
/// Weither this [`Env`] has been opened with an encryption/decryption algorithm.
pub use_encryption: bool,
/// The maximum size this [`Env`] with take in bytes or [`None`] if it was not specified.
pub map_size: Option<usize>,
/// The maximum number of concurrent readers or [`None`] if it was not specified.
Expand All @@ -58,12 +61,14 @@ pub struct SimplifiedOpenOptions {
pub flags: u32,
}

impl<E: Encrypt, C: Checksum> From<&EnvOpenOptions<E, C>> for SimplifiedOpenOptions {
impl<E: AeadMutInPlace + KeyInit, C: Checksum> From<&EnvOpenOptions<E, C>>
for SimplifiedOpenOptions
{
fn from(eoo: &EnvOpenOptions<E, C>) -> SimplifiedOpenOptions {
let EnvOpenOptions { checksum, encrypt, map_size, max_readers, max_dbs, flags } = eoo;
SimplifiedOpenOptions {
checksum_name: checksum.map(|_| E::name()),
encrypt_name: encrypt.as_ref().map(|_| C::name()),
checksum_name: checksum.map(|_| C::name()),
use_encryption: encrypt.is_some(),
map_size: *map_size,
max_readers: *max_readers,
max_dbs: *max_dbs,
Expand Down Expand Up @@ -156,72 +161,65 @@ impl Checksum for DummyChecksum {
fn checksum(_input: &[u8], _output: &mut [u8], _key: Option<&[u8]>) {}
}

/// Describes an encryption/decryption algorithm to ensure pages
/// are readable only with the right key.
pub trait Encrypt {
/// The name of the encryption/decryption algorithm, used for safety purposes.
///
/// Make sure that the name corresponds to the algorithm, it will be compared
/// whenever an [`Env`] is opened and tried to be opened a second time.
fn name() -> String;

/// Takes the input bytes and writes them in the output slice of the same length.
fn encrypt_decrypt(
action: EncryptDecrypt,
input: &[u8],
output: &mut [u8],
key: &[u8],
iv: &[u8],
auth: &[u8],
) -> StdResult<(), ()>;
}

/// The action to perform on the page.
#[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub enum EncryptDecrypt {
Encrypt,
Decrypt,
}

/// A dummy encryption/decryption algorithm that must never be used.
/// Only here for Rust API purposes.
pub enum DummyEncrypt {}

impl Encrypt for DummyEncrypt {
fn name() -> String {
String::new()
impl AeadMutInPlace for DummyEncrypt {
fn encrypt_in_place_detached(
&mut self,
_nonce: &Nonce<Self>,
_associated_data: &[u8],
_buffer: &mut [u8],
) -> aead::Result<Tag<Self>> {
Err(aead::Error)
}

fn decrypt_in_place_detached(
&mut self,
_nonce: &Nonce<Self>,
_associated_data: &[u8],
_buffer: &mut [u8],
_tag: &Tag<Self>,
) -> aead::Result<()> {
Err(aead::Error)
}
}

impl AeadCore for DummyEncrypt {
type NonceSize = U0;
type TagSize = U0;
type CiphertextOverhead = U0;
}

fn encrypt_decrypt(
_action: EncryptDecrypt,
_input: &[u8],
_output: &mut [u8],
_key: &[u8],
_iv: &[u8],
_auth: &[u8],
) -> StdResult<(), ()> {
Err(())
impl KeySizeUser for DummyEncrypt {
type KeySize = U0;
}

impl KeyInit for DummyEncrypt {
fn new(_key: &GenericArray<u8, Self::KeySize>) -> Self {
todo!()
}
}

/// Options and flags which can be used to configure how an environment is opened.
#[derive(Clone, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct EnvOpenOptions<E: Encrypt = DummyEncrypt, C: Checksum = DummyChecksum> {
pub struct EnvOpenOptions<E: AeadMutInPlace + KeyInit = DummyEncrypt, C: Checksum = DummyChecksum> {
checksum: Option<PhantomData<C>>,
encrypt: Option<(PhantomData<E>, Vec<u8>, u32)>,
encrypt: Option<(PhantomData<E>, Vec<u8>)>,
map_size: Option<usize>,
max_readers: Option<u32>,
max_dbs: Option<u32>,
flags: u32,
}

impl<E: Encrypt, C: Checksum> fmt::Debug for EnvOpenOptions<E, C> {
impl<E: AeadMutInPlace + KeyInit, C: Checksum> fmt::Debug for EnvOpenOptions<E, C> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
let EnvOpenOptions { checksum, encrypt, map_size, max_readers, max_dbs, flags } = self;
f.debug_struct("EnvOpenOptions")
.field("checksum", &checksum.map(|_| C::name()))
.field("encrypt", &encrypt.as_ref().map(|_| E::name()))
.field("encrypt", &encrypt.is_some())
.field("map_size", &map_size)
.field("max_readers", &max_readers)
.field("max_dbs", &max_dbs)
Expand All @@ -244,7 +242,7 @@ impl EnvOpenOptions {
}
}

impl<E: Encrypt, C: Checksum> EnvOpenOptions<E, C> {
impl<E: AeadMutInPlace + KeyInit, C: Checksum> EnvOpenOptions<E, C> {
/// Set the size of the memory map to use for this environment.
pub fn map_size(&mut self, size: usize) -> &mut Self {
self.map_size = Some(size);
Expand All @@ -268,11 +266,11 @@ impl<E: Encrypt, C: Checksum> EnvOpenOptions<E, C> {
///
/// It is advised to use a checksum algorithm when an encryption/decryption algorithm
/// is specified to get better error messages when the encryption key is wrong.
pub fn encrypt_with<F: Encrypt>(self, key: Vec<u8>, auth_size: u32) -> EnvOpenOptions<F, C> {
pub fn encrypt_with<F: AeadMutInPlace + KeyInit>(self, key: Vec<u8>) -> EnvOpenOptions<F, C> {
let EnvOpenOptions { checksum, encrypt: _, map_size, max_readers, max_dbs, flags } = self;
EnvOpenOptions {
checksum,
encrypt: Some((PhantomData, key, auth_size)),
encrypt: Some((PhantomData, key)),
map_size,
max_readers,
max_dbs,
Expand Down Expand Up @@ -373,13 +371,13 @@ impl<E: Encrypt, C: Checksum> EnvOpenOptions<E, C> {
))?;
}

if let Some((_marker, key, auth_size)) = &self.encrypt {
if let Some((_marker, key)) = &self.encrypt {
let key = crate::into_val(key);
mdb_result(ffi::mdb_env_set_encrypt(
env,
Some(encrypt_func_wrapper::<E>),
&key,
*auth_size,
<E as AeadCore>::TagSize::U32,
))?;
}

Expand Down Expand Up @@ -446,9 +444,41 @@ impl<E: Encrypt, C: Checksum> EnvOpenOptions<E, C> {
}
}

fn encrypt<A: AeadMutInPlace + KeyInit>(
key: &[u8],
nonce: &[u8],
aad: &[u8],
plaintext: &[u8],
chipertext_out: &mut [u8],
auth_out: &mut [u8],
) {
chipertext_out.copy_from_slice(plaintext);
let key: &Key<A> = key.try_into().unwrap();
let nonce: &Nonce<A> = nonce.try_into().unwrap();
let mut aead = A::new(key);
let tag = aead.encrypt_in_place_detached(nonce, aad, chipertext_out).unwrap();
auth_out.copy_from_slice(&tag);
}

fn decrypt<A: AeadMutInPlace + KeyInit>(
key: &[u8],
nonce: &[u8],
aad: &[u8],
chipher_text: &[u8],
output: &mut [u8],
auth_in: &[u8],
) -> aead::Result<()> {
output.copy_from_slice(chipher_text);
let key: &Key<A> = key.try_into().unwrap();
let nonce: &Nonce<A> = nonce.try_into().unwrap();
let tag: &Tag<A> = auth_in.try_into().unwrap();
let mut aead = A::new(key);
aead.decrypt_in_place_detached(nonce, aad, output, tag)
}

/// The wrapper function that is called by LMDB that directly calls
/// the Rust idiomatic function internally.
unsafe extern "C" fn encrypt_func_wrapper<E: Encrypt>(
unsafe extern "C" fn encrypt_func_wrapper<E: AeadMutInPlace + KeyInit>(
src: *const MDB_val,
dst: *mut MDB_val,
key_ptr: *const MDB_val,
Expand All @@ -458,18 +488,22 @@ unsafe extern "C" fn encrypt_func_wrapper<E: Encrypt>(
let input = std::slice::from_raw_parts((*src).mv_data as *const u8, (*src).mv_size);
let output = std::slice::from_raw_parts_mut((*dst).mv_data as *mut u8, (*dst).mv_size);
let key = std::slice::from_raw_parts((*key_ptr).mv_data as *const u8, (*key_ptr).mv_size);
let iv = std::slice::from_raw_parts(
let nonce = std::slice::from_raw_parts(
(*key_ptr.offset(1)).mv_data as *const u8,
(*key_ptr.offset(1)).mv_size,
);
let auth = std::slice::from_raw_parts(
(*key_ptr.offset(2)).mv_data as *const u8,
let auth = std::slice::from_raw_parts_mut(
(*key_ptr.offset(2)).mv_data as *mut u8,
(*key_ptr.offset(2)).mv_size,
);

let action = if encdec == 1 { EncryptDecrypt::Encrypt } else { EncryptDecrypt::Decrypt };

E::encrypt_decrypt(action, input, output, key, iv, auth).is_err() as i32
let aad = [];
if encdec == 1 {
encrypt::<E>(&key, nonce, &aad, input, output, auth);
1
} else {
decrypt::<E>(&key, nonce, &aad, input, output, auth).is_err() as i32
}
});

match result {
Expand Down
6 changes: 3 additions & 3 deletions heed/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -59,13 +59,13 @@ mod txn;
use std::{error, fmt, io, result};

use heed_traits as traits;
pub use {bytemuck, byteorder, heed_types as types};
pub use {aead, bytemuck, byteorder, heed_types as types};

use self::cursor::{RoCursor, RwCursor};
pub use self::db::{Database, PolyDatabase};
pub use self::env::{
env_closing_event, Checksum, CompactionOption, Encrypt, EncryptDecrypt, Env, EnvClosingEvent,
EnvOpenOptions, SimplifiedOpenOptions,
env_closing_event, Checksum, CompactionOption, Env, EnvClosingEvent, EnvOpenOptions,
SimplifiedOpenOptions,
};
pub use self::iter::{
RoIter, RoPrefix, RoRange, RoRevIter, RoRevPrefix, RoRevRange, RwIter, RwPrefix, RwRange,
Expand Down

0 comments on commit 7bf57bf

Please sign in to comment.