Fix auditor append-only bypass in batch node insertion

akd/src/append_only_zks.rs

@@ -18,7 +18,7 @@ use crate::tree_node::{
 };
 use crate::Configuration;
 use crate::{
-    errors::{AkdError, DirectoryError, ParallelismError, TreeNodeError},
+    errors::{AkdError, AzksError, DirectoryError, ParallelismError, TreeNodeError},
     storage::{Database, Storable},
     AppendOnlyProof, AzksElement, AzksValue, Digest, Direction, MembershipProof, NodeLabel,
     NonMembershipProof, PrefixOrdering, SiblingProof, SingleAppendOnlyProof, SizeOf, ARITY,
@@ -453,8 +453,27 @@ impl Azks {
         // nodes are located in with respect to the current node and call this
         // function recursively on the left and right child nodes. The current
         // node is updated with the new child nodes.
+        let element_count = azks_element_set.len();
         let (left_azks_element_set, right_azks_element_set) =
             azks_element_set.partition(current_node.label);
+        // `partition` silently discards any element not strictly beneath the
+        // pivot. In Auditor mode the only legitimate drop is the single element
+        // consumed to become this leaf; any other drop loses a committed value
+        // (e.g. an unchanged interior node that is a strict prefix of an inserted
+        // leaf) and breaks the append-only guarantee, so reject it. Scoped to
+        // Auditor mode, the only path that inserts interior nodes alongside
+        // leaves; the directory publish path inserts full-length leaves only.
+        if matches!(insert_mode, InsertMode::Auditor)
+            && left_azks_element_set.len() + right_azks_element_set.len() != element_count
+            && !(current_node.node_type == TreeNodeType::Leaf && element_count == 1)
+        {
+            return Err(AkdError::AzksErr(AzksError::BatchInsertDroppedNode(
+                format!(
+                    "a committed value was dropped while inserting under node {:?}",
+                    current_node.label
+                ),
+            )));
+        }
         let child_parallel_levels =
             parallel_levels.and_then(|x| if x <= 1 { None } else { Some(x - 1) });
 
@@ -1428,6 +1447,90 @@ mod tests {
         Ok(())
     }
 
+    // An `InsertMode::Auditor` node set that forces `partition` to drop a
+    // committed value must be rejected, while collision-free sets still insert.
+    test_config!(test_auditor_insert_rejects_dropped_nodes);
+    async fn test_auditor_insert_rejects_dropped_nodes<TC: Configuration>() -> Result<(), AkdError>
+    {
+        async fn try_auditor_insert<TC: Configuration>(
+            nodes: Vec<AzksElement>,
+        ) -> Result<(), AkdError> {
+            let db = StorageManager::new_no_cache(AsyncInMemoryDatabase::new());
+            let mut azks = Azks::new::<TC, _>(&db).await?;
+            azks.batch_insert_nodes::<TC, _>(
+                &db,
+                nodes,
+                InsertMode::Auditor,
+                AzksParallelismConfig::default(),
+            )
+            .await
+        }
+        let val = |s: &[u8]| AzksValue(TC::hash(s));
+
+        // (a) An interior node (1-bit label) that is a strict prefix of a leaf
+        // (256-bit label): the interior node's committed value is dropped when
+        // both are inserted together -- the reported bypass, one level deeper
+        // than the direct root child.
+        let prefix = try_auditor_insert::<TC>(vec![
+            AzksElement {
+                label: NodeLabel::new([0u8; 32], 1),
+                value: val(b"a"),
+            },
+            AzksElement {
+                label: NodeLabel::new([0u8; 32], 256),
+                value: val(b"b"),
+            },
+        ])
+        .await;
+        assert!(
+            matches!(
+                prefix,
+                Err(AkdError::AzksErr(AzksError::BatchInsertDroppedNode(_)))
+            ),
+            "strict-prefix collision must be rejected, got {prefix:?}"
+        );
+
+        // (b) Two elements sharing a label but committing to different values:
+        // one is dropped at the interior pivot created for their common prefix.
+        let dup = try_auditor_insert::<TC>(vec![
+            AzksElement {
+                label: NodeLabel::new([1u8; 32], 256),
+                value: val(b"a"),
+            },
+            AzksElement {
+                label: NodeLabel::new([1u8; 32], 256),
+                value: val(b"b"),
+            },
+        ])
+        .await;
+        assert!(
+            matches!(
+                dup,
+                Err(AkdError::AzksErr(AzksError::BatchInsertDroppedNode(_)))
+            ),
+            "duplicate-label collision must be rejected, got {dup:?}"
+        );
+
+        // Positive control: two distinct full-length leaves have no prefix
+        // relation, so nothing is dropped and the auditor insert succeeds.
+        let mut b_label = [0u8; 32];
+        b_label[0] = 1;
+        try_auditor_insert::<TC>(vec![
+            AzksElement {
+                label: NodeLabel::new([0u8; 32], 256),
+                value: val(b"a"),
+            },
+            AzksElement {
+                label: NodeLabel::new(b_label, 256),
+                value: val(b"b"),
+            },
+        ])
+        .await
+        .expect("disjoint leaves must insert cleanly in auditor mode");
+
+        Ok(())
+    }
+
     test_config!(test_batch_insert_root_hash);
     async fn test_batch_insert_root_hash<TC: Configuration>() -> Result<(), AkdError> {
         let database = AsyncInMemoryDatabase::new();

akd/src/auditor.rs

@@ -117,3 +117,111 @@ async fn verify_append_only_hash<TC: Configuration>(
     }
     Ok(())
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::client::verify_membership_for_tests_only;
+    use crate::test_config;
+    use crate::{
+        AzksValue, Direction, MembershipProof, NodeLabel, SiblingProof, SingleAppendOnlyProof,
+    };
+
+    // Regression test for the auditor append-only bypass: when an unchanged
+    // interior node's label is a strict prefix of an inserted leaf's label,
+    // `partition` used to silently drop the unchanged node, letting a malicious
+    // server rewrite a label's value (`val1` -> `val2`) while still producing a
+    // valid append-only proof. `audit_verify` must now reject the transition.
+    test_config!(test_auditor_rejects_prefix_collision_value_rewrite);
+    #[allow(non_snake_case)]
+    async fn test_auditor_rejects_prefix_collision_value_rewrite<TC: Configuration>(
+    ) -> Result<(), AkdError> {
+        const AUDIT_EPOCH: u64 = 2;
+        const START_EPOCH: u64 = AUDIT_EPOCH - 1;
+
+        // `shell_label` (1 bit) is a strict prefix of `label` (256 bits). This
+        // is what triggers the drop in `partition` when both are inserted.
+        let shell_label = NodeLabel::new([0u8; 32], 1);
+        let label = NodeLabel::new([0u8; 32], 256);
+
+        let empty_child = AzksElement {
+            label: TC::empty_label(),
+            value: TC::empty_node_hash(),
+        };
+        // Value of a node whose only (left) child is `(child_label, child_val)`.
+        let parent_hash = |child_label: NodeLabel, child_val: AzksValue| {
+            TC::compute_parent_hash_from_children(
+                &child_val,
+                &child_label.value::<TC>(),
+                &empty_child.value,
+                &empty_child.label.value::<TC>(),
+            )
+        };
+
+        let val1 = AzksValue(TC::hash(b"val1"));
+        let val2 = AzksValue(TC::hash(b"val2"));
+        // The auditor re-commits leaves with their corresponding epoch.
+        let leaf_val1 = AzksValue(TC::hash_leaf_with_commitment(val1, START_EPOCH).0);
+        let leaf_val2 = AzksValue(TC::hash_leaf_with_commitment(val2, AUDIT_EPOCH).0);
+
+        // root ->[left] Leaf(shell_label, shell_val), where shell_val commits to
+        // the subtree Interior(shell_label) ->[left] Leaf(label, val1).
+        let shell_val = parent_hash(label, leaf_val1);
+        let root_hash1 = TC::compute_root_hash_from_val(&parent_hash(shell_label, shell_val));
+        // root ->[left] Interior(shell_label) ->[left] Leaf(label, val2).
+        let root_hash2 = TC::compute_root_hash_from_val(&parent_hash(
+            shell_label,
+            parent_hash(label, leaf_val2),
+        ));
+
+        let update_proof = SingleAppendOnlyProof {
+            unchanged_nodes: vec![AzksElement {
+                label: shell_label,
+                value: shell_val,
+            }],
+            inserted: vec![AzksElement { label, value: val2 }],
+        };
+
+        // Sanity: the two membership proofs genuinely disagree on `label`'s
+        // value, so accepting this transition would break append-only-ness.
+        let sibling_path = vec![
+            SiblingProof {
+                label: NodeLabel::root(),
+                siblings: [empty_child],
+                direction: Direction::Left,
+            },
+            SiblingProof {
+                label: shell_label,
+                siblings: [empty_child],
+                direction: Direction::Left,
+            },
+        ];
+        let membership = |hash_val| MembershipProof {
+            label,
+            hash_val,
+            sibling_proofs: sibling_path.clone(),
+        };
+        verify_membership_for_tests_only::<TC>(root_hash1, &membership(leaf_val1)).unwrap();
+        verify_membership_for_tests_only::<TC>(root_hash2, &membership(leaf_val2)).unwrap();
+        assert_ne!(leaf_val1, leaf_val2);
+
+        // The core assertion: the auditor must reject the crafted transition.
+        let result = audit_verify::<TC>(
+            vec![root_hash1, root_hash2],
+            AppendOnlyProof {
+                proofs: vec![update_proof],
+                epochs: vec![START_EPOCH],
+            },
+        )
+        .await;
+        assert!(
+            matches!(
+                result,
+                Err(AkdError::AzksErr(AzksError::BatchInsertDroppedNode(_)))
+            ),
+            "auditor must reject the value-rewrite transition, got {result:?}"
+        );
+
+        Ok(())
+    }
+}

akd/src/errors.rs

@@ -196,6 +196,9 @@ pub enum AzksError {
     VerifyMembershipProof(String),
     /// Append-only proof did not verify
     VerifyAppendOnlyProof(String),
+    /// A batch insertion dropped one or more nodes (a node could not be placed
+    /// strictly beneath its parent). This indicates a malformed node set.
+    BatchInsertDroppedNode(String),
     /// Thrown when a place where an epoch is needed wasn't provided one.
     NoEpochGiven,
 }
@@ -211,6 +214,9 @@ impl fmt::Display for AzksError {
             Self::VerifyAppendOnlyProof(error_string) => {
                 write!(f, "Append only proof did not verify: {error_string}")
             }
+            Self::BatchInsertDroppedNode(error_string) => {
+                write!(f, "Batch insert dropped a node: {error_string}")
+            }
             Self::NoEpochGiven => {
                 write!(f, "An epoch was required but not supplied")
             }