Auditing improvements

akd/src/auditor.rs

@@ -55,47 +55,70 @@ pub async fn audit_verify<TC: Configuration>(
 }
 
 /// Helper for audit, verifies an append-only proof.
+///
+/// This function first creates a new AZKS instance with the unchanged nodes from the proof,
+/// then it verifies the start hash against the root hash of this AZKS instance.
+/// Next, it creates another AZKS instance with the unchanged nodes and inserted nodes,
+/// and verifies the end hash against the root hash of this second AZKS instance.
 #[cfg_attr(feature = "tracing_instrument", tracing::instrument(skip_all))]
 pub async fn verify_consecutive_append_only<TC: Configuration>(
     proof: &SingleAppendOnlyProof,
     start_hash: Digest,
     end_hash: Digest,
     end_epoch: u64,
 ) -> Result<(), AkdError> {
-    let db = AsyncInMemoryDatabase::new();
-    let manager = StorageManager::new_no_cache(db);
+    let manager1 = StorageManager::new_no_cache(
+        AsyncInMemoryDatabase::new_with_remove_child_nodes_on_insertion(),
+    );
+    let mut azks1 = Azks::new::<TC, _>(&manager1).await?;
+    azks1
+        .batch_insert_nodes::<TC, _>(
+            &manager1,
+            proof.unchanged_nodes.clone(),
+            InsertMode::Auditor,
+            AzksParallelismConfig::default(),
+        )
+        .await?;
+    let computed_start_root_hash: Digest = azks1.get_root_hash::<TC, _>(&manager1).await?;
+    if computed_start_root_hash != start_hash {
+        return Err(AkdError::AzksErr(AzksError::VerifyAppendOnlyProof(
+            format!(
+                "Start hash {} does not match computed root hash {}",
+                hex::encode(start_hash),
+                hex::encode(computed_start_root_hash)
+            ),
+        )));
+    }
 
-    let mut azks = Azks::new::<TC, _>(&manager).await?;
-    azks.batch_insert_nodes::<TC, _>(
-        &manager,
-        proof.unchanged_nodes.clone(),
-        InsertMode::Auditor,
-        AzksParallelismConfig::default(),
-    )
-    .await?;
-    let computed_start_root_hash: Digest = azks.get_root_hash::<TC, _>(&manager).await?;
-    let mut verified = computed_start_root_hash == start_hash;
-    azks.latest_epoch = end_epoch - 1;
-    let updated_inserted = proof
-        .inserted
-        .iter()
-        .map(|x| {
-            let mut y = *x;
-            y.value = AzksValue(TC::hash_leaf_with_commitment(x.value, end_epoch).0);
-            y
-        })
-        .collect();
-    azks.batch_insert_nodes::<TC, _>(
-        &manager,
-        updated_inserted,
-        InsertMode::Auditor,
-        AzksParallelismConfig::default(),
-    )
-    .await?;
-    let computed_end_root_hash: Digest = azks.get_root_hash::<TC, _>(&manager).await?;
-    verified = verified && (computed_end_root_hash == end_hash);
-    if !verified {
-        return Err(AkdError::AzksErr(AzksError::VerifyAppendOnlyProof));
+    let manager2 = StorageManager::new_no_cache(
+        AsyncInMemoryDatabase::new_with_remove_child_nodes_on_insertion(),
+    );
+    let mut azks2 = Azks::new::<TC, _>(&manager2).await?;
+    azks2.latest_epoch = end_epoch - 1;
+    let mut unchanged_with_inserted_nodes = proof.unchanged_nodes.clone();
+    unchanged_with_inserted_nodes.extend(proof.inserted.iter().map(|x| {
+        let mut y = *x;
+        y.value = AzksValue(TC::hash_leaf_with_commitment(x.value, end_epoch).0);
+        y
+    }));
+
+    azks2
+        .batch_insert_nodes::<TC, _>(
+            &manager2,
+            unchanged_with_inserted_nodes,
+            InsertMode::Auditor,
+            AzksParallelismConfig::default(),
+        )
+        .await?;
+    let computed_end_root_hash: Digest = azks2.get_root_hash::<TC, _>(&manager2).await?;
+    if computed_end_root_hash != end_hash {
+        return Err(AkdError::AzksErr(AzksError::VerifyAppendOnlyProof(
+            format!(
+                "End hash {} does not match computed root hash {}",
+                hex::encode(end_hash),
+                hex::encode(computed_end_root_hash)
+            ),
+        )));
     }
     Ok(())
 }

akd/src/errors.rs

@@ -195,7 +195,7 @@ pub enum AzksError {
     /// Membership proof did not verify
     VerifyMembershipProof(String),
     /// Append-only proof did not verify
-    VerifyAppendOnlyProof,
+    VerifyAppendOnlyProof(String),
     /// Thrown when a place where an epoch is needed wasn't provided one.
     NoEpochGiven,
 }
@@ -208,8 +208,8 @@ impl fmt::Display for AzksError {
             Self::VerifyMembershipProof(error_string) => {
                 write!(f, "{error_string}")
             }
-            Self::VerifyAppendOnlyProof => {
-                write!(f, "Append only proof did not verify!")
+            Self::VerifyAppendOnlyProof(error_string) => {
+                write!(f, "Append only proof did not verify: {error_string}")
             }
             Self::NoEpochGiven => {
                 write!(f, "An epoch was required but not supplied")

akd/src/storage/memory.rs

@@ -14,6 +14,7 @@ use crate::storage::types::{
     DbRecord, KeyData, StorageType, ValueState, ValueStateKey, ValueStateRetrievalFlag,
 };
 use crate::storage::{Database, Storable, StorageUtil};
+use crate::tree_node::{NodeKey, TreeNodeWithPreviousValue};
 use crate::{AkdLabel, AkdValue};
 use async_trait::async_trait;
 use dashmap::DashMap;
@@ -30,15 +31,42 @@ type UserValueMap = HashMap<Epoch, ValueState>;
 pub struct AsyncInMemoryDatabase {
     db: Arc<DashMap<Vec<u8>, DbRecord>>,
     user_info: Arc<DashMap<Vec<u8>, UserValueMap>>,
+    /// This flag is used to determine whether the database will automatically
+    /// (and aggressively) remove entries corresponding to left and right
+    /// children of a tree node when the node is inserted. The purpose behind this
+    /// is to reduce the size of the in-memory database by culling the child enries
+    /// once the parent node's hash has been calculated. The primary use case for this
+    /// is to improve auditing memory usage and time (since during auditing, we no longer
+    /// care about the child node hashes once its parent has been computed). Note that this
+    /// technique takes advantage of the way batch insertion of nodes into the tree works,
+    /// since we always process all of the children of a particular subtree before processing
+    /// the root of that subtree.
+    remove_child_nodes_on_insertion: bool,
 }
 
 unsafe impl Send for AsyncInMemoryDatabase {}
 unsafe impl Sync for AsyncInMemoryDatabase {}
 
 impl AsyncInMemoryDatabase {
+    /// Returns the size of the in-memory database
+    pub fn size_of_db(&self) -> usize {
+        self.db.len()
+    }
+
     /// Creates a new in memory db
     pub fn new() -> Self {
-        Self::default()
+        Self {
+            remove_child_nodes_on_insertion: false,
+            ..Self::default()
+        }
+    }
+
+    /// Creates a new in memory db with the flag set to remove child nodes on insertion
+    pub fn new_with_remove_child_nodes_on_insertion() -> Self {
+        Self {
+            remove_child_nodes_on_insertion: true,
+            ..Self::default()
+        }
     }
 
     #[cfg(test)]
@@ -103,6 +131,22 @@ impl Database for AsyncInMemoryDatabase {
                     }
                 }
             } else {
+                if self.remove_child_nodes_on_insertion {
+                    if let DbRecord::TreeNode(node) = record.clone() {
+                        if let Some(left_child) = node.latest_node.left_child {
+                            self.db
+                                .remove(&TreeNodeWithPreviousValue::get_full_binary_key_id(
+                                    &NodeKey(left_child),
+                                ));
+                        }
+                        if let Some(right_child) = node.latest_node.right_child {
+                            self.db
+                                .remove(&TreeNodeWithPreviousValue::get_full_binary_key_id(
+                                    &NodeKey(right_child),
+                                ));
+                        }
+                    }
+                }
                 self.db.insert(record.get_full_binary_id(), record);
             }
         }

akd_core/Cargo.toml

@@ -58,7 +58,7 @@ paste = { version = "1", optional = true }
 bincode = "1"
 itertools = "0.13"
 proptest = "1"
-proptest-derive = "0.4"
+proptest-derive = "0.6"
 rand = "0.8"
 serde = { version = "1", features = ["derive"] }
 criterion = "0.5"

examples/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "examples"
-version = "0.12.0-pre.11"
+version = "0.12.0-pre.12"
 authors = ["akd contributors"]
 license = "MIT OR Apache-2.0"
 edition = "2021"
@@ -20,6 +20,7 @@ runtime_metrics = []
 [dependencies]
 anyhow = "1"
 async-trait = "0.1"
+bytesize = "1"
 colored = "2"
 clap = { version = "4", features = ["derive"] }
 dialoguer = "0.11"

examples/README.md

@@ -21,6 +21,11 @@ You can also automatically audit the latest epoch with the `-l` parameter (for "
 ```
 cargo run -p examples --release -- whatsapp-kt-auditor -l
 ```
+or if you want to audit a specific epoch:
+```
+cargo run -p examples --release -- whatsapp-kt-auditor -e 42
+```
+
 
 ### MySQL Demo
 

examples/src/whatsapp_kt_auditor/auditor.rs

@@ -54,11 +54,7 @@ pub(crate) async fn audit_epoch(blob: akd::local_auditing::AuditBlob) -> Result<
     )
     .await
     {
-        bail!(
-            "Audit proof for epoch {} failed to verify with error: {}",
-            end_epoch,
-            akd_error
-        )
+        bail!("Audit proof for epoch {end_epoch} failed to verify with error: {akd_error}")
     } else {
         // verification passed, generate the appropriate QR code
         Ok(format!(
@@ -99,6 +95,22 @@ pub(crate) fn display_audit_proofs_info(info: &mut [EpochSummary]) -> Result<Str
     ))
 }
 
+pub(crate) async fn get_proof_from_epoch(url: &str, epoch: u64) -> Result<EpochSummary> {
+    let params: Vec<(String, String)> = vec![
+        ("list-type".to_string(), "2".to_string()),
+        ("prefix".to_string(), format!("{epoch}/")),
+    ];
+
+    let (keys, truncated_result) = get_xml(url, &params).await.unwrap();
+    if truncated_result || keys.len() > 1 {
+        bail!("Found multiple matches for epoch {epoch}, which is unexpected. Bailing...");
+    }
+    if keys.is_empty() {
+        bail!("Could not find epoch {epoch}");
+    }
+    Ok(keys[0].clone())
+}
+
 pub(crate) async fn list_proofs(url: &str) -> Result<Vec<EpochSummary>> {
     let mut results = vec![];
     let mut is_truncated = true;