Prevent Any Storage Errors from Clobbering aZKS Entry (#432)

* Prevent Any Storage Errors from Clobbering aZKS Entry

**Current**
--
Prior to this patch, AKD would incorrectly assume that any error during initialization of
a `Directory` should result in an `aZKS` being created and stored in the underlying database
implementation. The problem here is that this behavior happened on every error, even those which
do not definitively indicate that the read for an `aZKS` passed but a record was not found.

For example, if a connection issue occurs when connecting to the storage layer, the current implementation
could wrongly overwrite an existing `aZKS`. In such a scenario, the original read for the `aZKS` fails
during `Directory` initialization, but a subsequent write succeeds. This can result in clobbered records
and unexpected behavior.

**Proposed/New**
--
With this change, we update the initialization logic for the `Directory` type to only create an `aZKS` if
the returned `Err` indicates that the read was successful, but an `aZKS` was not found. In all other instances,
we propagate the `Err` back to the caller instead of swallowing it and wrongly creating an `aZKS`.

* Prepare for 0.12.0-pre.3 Release

---------

Co-authored-by: Dillon George <[email protected]>

CHANGELOG.md

@@ -1,5 +1,8 @@
 # Changelog
 
+## 0.12.0-pre.3 (April 4, 2024)
+* Eliminates a rare bug that can result in an aZKS being overwritten during Directory initialization
+
 ## 0.12.0-pre.2 (March 26, 2024)
 * Updated storage tombstone API params to be more ergonomic
 * Renamed HistoryParams enum variants to highlight caveats

README.md

@@ -22,7 +22,7 @@ Installation
 Add the following line to the dependencies of your `Cargo.toml`:
 
 ```
-akd = "0.12.0-pre.2"
+akd = "0.12.0-pre.3"
 ```
 
 ### Minimum Supported Rust Version
@@ -65,4 +65,4 @@ License
 
 This project is dual-licensed under either the [MIT license](https://github.com/facebook/akd/blob/main/LICENSE-MIT)
 or the [Apache License, Version 2.0](https://github.com/facebook/akd/blob/main/LICENSE-APACHE).
-You may select, at your option, one of the above-listed licenses.
+You may select, at your option, one of the above-listed licenses.

akd/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "akd"
-version = "0.12.0-pre.2"
+version = "0.12.0-pre.3"
 authors = ["akd contributors"]
 description = "An implementation of an auditable key directory"
 license = "MIT OR Apache-2.0"
@@ -51,7 +51,7 @@ default = [
 
 [dependencies]
 ## Required dependencies ##
-akd_core = { version = "0.12.0-pre.2", path = "../akd_core", default-features = false, features = [
+akd_core = { version = "0.12.0-pre.3", path = "../akd_core", default-features = false, features = [
     "vrf",
 ] }
 async-recursion = "1"

akd/src/directory.rs

@@ -65,11 +65,15 @@ where
     pub async fn new(storage: StorageManager<S>, vrf: V) -> Result<Self, AkdError> {
         let azks = Directory::<TC, S, V>::get_azks_from_storage(&storage, false).await;
 
-        if azks.is_err() {
-            // generate a new azks if one is not found
-            let azks = Azks::new::<TC, _>(&storage).await?;
-            // store it
-            storage.set(DbRecord::Azks(azks.clone())).await?;
+        if let Err(AkdError::Storage(StorageError::NotFound(e))) = azks {
+            info!("No aZKS was found in storage: {e}. Creating a new aZKS!");
+            // generate + store a new azks only if one is not found
+            let new_azks = Azks::new::<TC, _>(&storage).await?;
+            storage.set(DbRecord::Azks(new_azks)).await?;
+        } else {
+            // If the value is `Ok`, we drop it since we're not using it below
+            // In all other `Err` cases, we propagate the error to the caller
+            let _res = azks?;
         }
 
         Ok(Directory {

akd/src/storage/manager/mod.rs

@@ -181,7 +181,7 @@ impl<Db: Database> StorageManager<Db> {
         let started = self.transaction.begin_transaction();
 
         // disable the cache cleaning since we're in a write transaction
-        // and will want to keep cache'd objects for the life of the transaction
+        // and will want to keep cached objects for the life of the transaction
         if let Some(cache) = &self.cache {
             cache.disable_clean();
         }

akd/src/tests.rs

@@ -34,6 +34,7 @@ use crate::{
 pub struct LocalDatabase;
 
 unsafe impl Send for LocalDatabase {}
+
 unsafe impl Sync for LocalDatabase {}
 
 mockall::mock! {
@@ -130,6 +131,45 @@ fn setup_mocked_db(db: &mut MockLocalDatabase, test_db: &AsyncInMemoryDatabase)
         });
 }
 
+// A test to ensure that any database error at the time a Directory is created
+// does not automatically attempt to create a new aZKS. Only aZKS not found errors
+// should assume that a successful read happened and no aZKS exists.
+test_config!(test_directory_azks_bootstrapping);
+async fn test_directory_azks_bootstrapping<TC: Configuration>() -> Result<(), AkdError> {
+    let vrf = HardCodedAkdVRF {};
+
+    // Verify that a Storage error results in an error when attempting to create the Directory
+    let mut mock_db = MockLocalDatabase {
+        ..Default::default()
+    };
+    mock_db
+        .expect_get::<Azks>()
+        .returning(|_| Err(StorageError::Connection("Fire!".to_string())));
+    mock_db.expect_set().times(0);
+    let storage = StorageManager::new_no_cache(mock_db);
+
+    let maybe_akd = Directory::<TC, _, _>::new(storage, vrf.clone()).await;
+    assert!(maybe_akd.is_err());
+
+    // Verify that an aZKS not found error results in one being created with the Directory
+    // We're creating an empty directory here, so this is the expected behavior for NotFound
+    let mut mock_db = MockLocalDatabase {
+        ..Default::default()
+    };
+    let test_db = AsyncInMemoryDatabase::new();
+    setup_mocked_db(&mut mock_db, &test_db);
+    let storage = StorageManager::new_no_cache(mock_db);
+
+    let maybe_akd = Directory::<TC, _, _>::new(storage, vrf).await;
+    assert!(maybe_akd.is_ok());
+
+    let akd = maybe_akd.expect("Failed to get create a Directory!");
+    let azks = akd.retrieve_azks().await.expect("Failed to get aZKS!");
+    assert_eq!(0, azks.get_latest_epoch());
+
+    Ok(())
+}
+
 // A simple test to ensure that the empty tree hashes to the correct value
 test_config!(test_empty_tree_root_hash);
 async fn test_empty_tree_root_hash<TC: Configuration>() -> Result<(), AkdError> {

akd_core/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "akd_core"
-version = "0.12.0-pre.2"
+version = "0.12.0-pre.3"
 authors = ["akd contributors"]
 description = "Core utilities for the akd crate"
 license = "MIT OR Apache-2.0"

examples/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "examples"
-version = "0.12.0-pre.2"
+version = "0.12.0-pre.3"
 authors = ["akd contributors"]
 license = "MIT OR Apache-2.0"
 edition = "2021"