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`.

Author: Dillon George

akd/src/directory.rs

@@ -65,11 +65,12 @@ 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(_))) = 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.clone())).await?;
+        } else {
+            let _res = azks?;
         }
 
         Ok(Directory {
@@ -674,6 +675,7 @@ where
                 .get::<Azks>(&crate::append_only_zks::DEFAULT_AZKS_KEY)
                 .await?
         };
+
         match got {
             DbRecord::Azks(azks) => Ok(azks),
             _ => {

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,39 @@ 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())));
+    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
+    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());
+
+    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> {