Fix Key History Dirty Read Bug

akd/src/directory.rs

@@ -480,7 +480,7 @@ where
         }
 
         let mut start_version = user_data[0].version;
-        let mut end_version = 0;
+        let mut end_version = user_data[0].version;
         for user_state in &user_data {
             // Ignore states in storage that are ahead of current directory epoch
             if user_state.epoch <= current_epoch {
@@ -489,9 +489,10 @@ where
             }
         }
 
-        if start_version == 0 {
+        if start_version == 0 || end_version == 0 {
             return Err(AkdError::Directory(DirectoryError::InvalidVersion(
-                "Computed start version for the key history should be non-zero".to_string(),
+                "Computed start and end versions for the key history should be non-zero"
+                    .to_string(),
             )));
         }
 
@@ -1033,7 +1034,7 @@ impl<TC: Configuration, S: Database + 'static, V: VRFKeyStorage> Directory<TC, S
             return Ok(EpochHash(current_epoch, root_hash));
         }
 
-        if let false = self.storage.begin_transaction() {
+        if !self.storage.begin_transaction() {
             error!("Transaction is already active");
             return Err(AkdError::Storage(StorageError::Transaction(
                 "Transaction is already active".to_string(),

akd/src/storage/tests.rs

@@ -26,21 +26,6 @@ type Azks = crate::append_only_zks::Azks;
 type TreeNode = crate::tree_node::TreeNode;
 type PvTreeNode = crate::tree_node::TreeNodeWithPreviousValue;
 
-// *** Tests *** //
-
-#[cfg(test)]
-mod memory_storage_tests {
-    use crate::storage::memory::AsyncInMemoryDatabase;
-    use serial_test::serial;
-
-    #[tokio::test]
-    #[serial]
-    async fn test_in_memory_db() {
-        let db = AsyncInMemoryDatabase::new();
-        crate::storage::tests::run_test_cases_for_storage_impl(db).await;
-    }
-}
-
 // *** Run the test cases for a given data-layer impl *** //
 /// Run the storage-layer test suite for a given storage implementation.
 /// This is public because it can be used by other implemented storage layers
@@ -688,3 +673,18 @@ async fn test_tombstoning_data<S: Database>(
 
     Ok(())
 }
+
+// *** Tests *** //
+
+#[cfg(test)]
+mod memory_storage_tests {
+    use crate::storage::memory::AsyncInMemoryDatabase;
+    use serial_test::serial;
+
+    #[tokio::test]
+    #[serial]
+    async fn test_in_memory_db() {
+        let db = AsyncInMemoryDatabase::new();
+        crate::storage::tests::run_test_cases_for_storage_impl(db).await;
+    }
+}

akd/src/tests/test_core_protocol.rs

@@ -303,7 +303,7 @@ async fn test_simple_key_history<TC: Configuration>() -> Result<(), AkdError> {
         borked_proof,
         HistoryVerificationParams::default(),
     );
-    assert!(matches!(result, Err(_)), "{}", "{result:?}");
+    assert!(result.is_err(), "{}", "{result:?}");
 
     Ok(())
 }
@@ -662,15 +662,15 @@ async fn test_simple_audit<TC: Configuration>() -> Result<(), AkdError> {
 
     // The audit should be of more than 1 epoch
     let invalid_audit = akd.audit(3, 3).await;
-    assert!(matches!(invalid_audit, Err(_)));
+    assert!(invalid_audit.is_err());
 
     // The audit epochs must be increasing
     let invalid_audit = akd.audit(3, 2).await;
-    assert!(matches!(invalid_audit, Err(_)));
+    assert!(invalid_audit.is_err());
 
     // The audit should throw an error when queried for an epoch which hasn't yet taken place!
     let invalid_audit = akd.audit(6, 7).await;
-    assert!(matches!(invalid_audit, Err(_)));
+    assert!(invalid_audit.is_err());
 
     Ok(())
 }
@@ -781,7 +781,7 @@ async fn test_tombstoned_key_history<TC: Configuration>() -> Result<(), AkdError
         history_proof.clone(),
         HistoryVerificationParams::default(),
     );
-    assert!(matches!(tombstones, Err(_)));
+    assert!(tombstones.is_err());
 
     // We should be able to verify tombstones assuming the client is accepting
     // of tombstoned states

akd/src/tests/test_errors.rs

@@ -9,17 +9,23 @@
 //! by the API.
 
 use akd_core::configuration::Configuration;
+use std::default::Default;
 
+use crate::storage::types::KeyData;
+use crate::tree_node::TreeNodeWithPreviousValue;
 use crate::{
     auditor::audit_verify,
     client::{key_history_verify, lookup_verify},
     directory::{Directory, PublishCorruption, ReadOnlyDirectory},
     ecvrf::{HardCodedAkdVRF, VRFKeyStorage},
     errors::{AkdError, DirectoryError, StorageError},
-    storage::{manager::StorageManager, memory::AsyncInMemoryDatabase, types::DbRecord, Database},
+    storage::{
+        manager::StorageManager, memory::AsyncInMemoryDatabase, types::DbRecord, types::ValueState,
+        Database,
+    },
     test_config,
     tests::{setup_mocked_db, MockLocalDatabase},
-    AkdLabel, AkdValue, Azks, EpochHash, HistoryParams, HistoryVerificationParams,
+    AkdLabel, AkdValue, Azks, EpochHash, HistoryParams, HistoryVerificationParams, NodeLabel,
 };
 
 // This test is meant to test the function poll_for_azks_change
@@ -114,6 +120,55 @@ async fn test_directory_azks_bootstrapping<TC: Configuration>() -> Result<(), Ak
     Ok(())
 }
 
+// It is possible to perform a "dirty read" when reading states during a key history operation
+// that will result in an epoch from the dirty read being higher than the aZKS epoch. In such an
+// event, we ignore value states that are part of the dirty read. This test ensures that we do not
+// inadvertently panic when inspecting marker versions due to "start version" and "end version"
+// invariants being violated.
+test_config!(test_key_history_dirty_reads);
+async fn test_key_history_dirty_reads<TC: Configuration>() -> Result<(), AkdError> {
+    let committed_epoch = 10;
+    let dirty_epoch = 11;
+
+    let mut mock_db = MockLocalDatabase::default();
+    mock_db.expect_get::<Azks>().returning(move |_| {
+        Ok(DbRecord::Azks(Azks {
+            latest_epoch: committed_epoch,
+            num_nodes: 1,
+        }))
+    });
+    mock_db.expect_get_user_data().returning(move |_| {
+        Ok(KeyData {
+            states: vec![ValueState {
+                value: AkdValue(Vec::new()),
+                version: 2,
+                label: NodeLabel {
+                    label_val: [0u8; 32],
+                    label_len: 32,
+                },
+                epoch: dirty_epoch,
+                username: AkdLabel::from("ferris"),
+            }],
+        })
+    });
+    // We can just return some fake error at this point, as we're not validating
+    // actual history proof functionality.
+    mock_db
+        .expect_get::<TreeNodeWithPreviousValue>()
+        .returning(|_| Err(StorageError::Other("Fake!".to_string())));
+
+    let storage = StorageManager::new_no_cache(mock_db);
+    let vrf = HardCodedAkdVRF {};
+    let akd = Directory::<TC, _, _>::new(storage, vrf).await?;
+
+    // Ensure that we do not panic in this scenario, so we can just ignore the result.
+    let _res = akd
+        .key_history(&AkdLabel::from("ferris"), HistoryParams::MostRecent(1))
+        .await;
+
+    Ok(())
+}
+
 test_config!(test_read_during_publish);
 async fn test_read_during_publish<TC: Configuration>() -> Result<(), AkdError> {
     let db = AsyncInMemoryDatabase::new();
@@ -202,7 +257,7 @@ async fn test_read_during_publish<TC: Configuration>() -> Result<(), AkdError> {
         .unwrap();
 
     let invalid_audit = akd.audit(2, 3).await;
-    assert!(matches!(invalid_audit, Err(_)));
+    assert!(invalid_audit.is_err());
 
     Ok(())
 }
@@ -218,7 +273,7 @@ async fn test_directory_read_only_mode<TC: Configuration>() -> Result<(), AkdErr
     let vrf = HardCodedAkdVRF {};
     // There is no AZKS object in the storage layer, directory construction should fail
     let akd = ReadOnlyDirectory::<TC, _, _>::new(storage, vrf).await;
-    assert!(matches!(akd, Err(_)));
+    assert!(akd.is_err());
 
     Ok(())
 }
@@ -333,7 +388,7 @@ async fn test_key_history_verify_malformed<TC: Configuration>() -> Result<(), Ak
     for _ in 0..100 {
         let mut updates = vec![];
         updates.push((
-            AkdLabel(format!("label").as_bytes().to_vec()),
+            AkdLabel("label".to_string().as_bytes().to_vec()),
             AkdValue::random(&mut rng),
         ));
         akd.publish(updates.clone()).await?;
@@ -342,7 +397,7 @@ async fn test_key_history_verify_malformed<TC: Configuration>() -> Result<(), Ak
     for _ in 0..100 {
         let mut updates = vec![];
         updates.push((
-            AkdLabel(format!("another label").as_bytes().to_vec()),
+            AkdLabel("another label".to_string().as_bytes().to_vec()),
             AkdValue::random(&mut rng),
         ));
         akd.publish(updates.clone()).await?;
@@ -352,7 +407,7 @@ async fn test_key_history_verify_malformed<TC: Configuration>() -> Result<(), Ak
     let EpochHash(current_epoch, root_hash) = akd.get_epoch_hash().await?;
     // Get the VRF public key
     let vrf_pk = akd.get_public_key().await?;
-    let target_label = AkdLabel(format!("label").as_bytes().to_vec());
+    let target_label = AkdLabel("label".to_string().as_bytes().to_vec());
 
     let history_params_5 = HistoryParams::MostRecent(5);