Merge pull request #92 from bluele/fix-keymanager

Fix keymanager to be thread-safe

Signed-off-by: Jun Kimura <[email protected]>

Author: bluele

modules/enclave-api/src/enclave.rs

@@ -9,15 +9,15 @@ use store::host::{HostStore, IntoCommitStore};
 use store::transaction::{CommitStore, CreatedTx, UpdateKey};
 
 /// `Enclave` keeps an enclave id and reference to the host environement
-pub struct Enclave<S> {
+pub struct Enclave<S: CommitStore> {
     pub(crate) path: PathBuf,
     pub(crate) key_manager: EnclaveKeyManager,
     pub(crate) store: Arc<RwLock<HostStore>>,
     pub(crate) sgx_enclave: SgxEnclave,
     _marker: PhantomData<S>,
 }
 
-impl<S> Enclave<S> {
+impl<S: CommitStore> Enclave<S> {
     pub fn new(
         path: impl Into<PathBuf>,
         key_manager: EnclaveKeyManager,
@@ -49,7 +49,7 @@ impl<S> Enclave<S> {
 }
 
 /// `EnclaveInfo` is an accessor to enclave information
-pub trait EnclaveInfo {
+pub trait EnclaveInfo: Sync + Send {
     /// `get_eid` returns the enclave id
     fn get_eid(&self) -> sgx_enclave_id_t;
     /// `metadata` returns the metadata of the enclave
@@ -58,7 +58,7 @@ pub trait EnclaveInfo {
     fn get_key_manager(&self) -> &EnclaveKeyManager;
 }
 
-impl<S> EnclaveInfo for Enclave<S> {
+impl<S: CommitStore> EnclaveInfo for Enclave<S> {
     /// `get_eid` returns the enclave id
     fn get_eid(&self) -> sgx_enclave_id_t {
         self.sgx_enclave.geteid()

modules/keymanager/src/errors.rs

@@ -36,7 +36,15 @@ define_error! {
 
         Rusqlite
         [TraceError<rusqlite::Error>]
-        |_| { "rusqlite error" }
+        |_| { "rusqlite error" },
+
+        MutexLock
+        {
+            descr: String
+        }
+        |e| {
+            format_args!("mutex lock error: descr={}", e.descr)
+        }
     }
 }
 

modules/keymanager/src/lib.rs

@@ -7,37 +7,42 @@ use lcp_types::{Mrenclave, Time};
 use log::*;
 use rusqlite::{params, types::Type, Connection};
 use serde::{Deserialize, Serialize};
+use std::sync::Mutex;
 use std::{ops::Deref, path::Path, time::Duration};
 
 pub static KEY_MANAGER_DB: &str = "km.sqlite";
 
 pub struct EnclaveKeyManager {
-    conn: Connection,
+    conn: Mutex<Connection>,
 }
 
 impl EnclaveKeyManager {
     pub fn new(home_dir: &Path) -> Result<Self, Error> {
         let km_db = home_dir.join(KEY_MANAGER_DB);
         let db_exists = km_db.exists();
-        let conn = Connection::open(&km_db)?;
+        let conn = Mutex::new(Connection::open(&km_db)?);
         let this = Self { conn };
         if !db_exists {
-            this.setup()?;
+            this.init_db()?;
             info!("initialized Key Manager: {:?}", km_db);
         }
         Ok(this)
     }
 
     #[cfg(test)]
     pub fn new_in_memory() -> Result<Self, Error> {
-        let conn = Connection::open_in_memory()?;
+        let conn = Mutex::new(Connection::open_in_memory()?);
         let this = Self { conn };
-        this.setup()?;
+        this.init_db()?;
         Ok(this)
     }
 
-    fn setup(&self) -> Result<(), Error> {
-        self.conn.execute_batch(
+    fn init_db(&self) -> Result<(), Error> {
+        let conn = self
+            .conn
+            .lock()
+            .map_err(|e| Error::mutex_lock(e.to_string()))?;
+        conn.execute_batch(
             r#"
             BEGIN;
             CREATE TABLE enclave_keys (
@@ -61,7 +66,11 @@ impl EnclaveKeyManager {
 
     /// Load a sealed enclave key by address
     pub fn load(&self, address: Address) -> Result<SealedEnclaveKeyInfo, Error> {
-        let mut stmt = self.conn.prepare(
+        let conn = self
+            .conn
+            .lock()
+            .map_err(|e| Error::mutex_lock(e.to_string()))?;
+        let mut stmt = conn.prepare(
             "SELECT ek_sealed, mrenclave, avr, signature, signing_cert FROM enclave_keys WHERE ek_address = ?1",
         )?;
         let key_info = stmt.query_row(params![address.to_hex_string()], |row| {
@@ -98,7 +107,11 @@ impl EnclaveKeyManager {
         sealed_ek: SealedEnclaveKey,
         mrenclave: Mrenclave,
     ) -> Result<(), Error> {
-        let mut stmt = self.conn.prepare(
+        let conn = self
+            .conn
+            .lock()
+            .map_err(|e| Error::mutex_lock(e.to_string()))?;
+        let mut stmt = conn.prepare(
             "INSERT INTO enclave_keys (ek_address, ek_sealed, mrenclave) VALUES (?1, ?2, ?3)",
         )?;
         let _ = stmt.execute(params![
@@ -115,9 +128,13 @@ impl EnclaveKeyManager {
         address: Address,
         avr: EndorsedAttestationVerificationReport,
     ) -> Result<(), Error> {
+        let conn = self
+            .conn
+            .lock()
+            .map_err(|e| Error::mutex_lock(e.to_string()))?;
         let attested_at = avr.get_avr()?.attestation_time()?;
         // update avr and attested_at and signature and sigining_cert
-        let mut stmt = self.conn.prepare(
+        let mut stmt = conn.prepare(
             "UPDATE enclave_keys SET avr = ?1, attested_at = ?2, signature = ?3, signing_cert = ?4 WHERE ek_address = ?5",
         )?;
         stmt.execute(params![
@@ -132,7 +149,11 @@ impl EnclaveKeyManager {
 
     /// Returns a list of available enclave keys
     pub fn available_keys(&self, mrenclave: Mrenclave) -> Result<Vec<SealedEnclaveKeyInfo>, Error> {
-        let mut stmt = self.conn.prepare(
+        let conn = self
+            .conn
+            .lock()
+            .map_err(|e| Error::mutex_lock(e.to_string()))?;
+        let mut stmt = conn.prepare(
             r#"
             SELECT ek_address, ek_sealed, mrenclave, avr, signature, signing_cert
             FROM enclave_keys
@@ -164,7 +185,11 @@ impl EnclaveKeyManager {
 
     /// Returns a list of all enclave keys
     pub fn all_keys(&self) -> Result<Vec<SealedEnclaveKeyInfo>, Error> {
-        let mut stmt = self.conn.prepare(
+        let conn = self
+            .conn
+            .lock()
+            .map_err(|e| Error::mutex_lock(e.to_string()))?;
+        let mut stmt = conn.prepare(
             "SELECT ek_address, ek_sealed, mrenclave, avr, signature, signing_cert FROM enclave_keys ORDER BY updated_at DESC",
         )?;
         let key_infos = stmt
@@ -200,10 +225,12 @@ impl EnclaveKeyManager {
 
     /// Prune keys after the expiration time(secs) from the attestation time.
     pub fn prune(&self, expiration_time: u64) -> Result<usize, Error> {
-        let expired = (Time::now() - Duration::from_secs(expiration_time))?;
-        let mut stmt = self
+        let conn = self
             .conn
-            .prepare("DELETE FROM enclave_keys WHERE attested_at <= ?1")?;
+            .lock()
+            .map_err(|e| Error::mutex_lock(e.to_string()))?;
+        let expired = (Time::now() - Duration::from_secs(expiration_time))?;
+        let mut stmt = conn.prepare("DELETE FROM enclave_keys WHERE attested_at <= ?1")?;
         let count = stmt.execute(params![expired.as_unix_timestamp_secs()])?;
         Ok(count)
     }

modules/service/src/service.rs

@@ -33,19 +33,6 @@ where
     }
 }
 
-unsafe impl<E, S> Send for AppService<E, S>
-where
-    S: CommitStore + 'static,
-    E: EnclaveProtoAPI<S> + 'static,
-{
-}
-unsafe impl<E, S> Sync for AppService<E, S>
-where
-    S: CommitStore + 'static,
-    E: EnclaveProtoAPI<S> + 'static,
-{
-}
-
 impl<E, S> AppService<E, S>
 where
     S: CommitStore + 'static,

modules/store/src/memory.rs

@@ -3,17 +3,72 @@ use crate::store::TxId;
 use crate::transaction::{CommitStore, CreatedTx, Tx, TxAccessor};
 use crate::{KVStore, Result};
 use std::collections::HashMap;
+use std::sync::Mutex;
 
 // MemStore is only available for testing purposes
 #[derive(Default, Debug)]
-pub struct MemStore {
+pub struct MemStore(Mutex<InnerMemStore>);
+
+impl KVStore for MemStore {
+    fn get(&self, key: &[u8]) -> Option<Vec<u8>> {
+        self.0.lock().unwrap().get(key)
+    }
+
+    fn set(&mut self, key: Vec<u8>, value: Vec<u8>) {
+        self.0.lock().unwrap().set(key, value)
+    }
+
+    fn remove(&mut self, key: &[u8]) {
+        self.0.lock().unwrap().remove(key)
+    }
+}
+
+impl TxAccessor for MemStore {
+    fn run_in_tx<T>(&self, tx_id: TxId, f: impl FnOnce(&dyn KVStore) -> T) -> Result<T> {
+        self.0.lock().unwrap().run_in_tx(tx_id, f)
+    }
+
+    fn run_in_mut_tx<T>(
+        &mut self,
+        tx_id: TxId,
+        f: impl FnOnce(&mut dyn KVStore) -> T,
+    ) -> Result<T> {
+        self.0.lock().unwrap().run_in_mut_tx(tx_id, f)
+    }
+}
+
+impl CommitStore for MemStore {
+    type Tx = MemTx;
+
+    fn create_transaction(
+        &mut self,
+        _update_key: Option<crate::transaction::UpdateKey>,
+    ) -> Result<Self::Tx> {
+        self.0.lock().unwrap().create_transaction(_update_key)
+    }
+
+    fn begin(&mut self, tx: &<Self::Tx as CreatedTx>::PreparedTx) -> Result<()> {
+        self.0.lock().unwrap().begin(tx)
+    }
+
+    fn commit(&mut self, tx: <Self::Tx as CreatedTx>::PreparedTx) -> Result<()> {
+        self.0.lock().unwrap().commit(tx)
+    }
+
+    fn rollback(&mut self, tx: <Self::Tx as CreatedTx>::PreparedTx) {
+        self.0.lock().unwrap().rollback(tx)
+    }
+}
+
+#[derive(Default, Debug)]
+pub struct InnerMemStore {
     running_tx_exists: bool,
     latest_tx_id: TxId,
     uncommitted_data: HashMap<Vec<u8>, Option<Vec<u8>>>,
     committed_data: HashMap<Vec<u8>, Vec<u8>>,
 }
 
-impl KVStore for MemStore {
+impl KVStore for InnerMemStore {
     fn get(&self, key: &[u8]) -> Option<Vec<u8>> {
         if self.running_tx_exists {
             match self.uncommitted_data.get(key) {
@@ -42,7 +97,7 @@ impl KVStore for MemStore {
     }
 }
 
-impl TxAccessor for MemStore {
+impl TxAccessor for InnerMemStore {
     fn run_in_tx<T>(&self, _tx_id: TxId, f: impl FnOnce(&dyn KVStore) -> T) -> Result<T> {
         Ok(f(self))
     }
@@ -56,7 +111,7 @@ impl TxAccessor for MemStore {
     }
 }
 
-impl CommitStore for MemStore {
+impl CommitStore for InnerMemStore {
     type Tx = MemTx;
 
     fn create_transaction(

modules/store/src/transaction.rs

@@ -19,7 +19,7 @@ pub trait CreatedTx: Tx {
 }
 
 /// `CommitStore` is a store that supports transactions
-pub trait CommitStore {
+pub trait CommitStore: Sync + Send {
     type Tx: CreatedTx;
 
     /// `create_transaction` creates a transaction with a given `update_key`