From ea897593788aebed734454056acb709190143186 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cl=C3=A9ment=20Renault?= Date: Sat, 30 Nov 2024 12:10:38 +0100 Subject: [PATCH] Reintroduce #289 --- ...heed3-encryption.rs => heed3-encrypted.rs} | 12 +++---- heed/src/databases/database.rs | 19 +++------- heed/src/databases/encrypted_database.rs | 27 +++++++------- heed/src/envs/env.rs | 20 ++++++----- heed/src/envs/mod.rs | 35 +++++++++++++++++-- heed/src/txn.rs | 2 +- heed3/Cargo.toml | 4 +-- 7 files changed, 72 insertions(+), 47 deletions(-) rename examples/{heed3-encryption.rs => heed3-encrypted.rs} (85%) diff --git a/examples/heed3-encryption.rs b/examples/heed3-encrypted.rs similarity index 85% rename from examples/heed3-encryption.rs rename to examples/heed3-encrypted.rs index da16491f..acb3c671 100644 --- a/examples/heed3-encryption.rs +++ b/examples/heed3-encrypted.rs @@ -4,8 +4,8 @@ use std::path::Path; use argon2::Argon2; use chacha20poly1305::{ChaCha20Poly1305, Key}; -use heed3_encryption::types::*; -use heed3_encryption::{Database, EnvOpenOptions}; +use heed3::types::*; +use heed3::{Database, EnvOpenOptions}; fn main() -> Result<(), Box> { let env_path = Path::new("target").join("encrypt.mdb"); @@ -21,12 +21,12 @@ fn main() -> Result<(), Box> { Argon2::default().hash_password_into(password.as_bytes(), salt.as_bytes(), &mut key)?; // We open the environment - let mut options = EnvOpenOptions::::new_encrypted_with(key); + let mut options = EnvOpenOptions::new(); let env = unsafe { options .map_size(10 * 1024 * 1024) // 10MB .max_dbs(3) - .open(&env_path)? + .open_encrypted::(key, &env_path)? }; let key1 = "first-key"; @@ -36,11 +36,11 @@ fn main() -> Result<(), Box> { // We create database and write secret values in it let mut wtxn = env.write_txn()?; - let db: Database = env.create_database(&mut wtxn, Some("first"))?; + let db = env.create_database::(&mut wtxn, Some("first"))?; db.put(&mut wtxn, key1, val1)?; db.put(&mut wtxn, key2, val2)?; wtxn.commit()?; - env.prepare_for_closing().wait(); + // env.prepare_for_closing().wait(); // We reopen the environment now let env = unsafe { options.open(&env_path)? }; diff --git a/heed/src/databases/database.rs b/heed/src/databases/database.rs index 43e9c0f1..df8295ae 100644 --- a/heed/src/databases/database.rs +++ b/heed/src/databases/database.rs @@ -139,8 +139,7 @@ impl<'e, 'n, KC, DC, C> DatabaseOpenOptions<'e, 'n, KC, DC, C> { { assert_eq_env_txn!(self.env, rtxn); - let raw_txn = unsafe { rtxn.txn.unwrap().as_mut() }; - match self.env.raw_init_database::(raw_txn, self.name, self.flags) { + match self.env.raw_init_database::(rtxn.txn.unwrap(), self.name, self.flags) { Ok(dbi) => Ok(Some(Database::new(self.env.env_mut_ptr().as_ptr() as _, dbi))), Err(Error::Mdb(e)) if e.not_found() => Ok(None), Err(e) => Err(e), @@ -165,8 +164,7 @@ impl<'e, 'n, KC, DC, C> DatabaseOpenOptions<'e, 'n, KC, DC, C> { assert_eq_env_txn!(self.env, wtxn); let flags = self.flags | AllDatabaseFlags::CREATE; - let raw_txn = unsafe { wtxn.txn.txn.unwrap().as_mut() }; - match self.env.raw_init_database::(raw_txn, self.name, flags) { + match self.env.raw_init_database::(wtxn.txn.txn.unwrap(), self.name, flags) { Ok(dbi) => Ok(Database::new(self.env.env_mut_ptr().as_ptr() as _, dbi)), Err(e) => Err(e), } @@ -352,7 +350,6 @@ impl Database { let mut key_val = unsafe { crate::into_val(&key_bytes) }; let mut data_val = mem::MaybeUninit::uninit(); - let mut txn = txn.txn.unwrap(); let result = unsafe { mdb_result(ffi::mdb_get( @@ -1129,7 +1126,7 @@ impl Database { RoCursor::new(txn, self.dbi).map(|cursor| RoRevIter::new(cursor)) } - /// Return a mutable reversed lexicographically ordered iterator of all key-value\ + /// Return a mutable reverse ordered iterator of all key-value\ /// pairs in this database. /// /// ``` @@ -1842,7 +1839,6 @@ impl Database { let mut key_val = unsafe { crate::into_val(&key_bytes) }; let mut data_val = unsafe { crate::into_val(&data_bytes) }; let flags = 0; - let mut txn = txn.txn.txn.unwrap(); unsafe { mdb_result(ffi::mdb_put( @@ -1910,7 +1906,6 @@ impl Database { let mut key_val = unsafe { crate::into_val(&key_bytes) }; let mut reserved = ffi::reserve_size_val(data_size); let flags = ffi::MDB_RESERVE; - let mut txn = txn.txn.txn.unwrap(); unsafe { mdb_result(ffi::mdb_put( @@ -2009,7 +2004,6 @@ impl Database { let mut key_val = unsafe { crate::into_val(&key_bytes) }; let mut data_val = unsafe { crate::into_val(&data_bytes) }; let flags = flags.bits(); - let mut txn = txn.txn.txn.unwrap(); unsafe { mdb_result(ffi::mdb_put( @@ -2123,7 +2117,6 @@ impl Database { let mut key_val = unsafe { crate::into_val(&key_bytes) }; let mut data_val = unsafe { crate::into_val(&data_bytes) }; let flags = (flags | PutFlags::NO_OVERWRITE).bits(); - let mut txn = txn.txn.txn.unwrap(); let result = unsafe { mdb_result(ffi::mdb_put( @@ -2280,7 +2273,6 @@ impl Database { let mut key_val = unsafe { crate::into_val(&key_bytes) }; let mut reserved = ffi::reserve_size_val(data_size); let flags = (flags | PutFlags::NO_OVERWRITE).bits() | ffi::MDB_RESERVE; - let mut txn = txn.txn.txn.unwrap(); let result = unsafe { mdb_result(ffi::mdb_put( @@ -2364,7 +2356,6 @@ impl Database { let key_bytes: Cow<[u8]> = KC::bytes_encode(key).map_err(Error::Encoding)?; let mut key_val = unsafe { crate::into_val(&key_bytes) }; - let mut txn = txn.txn.txn.unwrap(); let result = unsafe { mdb_result(ffi::mdb_del( @@ -2457,7 +2448,6 @@ impl Database { let data_bytes: Cow<[u8]> = DC::bytes_encode(data).map_err(Error::Encoding)?; let mut key_val = unsafe { crate::into_val(&key_bytes) }; let mut data_val = unsafe { crate::into_val(&data_bytes) }; - let mut txn = txn.txn.txn.unwrap(); let result = unsafe { mdb_result(ffi::mdb_del( @@ -2589,7 +2579,6 @@ impl Database { /// ``` pub fn clear(&self, txn: &mut RwTxn) -> Result<()> { assert_eq_env_db_txn!(self, txn); - let mut txn = txn.txn.txn.unwrap(); unsafe { mdb_result(ffi::mdb_drop(txn.txn.txn.unwrap().as_mut(), self.dbi, 0)) @@ -2682,7 +2671,7 @@ mod tests { use heed_types::*; use super::*; - use crate::env::IntegerComparator; + use crate::IntegerComparator; #[test] fn put_overwrite() -> Result<()> { diff --git a/heed/src/databases/encrypted_database.rs b/heed/src/databases/encrypted_database.rs index d55a61da..f67522e7 100644 --- a/heed/src/databases/encrypted_database.rs +++ b/heed/src/databases/encrypted_database.rs @@ -1015,9 +1015,9 @@ impl EncryptedDatabase { self.inner.rev_iter_mut(txn) } - /// Return a lexicographically ordered iterator of a range of key-value pairs in this database. + /// Return an ordered iterator of a range of key-value pairs in this database. /// - /// Comparisons are made by using the bytes representation of the key. + /// Comparisons are made by using the comparator `C`. /// /// You can make this iterator `Send`able between threads by /// using the `read-txn-no-tls` crate feature. @@ -1062,7 +1062,7 @@ impl EncryptedDatabase { &self, txn: &'txn mut RoTxn, range: &'a R, - ) -> Result> + ) -> Result> where KC: BytesEncode<'a>, R: RangeBounds, @@ -1070,10 +1070,10 @@ impl EncryptedDatabase { self.inner.range(txn, range) } - /// Return a mutable lexicographically ordered iterator of a range of + /// Return a mutable ordered iterator of a range of /// key-value pairs in this database. /// - /// Comparisons are made by using the bytes representation of the key. + /// Comparisons are made by using the comparator `C`. /// /// ``` /// # use std::fs; @@ -1128,7 +1128,7 @@ impl EncryptedDatabase { &self, txn: &'txn mut RwTxn, range: &'a R, - ) -> Result> + ) -> Result> where KC: BytesEncode<'a>, R: RangeBounds, @@ -1136,10 +1136,10 @@ impl EncryptedDatabase { self.inner.range_mut(txn, range) } - /// Return a reversed lexicographically ordered iterator of a range of key-value + /// Return a reverse ordered iterator of a range of key-value /// pairs in this database. /// - /// Comparisons are made by using the bytes representation of the key. + /// Comparisons are made by using the comparator `C`. /// /// You can make this iterator `Send`able between threads by /// using the `read-txn-no-tls` crate feature. @@ -1184,7 +1184,7 @@ impl EncryptedDatabase { &self, txn: &'txn mut RoTxn, range: &'a R, - ) -> Result> + ) -> Result> where KC: BytesEncode<'a>, R: RangeBounds, @@ -1192,10 +1192,10 @@ impl EncryptedDatabase { self.inner.rev_range(txn, range) } - /// Return a mutable reversed lexicographically ordered iterator of a range of + /// Return a mutable reverse ordered iterator of a range of /// key-value pairs in this database. /// - /// Comparisons are made by using the bytes representation of the key. + /// Comparisons are made by using the comparator `C`. /// /// ``` /// # use std::fs; @@ -1250,7 +1250,7 @@ impl EncryptedDatabase { &self, txn: &'txn mut RwTxn, range: &'a R, - ) -> Result> + ) -> Result> where KC: BytesEncode<'a>, R: RangeBounds, @@ -2020,7 +2020,7 @@ impl EncryptedDatabase { /// /// Prefer using [`clear`] instead of a call to this method with a full range ([`..`]). /// - /// Comparisons are made by using the bytes representation of the key. + /// Comparisons are made by using the comparator `C`. /// /// [`clear`]: crate::Database::clear /// [`..`]: std::ops::RangeFull @@ -2068,6 +2068,7 @@ impl EncryptedDatabase { pub fn delete_range<'a, 'txn, R>(&self, txn: &'txn mut RwTxn, range: &'a R) -> Result where KC: BytesEncode<'a> + BytesDecode<'txn>, + C: Comparator, R: RangeBounds, { self.inner.delete_range(txn, range) diff --git a/heed/src/envs/env.rs b/heed/src/envs/env.rs index a178747c..3d8760d8 100644 --- a/heed/src/envs/env.rs +++ b/heed/src/envs/env.rs @@ -9,7 +9,7 @@ use heed_traits::Comparator; use super::{ custom_key_cmp_wrapper, get_file_fd, metadata_from_fd, DefaultComparator, EnvInfo, FlagSetMode, - OPENED_ENV, + IntegerComparator, OPENED_ENV, }; use crate::cursor::{MoveOperation, RoCursor}; use crate::mdb::ffi::{self, MDB_env}; @@ -158,8 +158,7 @@ impl Env { let rtxn = self.read_txn()?; // Open the main database - let raw_txn = unsafe { rtxn.txn.unwrap().as_mut() }; - let dbi = self.raw_open_dbi::(raw_txn, None, 0)?; + let dbi = self.raw_open_dbi::(rtxn.txn.unwrap(), None, 0)?; // We're going to iterate on the unnamed database let mut cursor = RoCursor::new(&rtxn, dbi)?; @@ -172,8 +171,9 @@ impl Env { let key = String::from_utf8(key.to_vec()).unwrap(); // Calling `ffi::db_stat` on a database instance does not involve key comparison // in LMDB, so it's safe to specify a noop key compare function for it. - let raw_txn = unsafe { rtxn.txn.unwrap().as_mut() }; - if let Ok(dbi) = self.raw_open_dbi::(raw_txn, Some(&key), 0) { + if let Ok(dbi) = + self.raw_open_dbi::(rtxn.txn.unwrap(), Some(&key), 0) + { let mut stat = mem::MaybeUninit::uninit(); unsafe { mdb_result(ffi::mdb_stat(rtxn.txn.unwrap().as_mut(), dbi, stat.as_mut_ptr()))? @@ -268,7 +268,7 @@ impl Env { fn raw_open_dbi( &self, - raw_txn: NonNull, + mut raw_txn: NonNull, name: Option<&str>, flags: u32, ) -> std::result::Result { @@ -282,13 +282,17 @@ impl Env { // safety: The name cstring is cloned by LMDB, we can drop it after. // If a read-only is used with the MDB_CREATE flag, LMDB will throw an error. unsafe { - mdb_result(ffi::mdb_dbi_open(raw_txn, name_ptr, flags, &mut dbi))?; + mdb_result(ffi::mdb_dbi_open(raw_txn.as_mut(), name_ptr, flags, &mut dbi))?; let cmp_type_id = TypeId::of::(); if cmp_type_id != TypeId::of::() && cmp_type_id != TypeId::of::() { - mdb_result(ffi::mdb_set_compare(raw_txn, dbi, Some(custom_key_cmp_wrapper::)))?; + mdb_result(ffi::mdb_set_compare( + raw_txn.as_mut(), + dbi, + Some(custom_key_cmp_wrapper::), + ))?; } }; diff --git a/heed/src/envs/mod.rs b/heed/src/envs/mod.rs index a5769213..9eb31dd3 100644 --- a/heed/src/envs/mod.rs +++ b/heed/src/envs/mod.rs @@ -8,9 +8,7 @@ use std::os::unix::io::{AsRawFd, BorrowedFd, RawFd}; use std::panic::catch_unwind; use std::path::{Path, PathBuf}; use std::process::abort; -use std::ptr::NonNull; use std::sync::{LazyLock, RwLock}; -use std::time::Duration; #[cfg(windows)] use std::{ ffi::OsStr, @@ -181,6 +179,39 @@ impl LexicographicComparator for DefaultComparator { } } +/// A representation of LMDB's `MDB_INTEGERKEY` comparator behavior. +/// +/// This enum is used to indicate a table should be sorted by the keys numeric +/// value in native byte order. When a [`Database`] is created or opened with +/// [`IntegerComparator`], it signifies that the comparator should not be explicitly +/// set via [`ffi::mdb_set_compare`], instead the flag [`AllDatabaseFlags::INTEGER_KEY`] +/// is set on the table. +/// +/// This can only be used on certain types: either `u32` or `usize`. +/// The keys must all be of the same size. +pub enum IntegerComparator {} + +impl Comparator for IntegerComparator { + fn compare(a: &[u8], b: &[u8]) -> Ordering { + #[cfg(target_endian = "big")] + return a.cmp(b); + + #[cfg(target_endian = "little")] + { + let len = a.len(); + + for i in (0..len).rev() { + match a[i].cmp(&b[i]) { + Ordering::Equal => continue, + other => return other, + } + } + + Ordering::Equal + } + } +} + /// Whether to perform compaction while copying an environment. #[derive(Debug, Copy, Clone)] pub enum CompactionOption { diff --git a/heed/src/txn.rs b/heed/src/txn.rs index 7e0d47de..07ba369e 100644 --- a/heed/src/txn.rs +++ b/heed/src/txn.rs @@ -47,7 +47,7 @@ use crate::Result; pub struct RoTxn<'e> { /// Makes the struct covariant and !Sync pub(crate) txn: Option>, - env: Cow<'e, Env>, + env: &'e Env, } impl<'e> RoTxn<'e> { diff --git a/heed3/Cargo.toml b/heed3/Cargo.toml index 1176d5e0..42a6a7b2 100644 --- a/heed3/Cargo.toml +++ b/heed3/Cargo.toml @@ -118,5 +118,5 @@ longer-keys = ["lmdb-master3-sys/longer-keys"] # Examples are located outside the standard heed/examples directory to prevent # conflicts between heed3 and heed examples when working on both crates. [[example]] -name = "heed3-encryption" -path = "../examples/heed3-encryption.rs" +name = "heed3-encrypted" +path = "../examples/heed3-encrypted.rs"