diff --git a/enclave-modules/environment/Cargo.toml b/enclave-modules/environment/Cargo.toml index ca09c9ec..cd9fea66 100644 --- a/enclave-modules/environment/Cargo.toml +++ b/enclave-modules/environment/Cargo.toml @@ -6,8 +6,8 @@ edition = "2021" [dependencies] light-client = { path = "../../modules/light-client", default-features = false } store = { path = "../../modules/store", default-features = false } -enclave-store = { path = "../store", optional = true } +host-api = { path = "../host-api", optional = true } [features] default = ["environment_impl"] -environment_impl = ["enclave-store"] +environment_impl = ["host-api"] diff --git a/enclave-modules/environment/src/environment_impl.rs b/enclave-modules/environment/src/environment_impl.rs index f8f46751..79cde767 100644 --- a/enclave-modules/environment/src/environment_impl.rs +++ b/enclave-modules/environment/src/environment_impl.rs @@ -1,6 +1,6 @@ use crate::{prelude::*, Env}; use alloc::sync::Arc; -use enclave_store::EnclaveStore; +use host_api::store::new_enclave_store; use light_client::{LightClient, LightClientResolver, MapLightClientRegistry}; use store::{KVStore, TxId}; @@ -24,7 +24,7 @@ impl LightClientResolver for Environment { impl Env for Environment { fn new_store(&self, tx_id: TxId) -> Box { - Box::new(EnclaveStore::new(tx_id)) + new_enclave_store(tx_id) } fn get_lc_registry(&self) -> Arc { diff --git a/enclave-modules/host-api/src/store.rs b/enclave-modules/host-api/src/store.rs index 989dca60..0c371077 100644 --- a/enclave-modules/host-api/src/store.rs +++ b/enclave-modules/host-api/src/store.rs @@ -1,9 +1,41 @@ use crate::prelude::*; use crate::{api::execute_command, Error}; use ocall_commands::{Command, CommandResult, StoreCommand, StoreResult}; -use store::TxId; +use store::cache::CacheKVS; +use store::{KVStore, TxId}; -pub fn get(tx_id: TxId, key: Vec) -> Result>, Error> { +/// The store guarantees that reads a value from the host store only once per key +pub fn new_enclave_store(tx_id: TxId) -> Box { + Box::new(CacheKVS::new(TxStore::new(tx_id))) +} + +/// TxStore is a KVStore implementation that uses the ocall_commands to interact with the +/// host's store. +struct TxStore { + tx_id: TxId, +} + +impl TxStore { + pub fn new(tx_id: TxId) -> Self { + Self { tx_id } + } +} + +impl KVStore for TxStore { + fn get(&self, key: &[u8]) -> Option> { + get(self.tx_id, key.to_vec()).unwrap() + } + + fn set(&mut self, key: Vec, value: Vec) { + set(self.tx_id, key.clone(), value.clone()).unwrap(); + } + + fn remove(&mut self, key: &[u8]) { + remove(self.tx_id, key.to_vec()).unwrap(); + } +} + +fn get(tx_id: TxId, key: Vec) -> Result>, Error> { let cmd = Command::Store(StoreCommand::Get(tx_id, key)); if let CommandResult::Store(StoreResult::Get(v)) = execute_command(cmd)? { Ok(v) @@ -12,7 +44,7 @@ pub fn get(tx_id: TxId, key: Vec) -> Result>, Error> { } } -pub fn set(tx_id: TxId, key: Vec, value: Vec) -> Result<(), Error> { +fn set(tx_id: TxId, key: Vec, value: Vec) -> Result<(), Error> { let cmd = Command::Store(StoreCommand::Set(tx_id, key, value)); if let CommandResult::Store(StoreResult::Set) = execute_command(cmd)? { Ok(()) @@ -21,7 +53,7 @@ pub fn set(tx_id: TxId, key: Vec, value: Vec) -> Result<(), Error> { } } -pub fn remove(tx_id: TxId, key: Vec) -> Result<(), Error> { +fn remove(tx_id: TxId, key: Vec) -> Result<(), Error> { let cmd = Command::Store(StoreCommand::Remove(tx_id, key)); if let CommandResult::Store(StoreResult::Remove) = execute_command(cmd)? { Ok(()) diff --git a/enclave-modules/store/Cargo.toml b/enclave-modules/store/Cargo.toml deleted file mode 100644 index 1788bad8..00000000 --- a/enclave-modules/store/Cargo.toml +++ /dev/null @@ -1,8 +0,0 @@ -[package] -name = "enclave-store" -version = "0.1.0" -edition = "2021" - -[dependencies] -host-api = { path = "../host-api" } -store = { path = "../../modules/store", default-features = false } diff --git a/enclave-modules/store/src/lib.rs b/enclave-modules/store/src/lib.rs deleted file mode 100644 index 6673a24a..00000000 --- a/enclave-modules/store/src/lib.rs +++ /dev/null @@ -1,23 +0,0 @@ -#![no_std] -extern crate alloc; -mod prelude { - pub use core::prelude::v1::*; - - // Re-export according to alloc::prelude::v1 because it is not yet stabilized - // https://doc.rust-lang.org/src/alloc/prelude/v1.rs.html - pub use alloc::borrow::ToOwned; - pub use alloc::boxed::Box; - pub use alloc::string::{String, ToString}; - pub use alloc::vec::Vec; - - pub use alloc::format; - pub use alloc::vec; - - // Those are exported by default in the std prelude in Rust 2021 - pub use core::convert::{TryFrom, TryInto}; - pub use core::iter::FromIterator; -} - -pub use crate::store::EnclaveStore; - -mod store; diff --git a/enclave-modules/store/src/store.rs b/enclave-modules/store/src/store.rs deleted file mode 100644 index ce569ae9..00000000 --- a/enclave-modules/store/src/store.rs +++ /dev/null @@ -1,24 +0,0 @@ -use crate::prelude::*; -use store::{KVStore, TxId}; - -pub struct EnclaveStore { - tx_id: TxId, -} - -impl EnclaveStore { - pub fn new(tx_id: TxId) -> Self { - Self { tx_id } - } -} - -impl KVStore for EnclaveStore { - fn get(&self, key: &[u8]) -> Option> { - host_api::store::get(self.tx_id, key.to_vec()).unwrap() - } - fn set(&mut self, key: Vec, value: Vec) { - host_api::store::set(self.tx_id, key, value).unwrap(); - } - fn remove(&mut self, key: &[u8]) { - host_api::store::remove(self.tx_id, key.to_vec()).unwrap(); - } -} diff --git a/enclave/Cargo.lock b/enclave/Cargo.lock index 2768ff23..70c1f5c8 100644 --- a/enclave/Cargo.lock +++ b/enclave/Cargo.lock @@ -535,7 +535,7 @@ dependencies = [ name = "enclave-environment" version = "0.1.0" dependencies = [ - "enclave-store", + "host-api", "light-client", "store", ] @@ -580,14 +580,6 @@ dependencies = [ "sgx_types", ] -[[package]] -name = "enclave-store" -version = "0.1.0" -dependencies = [ - "host-api", - "store", -] - [[package]] name = "enclave-utils" version = "0.1.0" diff --git a/modules/store/src/cache.rs b/modules/store/src/cache.rs new file mode 100644 index 00000000..709ab83e --- /dev/null +++ b/modules/store/src/cache.rs @@ -0,0 +1,122 @@ +use crate::prelude::*; +use crate::KVStore; +use alloc::collections::BTreeMap; +use core::cell::RefCell; + +/// A key-value store that caches reads in memory +pub struct CacheKVS { + parent: S, + cache: RefCell, Option>>>, +} + +impl CacheKVS { + pub fn new(parent: S) -> Self { + Self { + parent, + cache: RefCell::new(BTreeMap::new()), + } + } +} + +impl KVStore for CacheKVS { + fn set(&mut self, key: Vec, value: Vec) { + self.parent.set(key.clone(), value.clone()); + self.cache.borrow_mut().insert(key, Some(value)); + } + + fn get(&self, key: &[u8]) -> Option> { + let cache = self.cache.borrow(); + let res = cache.get(key); + match res { + Some(Some(v)) => Some(v.clone()), + Some(None) => None, + None => { + drop(cache); + let v = self.parent.get(key); + self.cache.borrow_mut().insert(key.to_vec(), v.clone()); + v + } + } + } + + fn remove(&mut self, key: &[u8]) { + self.parent.remove(key); + self.cache.borrow_mut().insert(key.to_vec(), None); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::KVStore; + use alloc::rc::Rc; + + pub struct MockStore { + db: BTreeMap, Vec>, + } + + impl MockStore { + pub fn new() -> Self { + Self { + db: BTreeMap::new(), + } + } + } + + impl KVStore for MockStore { + fn set(&mut self, key: Vec, value: Vec) { + self.db.insert(key, value); + } + + fn get(&self, key: &[u8]) -> Option> { + self.db.get(key).cloned() + } + + fn remove(&mut self, key: &[u8]) { + self.db.remove(key); + } + } + + #[allow(non_snake_case)] + fn B(s: &str) -> Vec { + s.as_bytes().to_vec() + } + + #[test] + fn test_cache_kvs() { + let mut mock = Rc::new(RefCell::new(MockStore::new())); + mock.set(B("k1"), B("v1")); + + let mut cache = CacheKVS::new(mock.clone()); + assert_eq!(cache.get(&B("k1")), Some(B("v1"))); + + cache.set(B("k1"), B("v2")); + assert_eq!(cache.get(&B("k1")), Some(B("v2"))); + assert_eq!(mock.get(&B("k1")), Some(B("v2"))); + + mock.set(B("k1"), B("v3")); + assert_eq!(mock.get(&B("k1")), Some(B("v3"))); + assert_eq!(cache.get(&B("k1")), Some(B("v2"))); + + mock.remove(&B("k1")); + assert_eq!(mock.get(&B("k1")), None); + assert_eq!(cache.get(&B("k1")), Some(B("v2"))); + + mock.set(B("k2"), B("v4")); + assert_eq!(cache.get(&B("k2")), Some(B("v4"))); + + assert_eq!(cache.get(&B("k3")), None); + mock.set(B("k3"), B("v5")); + assert_eq!(mock.get(&B("k3")), Some(B("v5"))); + assert_eq!(cache.get(&B("k3")), None); + cache.set(B("k3"), B("v6")); + assert_eq!(cache.get(&B("k3")), Some(B("v6"))); + + cache.remove(&B("k4")); + mock.set(B("k4"), B("v7")); + assert_eq!(cache.get(&B("k4")), None); + assert_eq!(mock.get(&B("k4")), Some(B("v7"))); + cache.set(B("k4"), B("v8")); + assert_eq!(cache.get(&B("k4")), Some(B("v8"))); + } +} diff --git a/modules/store/src/lib.rs b/modules/store/src/lib.rs index a80bb564..e6217bd4 100644 --- a/modules/store/src/lib.rs +++ b/modules/store/src/lib.rs @@ -22,6 +22,7 @@ mod prelude { pub use crate::errors::{Error, Result}; pub use crate::store::{KVStore, TxId}; +pub mod cache; mod errors; #[cfg(feature = "std")] pub mod host;