From 84674c20e2d1b6947858fdd0fb165ba22f55eef8 Mon Sep 17 00:00:00 2001
From: Niko <info@parlepeuple.fr>
Date: Wed, 31 Aug 2022 00:46:14 +0300
Subject: [PATCH] implement set_enc_key for builders

---
 Cargo.toml                           |  3 ++-
 src/backend/impl_lmdb/environment.rs |  7 +++++++
 src/backend/impl_safe/environment.rs |  7 +++++++
 src/backend/traits.rs                |  3 +++
 src/env.rs                           | 20 ++++++++++++++++++++
 src/manager.rs                       |  1 +
 tests/env-lmdb.rs                    | 16 ++++++++++++++++
 7 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/Cargo.toml b/Cargo.toml
index 5fa6ffd..e36c792 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -4,6 +4,7 @@ authors = [
   "Nan Jiang <najiang@mozilla.com>",
   "Myk Melez <myk@mykzilla.org>",
   "Victor Porof <vporof@mozilla.com>",
+  "NikoPLP <info@parlepeuple.fr>"
 ]
 categories = ["database"]
 description = "A simple, humane, typed key-value storage solution, with encryption at rest"
@@ -13,7 +14,7 @@ exclude = ["/tests/envs/*"]
 homepage = "https://github.com/mozilla/rkv"
 keywords = ["lmdb", "database", "storage"]
 license = "Apache-2.0"
-name = "rkv-lofire"
+name = "rkv"
 readme = "README.md"
 repository = "https://github.com/nikoPLP/rkv"
 version = "0.18.0"
diff --git a/src/backend/impl_lmdb/environment.rs b/src/backend/impl_lmdb/environment.rs
index 590e2c2..906b4ad 100644
--- a/src/backend/impl_lmdb/environment.rs
+++ b/src/backend/impl_lmdb/environment.rs
@@ -24,6 +24,8 @@ use crate::backend::traits::{
     BackendRoCursorTransaction, BackendStat,
 };
 
+use crate::env::Key;
+
 #[derive(Debug, PartialEq, Eq, Copy, Clone)]
 pub struct EnvironmentBuilderImpl {
     builder: lmdb::EnvironmentBuilder,
@@ -81,6 +83,11 @@ impl<'b> BackendEnvironmentBuilder<'b> for EnvironmentBuilderImpl {
         self
     }
 
+    fn set_enc_key(&mut self, key: Key) -> &mut Self {
+        self.builder.set_enc_key(key);
+        self
+    }
+
     fn set_make_dir_if_needed(&mut self, make_dir_if_needed: bool) -> &mut Self {
         self.make_dir_if_needed = make_dir_if_needed;
         self
diff --git a/src/backend/impl_safe/environment.rs b/src/backend/impl_safe/environment.rs
index edfeb09..dfaede5 100644
--- a/src/backend/impl_safe/environment.rs
+++ b/src/backend/impl_safe/environment.rs
@@ -25,6 +25,7 @@ use super::{
     RoTransactionImpl, RwTransactionImpl, StatImpl,
 };
 use crate::backend::traits::{BackendEnvironment, BackendEnvironmentBuilder};
+use crate::env::Key;
 
 const DEFAULT_DB_FILENAME: &str = "data.safe.bin";
 
@@ -75,6 +76,12 @@ impl<'b> BackendEnvironmentBuilder<'b> for EnvironmentBuilderImpl {
         self
     }
 
+    fn set_enc_key(&mut self, _key: Key) -> &mut Self {
+        // NOOP
+        warn!("set_enc_key is ignored by this storage backend.");
+        self
+    }
+
     fn set_map_size(&mut self, map_size: usize) -> &mut Self {
         self.map_size = Some(map_size);
         self
diff --git a/src/backend/traits.rs b/src/backend/traits.rs
index e034b6f..5befa51 100644
--- a/src/backend/traits.rs
+++ b/src/backend/traits.rs
@@ -15,6 +15,7 @@ use std::{
 
 use crate::{
     backend::common::{DatabaseFlags, EnvironmentFlags, WriteFlags},
+    env::Key,
     error::StoreError,
 };
 
@@ -81,6 +82,8 @@ pub trait BackendEnvironmentBuilder<'b>: Debug + Eq + PartialEq + Copy + Clone {
 
     fn set_map_size(&mut self, size: usize) -> &mut Self;
 
+    fn set_enc_key(&mut self, key: Key) -> &mut Self;
+
     fn set_make_dir_if_needed(&mut self, make_dir_if_needed: bool) -> &mut Self;
 
     fn set_discard_if_corrupted(&mut self, discard_if_corrupted: bool) -> &mut Self;
diff --git a/src/env.rs b/src/env.rs
index c1e74da..ce65fab 100644
--- a/src/env.rs
+++ b/src/env.rs
@@ -14,6 +14,8 @@ use std::{
     path::{Path, PathBuf},
 };
 
+pub type Key = [u8; 32];
+
 #[cfg(any(feature = "db-dup-sort", feature = "db-int-key"))]
 use crate::backend::{BackendDatabaseFlags, DatabaseFlags};
 use crate::{
@@ -79,6 +81,24 @@ where
         Rkv::from_builder(path, builder)
     }
 
+    /// Return a new Rkv environment that supports the specified number of open databases.
+    pub fn with_encryption_key_and_mapsize<B>(
+        path: &Path,
+        key: Key,
+        size: usize,
+    ) -> Result<Rkv<E>, StoreError>
+    where
+        B: BackendEnvironmentBuilder<'e, Environment = E>,
+    {
+        let mut builder = B::new();
+        builder.set_enc_key(key);
+        builder.set_map_size(size);
+        builder.set_max_dbs(DEFAULT_MAX_DBS);
+
+        // Future: set flags, maximum size, etc. here if necessary.
+        Rkv::from_builder(path, builder)
+    }
+
     /// Return a new Rkv environment from the provided builder.
     pub fn from_builder<B>(path: &Path, builder: B) -> Result<Rkv<E>, StoreError>
     where
diff --git a/src/manager.rs b/src/manager.rs
index 0455644..26665c7 100644
--- a/src/manager.rs
+++ b/src/manager.rs
@@ -190,6 +190,7 @@ impl Manager<SafeModeEnvironment> {
     }
 }
 
+#[allow(unused_imports)]
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/tests/env-lmdb.rs b/tests/env-lmdb.rs
index 2c32d69..e2956c2 100644
--- a/tests/env-lmdb.rs
+++ b/tests/env-lmdb.rs
@@ -81,6 +81,22 @@ fn test_open() {
     check_rkv(&k);
 }
 
+#[test]
+fn test_open_with_encryption() {
+    let root = Builder::new()
+        .prefix("test_open")
+        .tempdir()
+        .expect("tempdir");
+    println!("Root path: {:?}", root.path());
+    fs::create_dir_all(root.path()).expect("dir created");
+    assert!(root.path().is_dir());
+
+    let key: [u8; 32] = [0; 32];
+    let k = Rkv::with_encryption_key_and_mapsize::<Lmdb>(root.path(), key, 2 * 1024 * 1024 * 1024)
+        .expect("new succeeded");
+    check_rkv(&k);
+}
+
 #[test]
 fn test_open_from_builder() {
     let root = Builder::new()