Add akd_traits crate with KeyDirectory trait

Cargo.toml

@@ -1,4 +1,4 @@
 [workspace]
 
-members = ["akd", "akd_core", "examples", "xtask"]
+members = ["akd", "akd_core", "examples", "akd_traits", "xtask"]
 resolver = "2"

akd/Cargo.toml

@@ -23,7 +23,7 @@ default = [
     "experimental",
 ]
 
-bench = ["experimental", "public_tests", "tokio/rt-multi-thread"]
+bench = ["experimental", "public_tests", "tokio/rt-multi-thread", "akd_traits/bench"]
 # Greedy loading of lookup proof nodes
 greedy_lookup_preload = []
 public_auditing = ["dep:protobuf", "akd_core/protobuf"]
@@ -56,6 +56,7 @@ tracing_instrument = ["tracing/attributes"]
 akd_core = { version = "0.12.0-pre.12", path = "../akd_core", default-features = false, features = [
     "vrf",
 ] }
+akd_traits = { path = "../akd_traits" }
 async-recursion = "1"
 async-trait = "0.1"
 dashmap = "5"
@@ -102,3 +103,8 @@ required-features = ["bench"]
 name = "directory"
 harness = false
 required-features = ["bench"]
+
+[[bench]]
+name = "kd_benches"
+harness = false
+required-features = ["bench"]

akd/benches/kd_benches.rs

@@ -0,0 +1,80 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+//
+// This source code is dual-licensed under either the MIT license found in the
+// LICENSE-MIT file in the root directory of this source tree or the Apache
+// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
+// of this source tree. You may select, at your option, one of the above-listed licenses.
+
+use akd::append_only_zks::AzksParallelismConfig;
+use akd::ecvrf::HardCodedAkdVRF;
+use akd::storage::manager::StorageManager;
+use akd::storage::memory::AsyncInMemoryDatabase;
+use akd::{AkdLabel, AkdValue, Directory, LookupProof};
+use async_trait::async_trait;
+use criterion::Criterion;
+use rand::distributions::Alphanumeric;
+use rand::rngs::StdRng;
+use rand::{Rng, SeedableRng};
+
+type Config = akd::WhatsAppV1Configuration;
+
+struct AkdBenchSetup;
+
+#[async_trait]
+impl akd_traits::bench::BenchmarkSetup for AkdBenchSetup {
+    type Directory = Directory<Config, AsyncInMemoryDatabase, HardCodedAkdVRF>;
+
+    async fn create_directory() -> Self::Directory {
+        let db = AsyncInMemoryDatabase::new();
+        let storage = StorageManager::new_no_cache(db);
+        let vrf = HardCodedAkdVRF {};
+        Directory::<Config, _, _>::new(storage, vrf, AzksParallelismConfig::default())
+            .await
+            .unwrap()
+    }
+
+    fn generate_test_data(num_entries: usize, seed: u64) -> Vec<(AkdLabel, AkdValue)> {
+        let mut rng = StdRng::seed_from_u64(seed);
+        (0..num_entries)
+            .map(|i| {
+                let label = format!("user_{}", i);
+                let value: String = (0..16)
+                    .map(|_| rng.sample(Alphanumeric))
+                    .map(char::from)
+                    .collect();
+                (AkdLabel::from(&label), AkdValue::from(&value))
+            })
+            .collect()
+    }
+
+    fn name() -> &'static str {
+        "AKD (WhatsAppV1)"
+    }
+}
+
+/// Compute the approximate size of an AKD lookup proof in bytes.
+fn akd_lookup_proof_size(proof: &LookupProof) -> usize {
+    proof.existence_vrf_proof.len()
+        + proof.marker_vrf_proof.len()
+        + proof.freshness_vrf_proof.len()
+        + proof.commitment_nonce.len()
+        + std::mem::size_of_val(&proof.existence_proof)
+        + std::mem::size_of_val(&proof.marker_proof)
+        + std::mem::size_of_val(&proof.freshness_proof)
+}
+
+fn main() {
+    let mut criterion = Criterion::default().configure_from_args();
+
+    akd_traits::bench::bench_publish::<AkdBenchSetup>(&mut criterion);
+    akd_traits::bench::bench_lookup::<AkdBenchSetup>(&mut criterion);
+    akd_traits::bench::bench_lookup_verify::<AkdBenchSetup, _>(
+        &mut criterion,
+        akd_lookup_proof_size,
+    );
+    akd_traits::bench::bench_key_history::<AkdBenchSetup>(&mut criterion);
+    akd_traits::bench::bench_audit::<AkdBenchSetup>(&mut criterion);
+    akd_traits::bench::bench_audit_verify::<AkdBenchSetup>(&mut criterion);
+
+    criterion.final_summary();
+}

