diff --git a/db/src/db.rs b/db/src/db.rs index 327ef4ce2..2d02d696a 100644 --- a/db/src/db.rs +++ b/db/src/db.rs @@ -92,19 +92,25 @@ fn make_connection(uri: &Path, maybe_encryption_key: Option<&str>) -> rusqlite:: _ => rusqlite::Connection::open(uri)?, }; - let page_size = 32768; + init_connection(&conn, maybe_encryption_key)?; + + Ok(conn) +} + +pub fn init_connection(conn: &rusqlite::Connection, maybe_encryption_key: Option<&str>) -> rusqlite::Result<()> { let initial_pragmas = if let Some(encryption_key) = maybe_encryption_key { assert!(cfg!(feature = "sqlcipher"), "This function shouldn't be called with a key unless we have sqlcipher support"); + let page_size = 32768; // Important: The `cipher_page_size` cannot be changed without breaking // the ability to open databases that were written when using a // different `cipher_page_size`. Additionally, it (AFAICT) must be a // positive multiple of `page_size`. We use the same value for both here. format!(" - PRAGMA key='{}'; - PRAGMA cipher_page_size={}; - ", escape_string_for_pragma(encryption_key), page_size) + PRAGMA key='{}'; + PRAGMA cipher_page_size={}; + ", escape_string_for_pragma(encryption_key), page_size) } else { String::new() }; @@ -122,9 +128,7 @@ fn make_connection(uri: &Path, maybe_encryption_key: Option<&str>) -> rusqlite:: PRAGMA journal_size_limit=3145728; PRAGMA foreign_keys=ON; PRAGMA temp_store=2; - ", initial_pragmas))?; - - Ok(conn) + ", initial_pragmas)) } pub fn new_connection(uri: T) -> rusqlite::Result where T: AsRef { diff --git a/db/src/lib.rs b/db/src/lib.rs index c7cf40eb6..607ac3fbd 100644 --- a/db/src/lib.rs +++ b/db/src/lib.rs @@ -78,6 +78,7 @@ pub use entids::{ pub use db::{ TypedSQLValue, new_connection, + init_connection, }; #[cfg(feature = "sqlcipher")] diff --git a/ffi/src/lib.rs b/ffi/src/lib.rs index 96c35e132..4e231bdd4 100644 --- a/ffi/src/lib.rs +++ b/ffi/src/lib.rs @@ -107,6 +107,7 @@ pub use mentat::{ QueryResults, RelResult, Store, + Stores, Syncable, TypedValue, TxObserver, @@ -212,7 +213,12 @@ impl<'a, 'c> InProgressTransactResult<'a, 'c> { pub extern "C" fn store_open(uri: *const c_char) -> *mut Store { assert_not_null!(uri); let uri = c_char_to_string(uri); - let store = Store::open(&uri).expect("expected a store"); + let store: Store; + if uri.len() == 0 { + store = Stores::open_named_in_memory_store("").expect("expected a store"); + } else { + store = Stores::open_store(&uri).expect("expected a store"); + } Box::into_raw(Box::new(store)) } diff --git a/src/stores.rs b/src/stores.rs index fa120b57a..7b820b3c7 100644 --- a/src/stores.rs +++ b/src/stores.rs @@ -33,6 +33,8 @@ use r2d2_sqlite::{ SqliteConnectionManager, }; +use mentat_db::init_connection; + use conn::{ Conn, }; @@ -59,23 +61,29 @@ struct StoreConnection { } impl StoreConnection { - fn new(path: T) -> Result<(StoreConnection, PooledConnection)> where T: AsRef { + fn new(path: T, maybe_encryption_key: Option<&str>) -> Result<(StoreConnection, PooledConnection)> where T: AsRef { let path = path.as_ref().to_path_buf(); - let manager = SqliteConnectionManager::file(path); - StoreConnection::new_from_manager(manager) + let manager: SqliteConnectionManager; + if path.as_os_str().is_empty() { + manager = SqliteConnectionManager::file("file::memory:?cache=shared"); + } else { + manager = SqliteConnectionManager::file(path); + } + StoreConnection::new_from_manager(manager, maybe_encryption_key) } - fn new_named_in_memory_connection(name: &str) -> Result<(StoreConnection, PooledConnection)> { - let file = format!("{}?mode=memory&cache=shared", name); + fn new_named_in_memory_connection(name: &str, maybe_encryption_key: Option<&str>) -> Result<(StoreConnection, PooledConnection)> { + let file = format!("file::{}?mode=memory&cache=shared", name); let manager = SqliteConnectionManager::file(&file); - StoreConnection::new_from_manager(manager) + StoreConnection::new_from_manager(manager, maybe_encryption_key) } - fn new_from_manager(manager: SqliteConnectionManager) -> Result<(StoreConnection, PooledConnection)> { + fn new_from_manager(manager: SqliteConnectionManager, maybe_encryption_key: Option<&str>) -> Result<(StoreConnection, PooledConnection)> { let pool = Pool::builder() .max_size(15) .build(manager)?; let mut sqlite = pool.get()?; + init_connection(&sqlite, maybe_encryption_key)?; Ok((StoreConnection { conn: Arc::new(Conn::connect(&mut sqlite)?), pool: pool, @@ -83,8 +91,15 @@ impl StoreConnection { } fn store(& mut self) -> Result { - let sqlite = self.pool.get(); - Store::new(self.conn.clone(), sqlite?) + let sqlite = self.pool.get()?; + init_connection(&sqlite, None)?; + Store::new(self.conn.clone(), sqlite) + } + + fn encrypted_store(& mut self, encryption_key: &str) -> Result { + let sqlite = self.pool.get()?; + init_connection(&sqlite, Some(encryption_key))?; + Store::new(self.conn.clone(), sqlite) } fn store_with_connection(& mut self, sqlite: PooledConnection) -> Result { @@ -115,21 +130,34 @@ impl Stores { Stores::singleton().read().unwrap().is_open(&name) } - pub fn open_store< T>(path: T) -> Result where T: AsRef { + pub fn open_store(path: T) -> Result where T: AsRef { let path_ref = path.as_ref(); let name: String = path_ref.to_string_lossy().into(); - Stores::singleton().write().map_err(|_| MentatError::StoresLockPoisoned(name.clone()))?.open(&name, path_ref.to_path_buf()) + Stores::singleton().write().map_err(|_| MentatError::StoresLockPoisoned(name.clone()))?.open(&name, path_ref) } pub fn open_named_in_memory_store(name: &str) -> Result { Stores::singleton().write().map_err(|_| MentatError::StoresLockPoisoned(name.to_string()))?.open(name, "") } + #[cfg(feature = "sqlcipher")] + pub fn open_store_with_key(path: T, encryption_key: &str) -> Result where T: AsRef { + let path_ref = path.as_ref(); + let name: String = path_ref.to_string_lossy().into(); + Stores::singleton().write().map_err(|_| MentatError::StoresLockPoisoned(name.clone()))?.open_with_key(&name, path_ref, encryption_key) + } + pub fn get_store(path: T) -> Result> where T: AsRef { let name: String = path.as_ref().to_string_lossy().into(); Stores::singleton().write().map_err(|_| MentatError::StoresLockPoisoned(name.clone()))?.get(&name) } + #[cfg(feature = "sqlcipher")] + pub fn get_store_with_key(path: T, encryption_key: &str) -> Result> where T: AsRef { + let name: String = path.as_ref().to_string_lossy().into(); + Stores::singleton().write().map_err(|_| MentatError::StoresLockPoisoned(name.clone()))?.get_with_key(&name, encryption_key) + } + pub fn get_named_in_memory_store(name: &str) -> Result> { Stores::singleton().write().map_err(|_| MentatError::StoresLockPoisoned(name.to_string()))?.get(name) } @@ -139,6 +167,12 @@ impl Stores { Stores::singleton().write().map_err(|_| MentatError::StoresLockPoisoned(name.clone()))?.connect(&name) } + #[cfg(feature = "sqlcipher")] + pub fn connect_store_with_key< T>(path: T, encryption_key: &str) -> Result where T: AsRef { + let name: String = path.as_ref().to_string_lossy().into(); + Stores::singleton().write().map_err(|_| MentatError::StoresLockPoisoned(name.clone()))?.connect_with_key(&name, encryption_key) + } + pub fn connect_named_in_memory_store(name: &str) -> Result { Stores::singleton().write().map_err(|_| MentatError::StoresLockPoisoned(name.to_string()))?.connect(name) } @@ -169,7 +203,13 @@ impl Stores { connection.store() }, Entry::Vacant(entry) =>{ - let (mut store_connection, connection) = StoreConnection::new(path)?; + let path = path.as_ref().to_path_buf(); + + let (mut store_connection, connection) = if !name.is_empty() && path.as_os_str().is_empty() { + StoreConnection::new_named_in_memory_connection(name, None)? + } else { + StoreConnection::new(path, None)? + }; let store = store_connection.store_with_connection(connection); entry.insert(store_connection); store @@ -177,14 +217,23 @@ impl Stores { } } - pub fn open_in_memory(&mut self, name: &str) -> Result { + // Open an encrypted store with an existing connection if available, or + // create a new connection if not. + #[cfg(feature = "sqlcipher")] + pub fn open_with_key(&mut self, name: &str, path: T, encryption_key: &str) -> Result where T: AsRef { match self.connections.entry(name.to_string()) { Entry::Occupied(mut entry) => { let connection = entry.get_mut(); connection.store() }, Entry::Vacant(entry) =>{ - let (mut store_connection, connection) = StoreConnection::new_named_in_memory_connection(name)?; + let path = path.as_ref().to_path_buf(); + + let (mut store_connection, connection) = if !name.is_empty() && path.as_os_str().is_empty() { + StoreConnection::new_named_in_memory_connection(name, Some(encryption_key))? + } else { + StoreConnection::new(path, Some(encryption_key))? + }; let store = store_connection.store_with_connection(connection); entry.insert(store_connection); store @@ -200,6 +249,15 @@ impl Stores { .map(|s| Some(s))) } + // Returns an encrypted store with an existing connection to path, if available, or None if a + // store at the provided path has not yet been opened. + #[cfg(feature = "sqlcipher")] + pub fn get_with_key(&mut self, name: &str, encryption_key: &str) -> Result> { + self.connections.get_mut(name) + .map_or(Ok(None), |store_conn| store_conn.encrypted_store(encryption_key) + .map(|s| Some(s))) + } + // Creates a new store on an existing connection with a new rusqlite connection. // Equivalent to forking an existing store. pub fn connect(&mut self, name: &str) -> Result { @@ -208,6 +266,15 @@ impl Stores { .and_then(|store_conn| store_conn.store()) } + // Creates a new store on an existing connection with a new encrypted rusqlite connection. + // Equivalent to forking an existing store. + #[cfg(feature = "sqlcipher")] + pub fn connect_with_key(&mut self, name: &str, encryption_key: &str) -> Result { + self.connections.get_mut(name) + .ok_or(MentatError::StoreNotFound(name.to_string()).into()) + .and_then(|store_conn| store_conn.encrypted_store(encryption_key)) + } + // Drops the weak reference we have stored to an opened store there is no more than // one Store with a reference to the Conn for the provided path. pub fn close(&mut self, name: &str) -> Result<()> { @@ -364,4 +431,13 @@ mod tests { let result = store1.q_once(r#"[:find ?e . :where [?e :foo/baz false]]"#, None).expect("succeeded"); assert!(result.into_scalar().expect("succeeded").is_some()); } + + #[test] + #[cfg(feature = "sqlcipher")] + fn test_open_store_with_key() { + let secret_key = "key"; + let name = "../fixtures/v1encrypted.db"; + let _store = Stores::open_store_with_key(name, secret_key).expect("Expected a store to be opened"); + assert!(Stores::is_store_open(name)); + } }