Reducing visibility of membership verification functions (facebook#409)

haochenuw · Sep 20, 2023 · 68a4c97 · 68a4c97
1 parent 83a6dc9
commit 68a4c97
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 13 deletions.
diff --git a/akd/src/append_only_zks.rs b/akd/src/append_only_zks.rs
@@ -1173,7 +1173,7 @@ mod tests {
     use crate::utils::byte_arr_from_u64;
     use crate::{
         auditor::audit_verify,
-        client::{verify_membership, verify_nonmembership},
+        client::{verify_membership_for_tests_only, verify_nonmembership_for_tests_only},
         storage::memory::AsyncInMemoryDatabase,
     };
     use itertools::Itertools;
@@ -1729,7 +1729,7 @@ mod tests {
             .get_membership_proof::<TC, _>(&db, azks_element_set[0].label)
             .await?;
 
-        verify_membership::<TC>(azks.get_root_hash::<TC, _>(&db).await?, &proof)?;
+        verify_membership_for_tests_only::<TC>(azks.get_root_hash::<TC, _>(&db).await?, &proof)?;
 
         Ok(())
     }
@@ -1760,7 +1760,10 @@ mod tests {
                 .get_membership_proof::<TC, _>(&db, azks_element_set[0].label)
                 .await?;
 
-            verify_membership::<TC>(azks.get_root_hash::<TC, _>(&db).await?, &proof)?;
+            verify_membership_for_tests_only::<TC>(
+                azks.get_root_hash::<TC, _>(&db).await?,
+                &proof,
+            )?;
         }
         Ok(())
     }
@@ -1790,7 +1793,8 @@ mod tests {
             sibling_proofs: proof.sibling_proofs,
         };
         assert!(
-            verify_membership::<TC>(azks.get_root_hash::<TC, _>(&db).await?, &proof).is_err(),
+            verify_membership_for_tests_only::<TC>(azks.get_root_hash::<TC, _>(&db).await?, &proof)
+                .is_err(),
             "Membership proof does verify, despite being wrong"
         );
 
@@ -1833,7 +1837,11 @@ mod tests {
             .get_non_membership_proof::<TC, _>(&db, search_label)
             .await?;
         assert!(
-            verify_nonmembership::<TC>(azks.get_root_hash::<TC, _>(&db).await?, &proof).is_ok(),
+            verify_nonmembership_for_tests_only::<TC>(
+                azks.get_root_hash::<TC, _>(&db).await?,
+                &proof
+            )
+            .is_ok(),
             "Nonmembership proof does not verify"
         );
         Ok(())
@@ -1872,7 +1880,7 @@ mod tests {
             .get_non_membership_proof::<TC, _>(&db, search_label)
             .await?;
 
-        verify_nonmembership::<TC>(azks.get_root_hash::<TC, _>(&db).await?, &proof)?;
+        verify_nonmembership_for_tests_only::<TC>(azks.get_root_hash::<TC, _>(&db).await?, &proof)?;
 
         Ok(())
     }
@@ -1899,7 +1907,7 @@ mod tests {
             .get_non_membership_proof::<TC, _>(&db, search_label)
             .await?;
 
-        verify_nonmembership::<TC>(azks.get_root_hash::<TC, _>(&db).await?, &proof)?;
+        verify_nonmembership_for_tests_only::<TC>(azks.get_root_hash::<TC, _>(&db).await?, &proof)?;
 
         Ok(())
     }
@@ -1924,7 +1932,7 @@ mod tests {
             .get_non_membership_proof::<TC, _>(&db, search_label)
             .await?;
 
-        verify_nonmembership::<TC>(azks.get_root_hash::<TC, _>(&db).await?, &proof)?;
+        verify_nonmembership_for_tests_only::<TC>(azks.get_root_hash::<TC, _>(&db).await?, &proof)?;
 
         Ok(())
     }

diff --git a/akd_core/src/verify/base.rs b/akd_core/src/verify/base.rs
@@ -23,8 +23,19 @@ use alloc::format;
 use alloc::string::ToString;
 use core::convert::TryFrom;
 
-/// Verify the membership proof
-pub fn verify_membership<TC: Configuration>(
+/// Verifies a membership proof with respect to a root hash
+///
+/// Note: this function is only exposed for testing purposes and not meant to be used
+/// in a production setting on its own
+#[cfg(feature = "public_tests")]
+pub fn verify_membership_for_tests_only<TC: Configuration>(
+    root_hash: Digest,
+    proof: &MembershipProof,
+) -> Result<(), VerificationError> {
+    verify_membership::<TC>(root_hash, proof)
+}
+
+pub(crate) fn verify_membership<TC: Configuration>(
     root_hash: Digest,
     proof: &MembershipProof,
 ) -> Result<(), VerificationError> {
@@ -62,8 +73,19 @@ pub fn verify_membership<TC: Configuration>(
     }
 }
 
-/// Verifies the non-membership proof with respect to the root hash
-pub fn verify_nonmembership<TC: Configuration>(
+/// Verifies a non-membership proof with respect to a root hash
+///
+/// Note: this function is only exposed for testing purposes and not meant to be used
+/// in a production setting on its own
+#[cfg(feature = "public_tests")]
+pub fn verify_nonmembership_for_tests_only<TC: Configuration>(
+    root_hash: Digest,
+    proof: &NonMembershipProof,
+) -> Result<(), VerificationError> {
+    verify_nonmembership::<TC>(root_hash, proof)
+}
+
+pub(crate) fn verify_nonmembership<TC: Configuration>(
     root_hash: Digest,
     proof: &NonMembershipProof,
 ) -> Result<(), VerificationError> {

diff --git a/akd_core/src/verify/mod.rs b/akd_core/src/verify/mod.rs
@@ -78,6 +78,9 @@ impl From<protobuf::Error> for VerificationError {
 }
 
 // Re-export the necessary verification functions
-pub use base::{verify_membership, verify_nonmembership};
+
+#[cfg(feature = "public_tests")]
+pub use base::{verify_membership_for_tests_only, verify_nonmembership_for_tests_only};
+
 pub use history::{key_history_verify, HistoryVerificationParams};
 pub use lookup::lookup_verify;