akd/src/directory.rs

@@ -22,8 +22,11 @@ use crate::{
 
 use crate::VersionFreshness;
 use akd_core::configuration::Configuration;
+use akd_core::types::VerifyResult;
 use akd_core::utils::get_marker_versions;
 use akd_core::verify::history::HistoryParams;
+use akd_traits::KeyDirectory;
+use async_trait::async_trait;
 use std::collections::{HashMap, HashSet};
 use std::marker::PhantomData;
 use std::sync::Arc;
@@ -847,6 +850,102 @@ where
     }
 }
 
+#[async_trait]
+impl<TC, S, V> KeyDirectory for Directory<TC, S, V>
+where
+    TC: Configuration,
+    S: Database + 'static,
+    V: VRFKeyStorage,
+{
+    type LookupProof = LookupProof;
+    type HistoryProof = HistoryProof;
+    type AuditProof = AppendOnlyProof;
+    type PublicKey = VRFPublicKey;
+    type HistoryParams = akd_core::verify::history::HistoryParams;
+    type HistoryVerificationParams = crate::client::HistoryVerificationParams;
+    type Error = AkdError;
+
+    async fn publish(&self, updates: Vec<(AkdLabel, AkdValue)>) -> Result<EpochHash, AkdError> {
+        Directory::publish(self, updates).await
+    }
+
+    async fn lookup(&self, label: AkdLabel) -> Result<(LookupProof, EpochHash), AkdError> {
+        Directory::lookup(self, label).await
+    }
+
+    async fn batch_lookup(
+        &self,
+        labels: &[AkdLabel],
+    ) -> Result<(Vec<LookupProof>, EpochHash), AkdError> {
+        Directory::batch_lookup(self, labels).await
+    }
+
+    async fn key_history(
+        &self,
+        label: &AkdLabel,
+        params: akd_core::verify::history::HistoryParams,
+    ) -> Result<(HistoryProof, EpochHash), AkdError> {
+        Directory::key_history(self, label, params).await
+    }
+
+    async fn audit(&self, start_epoch: u64, end_epoch: u64) -> Result<AppendOnlyProof, AkdError> {
+        Directory::audit(self, start_epoch, end_epoch).await
+    }
+
+    async fn get_public_key(&self) -> Result<VRFPublicKey, AkdError> {
+        Directory::get_public_key(self).await
+    }
+
+    async fn get_epoch_hash(&self) -> Result<EpochHash, AkdError> {
+        Directory::get_epoch_hash(self).await
+    }
+
+    fn lookup_verify(
+        public_key: &VRFPublicKey,
+        root_hash: Digest,
+        current_epoch: u64,
+        label: AkdLabel,
+        proof: LookupProof,
+    ) -> Result<VerifyResult, akd_traits::KeyDirectoryError> {
+        akd_core::verify::lookup::lookup_verify::<TC>(
+            public_key.as_bytes(),
+            root_hash,
+            current_epoch,
+            label,
+            proof,
+        )
+        .map_err(|e| akd_traits::KeyDirectoryError::Verification(format!("{e:?}")))
+    }
+
+    fn key_history_verify(
+        public_key: &VRFPublicKey,
+        root_hash: Digest,
+        current_epoch: u64,
+        label: AkdLabel,
+        proof: HistoryProof,
+        params: crate::client::HistoryVerificationParams,
+    ) -> Result<Vec<VerifyResult>, akd_traits::KeyDirectoryError> {
+        akd_core::verify::history::key_history_verify::<TC>(
+            public_key.as_bytes(),
+            root_hash,
+            current_epoch,
+            label,
+            proof,
+            params,
+        )
+        .map_err(|e| akd_traits::KeyDirectoryError::Verification(format!("{e:?}")))
+    }
+
+    async fn audit_verify(
+        hashes: Vec<Digest>,
+        proof: AppendOnlyProof,
+    ) -> Result<(), akd_traits::KeyDirectoryError> {
+        crate::auditor::audit_verify::<TC>(hashes, proof)
+            .await
+            .map_err(|e| e.into())
+    }
+}
+
 /// A thin newtype which offers read-only interactivity with a [Directory].
 #[derive(Clone)]
 pub struct ReadOnlyDirectory<TC, S, V>(Directory<TC, S, V>)

