From 9cc209c26202efb4f493475879146ce7601f12a1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cl=C3=A9ment=20Renault?= Date: Sat, 30 Nov 2024 13:30:40 +0100 Subject: [PATCH] Factorize Env opening function --- heed/src/envs/env_open_options.rs | 112 ++++++++---------------------- heed/src/mdb/lmdb_ffi.rs | 4 +- 2 files changed, 31 insertions(+), 85 deletions(-) diff --git a/heed/src/envs/env_open_options.rs b/heed/src/envs/env_open_options.rs index 859a91a2..1ec1a8f8 100644 --- a/heed/src/envs/env_open_options.rs +++ b/heed/src/envs/env_open_options.rs @@ -18,6 +18,7 @@ use super::{canonicalize_path, OPENED_ENV}; #[cfg(windows)] use crate::envs::OsStrExtLmdb as _; use crate::mdb::error::mdb_result; +#[cfg(master3)] use crate::mdb::ffi; use crate::{EnvFlags, Error, Result}; @@ -151,79 +152,11 @@ impl EnvOpenOptions { /// [^7]: /// [^8]: pub unsafe fn open>(&self, path: P) -> Result { - /// TODO change the function description - /// and deduplicate the code - let mut lock = OPENED_ENV.write().unwrap(); - - let path = match canonicalize_path(path.as_ref()) { - Err(err) => { - if err.kind() == NotFound && self.flags.contains(EnvFlags::NO_SUB_DIR) { - let path = path.as_ref(); - match path.parent().zip(path.file_name()) { - Some((dir, file_name)) => canonicalize_path(dir)?.join(file_name), - None => return Err(err.into()), - } - } else { - return Err(err.into()); - } - } - Ok(path) => path, - }; - - if lock.contains(&path) { - Err(Error::EnvAlreadyOpened) - } else { - let path_str = CString::new(path.as_os_str().as_bytes()).unwrap(); - - unsafe { - let mut env: *mut ffi::MDB_env = ptr::null_mut(); - mdb_result(ffi::mdb_env_create(&mut env))?; - - if let Some(size) = self.map_size { - if size % page_size::get() != 0 { - let msg = format!( - "map size ({}) must be a multiple of the system page size ({})", - size, - page_size::get() - ); - return Err(Error::Io(io::Error::new(io::ErrorKind::InvalidInput, msg))); - } - mdb_result(ffi::mdb_env_set_mapsize(env, size))?; - } - - if let Some(readers) = self.max_readers { - mdb_result(ffi::mdb_env_set_maxreaders(env, readers))?; - } - - if let Some(dbs) = self.max_dbs { - mdb_result(ffi::mdb_env_set_maxdbs(env, dbs))?; - } - - // When the `read-txn-no-tls` feature is enabled, we must force LMDB - // to avoid using the thread local storage, this way we allow users - // to use references of RoTxn between threads safely. - let flags = if cfg!(feature = "read-txn-no-tls") { - // TODO make this a ZST flag on the Env and on RoTxn (make them Send when we can) - self.flags | EnvFlags::NO_TLS - } else { - self.flags - }; - - let result = ffi::mdb_env_open(env, path_str.as_ptr(), flags.bits(), 0o600); - match mdb_result(result) { - Ok(()) => { - let env_ptr = NonNull::new(env).unwrap(); - let inserted = lock.insert(path.clone()); - debug_assert!(inserted); - Ok(Env::new(env_ptr, path)) - } - Err(e) => { - ffi::mdb_env_close(env); - Err(e.into()) - } - } - } - } + self.raw_open_with_encryption( + path.as_ref(), + #[cfg(master3)] + None, + ) } /// Creates a blank new set of options ready for configuration and specifies that @@ -337,12 +270,23 @@ impl EnvOpenOptions { E: AeadMutInPlace + KeyInit, P: AsRef, { + self.raw_open_with_encryption( + path.as_ref(), + Some((Some(encrypt_func_wrapper::), &key, ::TagSize::U32)), + ) + .map(|inner| EncryptedEnv { inner }) + } + + fn raw_open_with_encryption( + &self, + path: &Path, + #[cfg(master3)] enc: Option<(ffi::MDB_enc_func, &[u8], u32)>, + ) -> Result { let mut lock = OPENED_ENV.write().unwrap(); - let path = match canonicalize_path(path.as_ref()) { + let path = match canonicalize_path(path) { Err(err) => { if err.kind() == NotFound && self.flags.contains(EnvFlags::NO_SUB_DIR) { - let path = path.as_ref(); match path.parent().zip(path.file_name()) { Some((dir, file_name)) => canonicalize_path(dir)?.join(file_name), None => return Err(err.into()), @@ -363,13 +307,15 @@ impl EnvOpenOptions { let mut env: *mut ffi::MDB_env = ptr::null_mut(); mdb_result(ffi::mdb_env_create(&mut env))?; - let encrypt_key = crate::into_val(&key); - mdb_result(ffi::mdb_env_set_encrypt( - env, - Some(encrypt_func_wrapper::), - &encrypt_key, - ::TagSize::U32, - ))?; + #[cfg(master3)] + if let Some((encrypt_func, key, tag_size)) = enc { + mdb_result(ffi::mdb_env_set_encrypt( + env, + encrypt_func, + &crate::into_val(key), + tag_size, + ))?; + } if let Some(size) = self.map_size { if size % page_size::get() != 0 { @@ -407,7 +353,7 @@ impl EnvOpenOptions { let env_ptr = NonNull::new(env).unwrap(); let inserted = lock.insert(path.clone()); debug_assert!(inserted); - Ok(EncryptedEnv { inner: Env::new(env_ptr, path) }) + Ok(Env::new(env_ptr, path)) } Err(e) => { ffi::mdb_env_close(env); diff --git a/heed/src/mdb/lmdb_ffi.rs b/heed/src/mdb/lmdb_ffi.rs index cb8b3ba5..1f68da0e 100644 --- a/heed/src/mdb/lmdb_ffi.rs +++ b/heed/src/mdb/lmdb_ffi.rs @@ -1,7 +1,5 @@ use std::ptr; -#[cfg(master3)] -pub use ffi::mdb_env_set_encrypt; pub use ffi::{ mdb_cursor_close, mdb_cursor_del, mdb_cursor_get, mdb_cursor_open, mdb_cursor_put, mdb_dbi_open, mdb_del, mdb_drop, mdb_env_close, mdb_env_copyfd2, mdb_env_create, @@ -13,6 +11,8 @@ pub use ffi::{ MDB_RDONLY, MDB_RESERVE, }; #[cfg(master3)] +pub use ffi::{mdb_env_set_encrypt, MDB_enc_func}; +#[cfg(master3)] use lmdb_master3_sys as ffi; #[cfg(not(master3))] use lmdb_master_sys as ffi;