Adding non-existence proofs for directory lookup

akd/src/directory.rs

@@ -465,6 +465,45 @@ where
         }
     }
 
+    /// Generates VRF proofs and non-membership proofs for future marker versions.
+    /// This is a helper method to eliminate code duplication between the non-existent
+    /// user case and the normal key history generation.
+    async fn generate_future_marker_proofs(
+        &self,
+        akd_label: &AkdLabel,
+        current_azks: &Azks,
+        future_marker_versions: &[u64],
+    ) -> Result<(Vec<Vec<u8>>, Vec<crate::NonMembershipProof>), AkdError> {
+        let mut future_marker_vrf_proofs = Vec::with_capacity(future_marker_versions.len());
+        let mut non_existence_of_future_marker_proofs =
+            Vec::with_capacity(future_marker_versions.len());
+
+        for &version in future_marker_versions {
+            let node_label = self
+                .vrf
+                .get_node_label::<TC>(akd_label, VersionFreshness::Fresh, version)
+                .await?;
+
+            let non_membership_proof = current_azks
+                .get_non_membership_proof::<TC, _>(&self.storage, node_label)
+                .await?;
+            non_existence_of_future_marker_proofs.push(non_membership_proof);
+
+            let vrf_proof = self
+                .vrf
+                .get_label_proof::<TC>(akd_label, VersionFreshness::Fresh, version)
+                .await?
+                .to_bytes()
+                .to_vec();
+            future_marker_vrf_proofs.push(vrf_proof);
+        }
+
+        Ok((
+            future_marker_vrf_proofs,
+            non_existence_of_future_marker_proofs,
+        ))
+    }
+
     /// Takes in the current state of the server and a label.
     /// If the label is present in the current state,
     /// this function returns all the values ever associated with it,
@@ -488,7 +527,42 @@ where
 
         let current_azks = self.retrieve_azks().await?;
         let current_epoch = current_azks.get_latest_epoch();
-        let mut user_data = self.storage.get_user_data(akd_label).await?.states;
+
+        let mut user_data = match self.storage.get_user_data(akd_label).await {
+            Ok(data) => data.states,
+            Err(StorageError::NotFound(_)) => {
+                // For non-existent labels, return a non-inclusion proof with future marker versions
+                // from the skiplist up to the current epoch
+                let epoch_index: usize = akd_core::utils::find_max_index_in_skiplist(current_epoch);
+                let future_marker_versions: Vec<u64> =
+                    akd_core::utils::MARKER_VERSION_SKIPLIST[0..=epoch_index].to_vec();
+
+                let (future_marker_vrf_proofs, non_existence_of_future_marker_proofs) = self
+                    .generate_future_marker_proofs(
+                        akd_label,
+                        &current_azks,
+                        &future_marker_versions,
+                    )
+                    .await?;
+
+                let root_hash = EpochHash(
+                    current_epoch,
+                    current_azks.get_root_hash::<TC, _>(&self.storage).await?,
+                );
+
+                return Ok((
+                    HistoryProof {
+                        update_proofs: vec![],
+                        past_marker_vrf_proofs: vec![],
+                        existence_of_past_marker_proofs: vec![],
+                        future_marker_vrf_proofs,
+                        non_existence_of_future_marker_proofs,
+                    },
+                    root_hash,
+                ));
+            }
+            Err(e) => return Err(AkdError::Storage(e)),
+        };
 
         // Ignore states in storage which are ahead of the current directory epoch
         user_data.retain(|vs| vs.epoch <= current_epoch);
@@ -582,27 +656,9 @@ where
             );
         }
 