akd/src/errors.rs

@@ -338,3 +338,18 @@ impl fmt::Display for ParallelismError {
         }
     }
 }
+
+impl From<AkdError> for akd_traits::KeyDirectoryError {
+    fn from(e: AkdError) -> Self {
+        match e {
+            AkdError::TreeNode(e) => akd_traits::KeyDirectoryError::Directory(format!("{e:?}")),
+            AkdError::Directory(e) => akd_traits::KeyDirectoryError::Directory(format!("{e:?}")),
+            AkdError::AzksErr(e) => akd_traits::KeyDirectoryError::Directory(format!("{e:?}")),
+            AkdError::Vrf(e) => akd_traits::KeyDirectoryError::Directory(format!("{e:?}")),
+            AkdError::Storage(e) => akd_traits::KeyDirectoryError::Storage(format!("{e:?}")),
+            AkdError::AuditErr(e) => akd_traits::KeyDirectoryError::Audit(format!("{e:?}")),
+            AkdError::Parallelism(e) => akd_traits::KeyDirectoryError::Other(format!("{e:?}")),
+            AkdError::TestErr(s) => akd_traits::KeyDirectoryError::Other(s),
+        }
+    }
+}

akd/src/helper_structs.rs

@@ -8,24 +8,8 @@
 //! Helper structs that are used for various data structures,
 //! to make it easier to pass arguments around.
 
-use crate::Digest;
 use crate::{storage::types::ValueState, NodeLabel};
 
-/// Root hash of the tree and its associated epoch
-#[derive(Debug, Clone, Hash, PartialEq, Eq)]
-pub struct EpochHash(pub u64, pub Digest);
-
-impl EpochHash {
-    /// Get the contained epoch
-    pub fn epoch(&self) -> u64 {
-        self.0
-    }
-    /// Get the contained hash
-    pub fn hash(&self) -> Digest {
-        self.1
-    }
-}
-
 #[derive(Clone, Debug)]
 /// Info needed for a lookup of a user for an epoch
 pub struct LookupInfo {

akd/src/lib.rs

@@ -565,10 +565,10 @@ pub use akd_core::{
 mod utils;
 
 // ========== Type re-exports which are commonly used ========== //
+pub use akd_traits::KeyDirectory;
 pub use append_only_zks::{Azks, AzksParallelismConfig, AzksParallelismOption};
 pub use client::HistoryVerificationParams;
 pub use directory::Directory;
-pub use helper_structs::EpochHash;
 
 // ========== Constants and type aliases ========== //
 #[cfg(any(test, feature = "public_tests"))]

akd_core/Cargo.toml

@@ -23,19 +23,21 @@ whatsapp_v1 = ["dep:blake3"]
 experimental = ["dep:blake3"]
 # Include the VRF verification logic
 vrf = ["ed25519-dalek", "curve25519-dalek"]
-serde_serialization = ["dep:serde", "dep:serde_bytes", "ed25519-dalek/serde"]
+serde_serialization = ["dep:serde", "dep:serde_bytes", "ed25519-dalek/serde", "akd_traits/serde"]
 # Parallelize VRF calculations during publish
 parallel_vrf = ["tokio"]
 
 bench = ["parallel_vrf", "experimental", "vrf", "tokio/rt-multi-thread"]
 public_tests = ["dep:paste"]
 protobuf = ["dep:protobuf"]
+rand = ["dep:rand", "akd_traits/rand"]
 
 # Default features mix
 default = ["vrf", "experimental"]
 
 [dependencies]
 ## Required dependencies ##
+akd_traits = { path = "../akd_traits", default-features = false }
 async-trait = "0.1"
 curve25519-dalek = { version = "4", optional = true }
 ed25519-dalek = { version = "2", features = [