diff --git a/akd/src/directory.rs b/akd/src/directory.rs index f72c8201..63ec3312 100644 --- a/akd/src/directory.rs +++ b/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), 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::(akd_label, VersionFreshness::Fresh, version) + .await?; + + let non_membership_proof = current_azks + .get_non_membership_proof::(&self.storage, node_label) + .await?; + non_existence_of_future_marker_proofs.push(non_membership_proof); + + let vrf_proof = self + .vrf + .get_label_proof::(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 = + 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, + ¤t_azks, + &future_marker_versions, + ) + .await?; + + let root_hash = EpochHash( + current_epoch, + current_azks.get_root_hash::(&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::(akd_label, VersionFreshness::Fresh, version) - .await?; - non_existence_of_future_marker_proofs.push( - current_azks - .get_non_membership_proof::(&self.storage, node_label) - .await?, - ); - future_marker_vrf_proofs.push( - self.vrf - .get_label_proof::(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, ¤t_azks, &future_marker_versions) + .await?; let root_hash = EpochHash( current_epoch, diff --git a/akd/src/storage/manager/mod.rs b/akd/src/storage/manager/mod.rs index 7328f5c4..5ecaf190 100644 --- a/akd/src/storage/manager/mod.rs +++ b/akd/src/storage/manager/mod.rs @@ -518,7 +518,7 @@ impl StorageManager { 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() { diff --git a/akd/src/storage/transaction.rs b/akd/src/storage/transaction.rs index 3359551d..5115e4cd 100644 --- a/akd/src/storage/transaction.rs +++ b/akd/src/storage/transaction.rs @@ -218,7 +218,7 @@ impl Transaction { flag: ValueStateRetrievalFlag, ) -> Option { 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); diff --git a/akd/src/tests/test_core_protocol.rs b/akd/src/tests/test_core_protocol.rs index 2da31d4a..52981e94 100644 --- a/akd/src/tests/test_core_protocol.rs +++ b/akd/src/tests/test_core_protocol.rs @@ -732,6 +732,135 @@ async fn test_simple_lookup_for_small_tree() -> Result<(), Ak Ok(()) } +test_config!(test_key_history_nonexistent_label); +async fn test_key_history_nonexistent_label() -> 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::::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::( + 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::( + &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() -> Result<(), AkdError> { let db = AsyncInMemoryDatabase::new(); diff --git a/akd_core/src/utils.rs b/akd_core/src/utils.rs index 62392758..0d0a26a9 100644 --- a/akd_core/src/utils.rs +++ b/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; @@ -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"); } diff --git a/akd_core/src/verify/history.rs b/akd_core/src/verify/history.rs index 19e4983a..b6348a30 100644 --- a/akd_core/src/verify/history.rs +++ b/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, Vec), 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 = + 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( 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;