-        let mut future_marker_vrf_proofs = vec![];
-        let mut non_existence_of_future_marker_proofs = vec![];
-
-        for version in future_marker_versions {
-            let node_label = self
-                .vrf
-                .get_node_label::<TC>(akd_label, VersionFreshness::Fresh, version)
-                .await?;
-            non_existence_of_future_marker_proofs.push(
-                current_azks
-                    .get_non_membership_proof::<TC, _>(&self.storage, node_label)
-                    .await?,
-            );
-            future_marker_vrf_proofs.push(
-                self.vrf
-                    .get_label_proof::<TC>(akd_label, VersionFreshness::Fresh, version)
-                    .await?
-                    .to_bytes()
-                    .to_vec(),
-            );
-        }
+        let (future_marker_vrf_proofs, non_existence_of_future_marker_proofs) = self
+            .generate_future_marker_proofs(akd_label, &current_azks, &future_marker_versions)
+            .await?;
 
         let root_hash = EpochHash(
             current_epoch,

akd/src/storage/manager/mod.rs

@@ -518,7 +518,7 @@ impl<Db: Database> StorageManager<Db> {
 
             let transaction_records = self
                 .transaction
-                .get_users_data(&[username.clone()])
+                .get_users_data(std::slice::from_ref(username))
                 .remove(username)
                 .unwrap_or_default();
             for transaction_record in transaction_records.into_iter() {

akd/src/storage/transaction.rs

@@ -218,7 +218,7 @@ impl Transaction {
         flag: ValueStateRetrievalFlag,
     ) -> Option<ValueState> {
         let intermediate = self
-            .get_users_data(&[username.clone()])
+            .get_users_data(std::slice::from_ref(username))
             .remove(username)
             .unwrap_or_default();
         let out = Self::find_appropriate_item(intermediate, flag);

akd/src/tests/test_core_protocol.rs

@@ -732,6 +732,135 @@ async fn test_simple_lookup_for_small_tree<TC: Configuration>() -> Result<(), Ak
     Ok(())
 }
 
+test_config!(test_key_history_nonexistent_label);
+async fn test_key_history_nonexistent_label<TC: Configuration>() -> Result<(), AkdError> {
+    // Test at different epoch levels to ensure robustness
+    let test_cases = vec![
+        (1, vec![1]),
+        (2, vec![1, 2]),
+        (5, vec![1, 2, 4]),
+        (16, vec![1, 2, 4, 16]),
+        (100, vec![1, 2, 4, 16]),
+        (256, vec![1, 2, 4, 16, 256]),
+        (1000, vec![1, 2, 4, 16, 256]),
+    ];
+
+    for (target_epoch, expected_future_versions) in test_cases {
+        let db = AsyncInMemoryDatabase::new();
+        let storage = StorageManager::new_no_cache(db);
+        let vrf = HardCodedAkdVRF {};
+        let akd =
+            Directory::<TC, _, _>::new(storage, vrf.clone(), AzksParallelismConfig::default())
+                .await?;
+
+        // Publish enough updates to reach the target epoch
+        for epoch in 1..=target_epoch {
+            let user_label = format!("user_epoch_{epoch}");
+            let user_value = format!("value_{epoch}");
+            akd.publish(vec![(
+                AkdLabel::from(user_label.as_str()),
+                AkdValue::from(user_value.as_str()),
+            )])
+            .await?;
+        }
+
+        let current_epoch = akd.retrieve_azks().await?.get_latest_epoch();
+        assert_eq!(
+            current_epoch, target_epoch,
+            "Should be at epoch {target_epoch}"
+        );
+
+        // Try to get history for a label that does not exist
+        let nonexistent_label_str = format!("nonexistent_user_epoch_{target_epoch}");
+        let nonexistent_label = AkdLabel::from(nonexistent_label_str.as_str());
+        let (history_proof, root_hash) = akd
+            .key_history(&nonexistent_label, HistoryParams::default())
+            .await?;
+
+        // The proof should be a non-inclusion proof with empty update_proofs
+        assert_eq!(
+            0,
+            history_proof.update_proofs.len(),
+            "Update proofs should be empty for nonexistent user at epoch {target_epoch}"
+        );
+        assert_eq!(
+            0,
+            history_proof.past_marker_vrf_proofs.len(),
+            "Past marker VRF proofs should be empty for nonexistent user at epoch {target_epoch}"
+        );
+        assert_eq!(0, history_proof.existence_of_past_marker_proofs.len(),
+                   "Past marker existence proofs should be empty for nonexistent user at epoch {target_epoch}");
+
+        // Check the exact expected future marker versions
+        assert_eq!(
+            history_proof.future_marker_vrf_proofs.len(), 
+            expected_future_versions.len(),
+            "Future marker VRF proofs count should match expected at epoch {}: expected {:?}, got {}",
+            target_epoch, expected_future_versions, history_proof.future_marker_vrf_proofs.len()
+        );
+        assert_eq!(
+            history_proof.non_existence_of_future_marker_proofs.len(), 
+            expected_future_versions.len(),
+            "Future marker non-existence proofs count should match expected at epoch {}: expected {:?}, got {}",
+            target_epoch, expected_future_versions, history_proof.non_existence_of_future_marker_proofs.len()
+        );
+
+        // Verify the structure makes sense and passes verification
+        let vrf_pk = akd.get_public_key().await?;
+        let result = crate::client::key_history_verify::<TC>(
+            vrf_pk.as_bytes(),
+            root_hash.hash(),
+            root_hash.epoch(),
+            nonexistent_label.clone(),
+            history_proof.clone(),
+            crate::HistoryVerificationParams::default(),
+        );
+
+        match &result {
+            Ok(verify_results) => {
+                assert_eq!(
+                    0,
+                    verify_results.len(),
+                    "Verification results should be empty for nonexistent user at epoch {}",
+                    target_epoch
+                );
+            }
+            Err(e) => {
+                panic!(
+                    "Key history verification failed at epoch {}: {:?}",
+                    target_epoch, e
+                );
+            }
+        }
+
+        // Let's also verify that we can extract the actual NodeLabels from the VRF proofs
+        // and they match our expectations for the specific versions
+        let mut actual_node_labels = Vec::new();
+        for (i, vrf_proof_bytes) in history_proof.future_marker_vrf_proofs.iter().enumerate() {
+            // Parse VRF proof and extract NodeLabel using the VRF instance
+            if let Ok(proof) = akd_core::ecvrf::Proof::try_from(vrf_proof_bytes.as_slice()) {
+                let node_label = vrf.get_node_label_from_vrf_proof(proof).await;
+                actual_node_labels.push(node_label);
+
+                // Also verify this matches what we'd expect for this specific version
+                let expected_version = expected_future_versions[i];
+                let expected_node_label = vrf
+                    .get_node_label::<TC>(
+                        &nonexistent_label,
+                        akd_core::VersionFreshness::Fresh,
+                        expected_version,
+                    )
+                    .await?;
+
+                assert_eq!(node_label, expected_node_label,
+                          "VRF proof node label should match expected node label for version {expected_version} at epoch {target_epoch}");
+            }
+        }
+    }
+
+    Ok(())
+}
+
 test_config!(test_tombstoned_key_history);
 async fn test_tombstoned_key_history<TC: Configuration>() -> Result<(), AkdError> {
     let db = AsyncInMemoryDatabase::new();

akd_core/src/utils.rs

@@ -12,7 +12,7 @@ use alloc::vec::Vec;
 
 /// This array is: [2, 4, 16, 256, 65536, 2^32] and is used in get_marker_versions() as
 /// an efficiency optimization
-const MARKER_VERSION_SKIPLIST: [u64; 7] = [1, 1 << 1, 1 << 2, 1 << 4, 1 << 8, 1 << 16, 1 << 32];
+pub const MARKER_VERSION_SKIPLIST: [u64; 7] = [1, 1 << 1, 1 << 2, 1 << 4, 1 << 8, 1 << 16, 1 << 32];
 
 /// a list of past marker versions used by history proofs
 pub type PastMarkerVersions = Vec<u64>;
@@ -167,18 +167,18 @@ pub fn get_marker_versions(
     (past_marker_versions, future_marker_versions)
 }
 
-// Given an input u64 and a sorted array of u64s, find the largest index for which
-// the corresponding array element is less than the input.
-//
-// This implementation performs a linear search over the MARKER_VERSION_SKIPLIST array,
-// but since it is sorted, it could be faster to do a binary search. However, given that
-// the array is small, there shouldn't be too much of a difference
-// between a binary search and a linear one. But if this ends up being problematic
-// in the future, it could certainly be optimized.
-//
-// Note that if the input is less than the smallest element of the array, then this
-// function will panic.
-fn find_max_index_in_skiplist(input: u64) -> usize {
+/// Given an input u64 and a sorted array of u64s, find the largest index for which
+/// the corresponding array element is less than the input.
+///
+/// This implementation performs a linear search over the MARKER_VERSION_SKIPLIST array,
+/// but since it is sorted, it could be faster to do a binary search. However, given that
+/// the array is small, there shouldn't be too much of a difference
+/// between a binary search and a linear one. But if this ends up being problematic
+/// in the future, it could certainly be optimized.
+///
+/// Note that if the input is less than the smallest element of the array, then this
+/// function will panic.
+pub fn find_max_index_in_skiplist(input: u64) -> usize {
     if input < MARKER_VERSION_SKIPLIST[0] {
         panic!("find_max_index_in_skiplist called with input less than smallest element of MARKER_VERSION_SKIPLIST");
     }

akd_core/src/verify/history.rs

@@ -67,17 +67,23 @@ impl Default for HistoryVerificationParams {
 
 fn verify_with_history_params(
     current_epoch: u64,
-    akd_label: &AkdLabel,
     proof: &HistoryProof,
     params: HistoryParams,
 ) -> Result<(Vec<u64>, Vec<u64>), VerificationError> {
     let num_proofs = proof.update_proofs.len();
 
-    // Make sure the update proofs are non-empty
+    // If we have no update proofs, this is likely a non-inclusion proof for a non-existent user
     if num_proofs == 0 {
-        return Err(VerificationError::HistoryProof(format!(
-            "No update proofs included in the proof of user {akd_label:?} at epoch {current_epoch:?}!"
-        )));
+        // For non-existent users, we should only have future marker proofs for non-existence
+        // The future marker versions should be elements from the skiplist up to current_epoch
+        let epoch_index: usize = crate::utils::find_max_index_in_skiplist(current_epoch);
+        let expected_future_versions: Vec<u64> =
+            crate::utils::MARKER_VERSION_SKIPLIST[0..=epoch_index].to_vec();
+
+        // TODO: add checks for proof lengths
+
+        // Return empty past markers and computed future markers for verification
+        return Ok((vec![], expected_future_versions));
     }
 
     // Check that the sent proofs are for a contiguous sequence of decreasing versions
@@ -206,7 +212,7 @@ pub fn key_history_verify<TC: Configuration>(
         HistoryVerificationParams::AllowMissingValues { history_params } => history_params,
     };
     let (past_marker_versions, future_marker_versions) =
-        verify_with_history_params(current_epoch, &akd_label, &proof, params)?;
+        verify_with_history_params(current_epoch, &proof, params)?;
 
     // Verify all individual update proofs
     let mut maybe_previous_update_epoch = None;