akd_core/ecvrf/
traits.rs

1// Copyright (c) Meta Platforms, Inc. and affiliates.
2//
3// This source code is dual-licensed under either the MIT license found in the
4// LICENSE-MIT file in the root directory of this source tree or the Apache
5// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
6// of this source tree. You may select, at your option, one of the above-listed licenses.
7
8//! This module implements traits for managing ECVRF, mainly pertaining to storage
9//! of public and private keys
10use super::{Output, Proof, VRFExpandedPrivateKey, VRFPrivateKey, VRFPublicKey, VrfError};
11use crate::configuration::Configuration;
12use crate::{AkdLabel, AkdValue, NodeLabel, VersionFreshness};
13
14#[cfg(feature = "nostd")]
15use alloc::boxed::Box;
16#[cfg(feature = "nostd")]
17use alloc::vec::Vec;
18use async_trait::async_trait;
19use core::convert::TryInto;
20
21/// Represents a secure storage of the VRF private key. Since the VRF private key
22/// should change never (if it does, the entire tree is no longer a consistent mapping
23/// of user -> node label), it is highly recommended to back this implementation with a
24/// static cache of the private key bytes which lives for the life of the process.
25///
26/// I.e. retrieve the byte vector 1 time, and simply keep serving it up without doing
27/// network access calls
28#[async_trait]
29pub trait VRFKeyStorage: Clone + Sync + Send {
30    /* ======= To be implemented ====== */
31
32    /// Retrieve the VRF Private key as a vector of bytes
33    async fn retrieve(&self) -> Result<Vec<u8>, VrfError>;
34
35    /* ======= Common trait functionality ====== */
36
37    /// Retrieve the properly constructed VRF Private key
38    async fn get_vrf_private_key(&self) -> Result<VRFPrivateKey, VrfError> {
39        match self.retrieve().await {
40            Ok(bytes) => {
41                let pk_ref: &[u8] = &bytes;
42                pk_ref.try_into()
43            }
44            Err(other) => Err(other),
45        }
46    }
47
48    /// Retrieve the VRF public key
49    async fn get_vrf_public_key(&self) -> Result<VRFPublicKey, VrfError> {
50        self.get_vrf_private_key().await.map(|key| (&key).into())
51    }
52
53    /// Returns the [NodeLabel] that corresponds to a version of the label argument.
54    ///
55    /// The stale boolean here is to indicate whether we are getting the [NodeLabel] for a fresh version,
56    /// or a version that we are retiring.
57    async fn get_node_label<TC: Configuration>(
58        &self,
59        label: &AkdLabel,
60        freshness: VersionFreshness,
61        version: u64,
62    ) -> Result<NodeLabel, VrfError> {
63        let key = self.get_vrf_private_key().await?;
64        let expanded_key = VRFExpandedPrivateKey::from(&key);
65        let pk = VRFPublicKey::from(&key);
66        Ok(Self::get_node_label_with_expanded_key::<TC>(
67            &expanded_key,
68            &pk,
69            label,
70            freshness,
71            version,
72        ))
73    }
74
75    /// Returns the [NodeLabel] that corresponds to a version of the label argument utilizing the provided
76    /// private key.
77    ///
78    /// The stale boolean here is to indicate whether we are getting the [NodeLabel] for a fresh version,
79    /// or a version that we are retiring.
80    fn get_node_label_with_expanded_key<TC: Configuration>(
81        expanded_private_key: &VRFExpandedPrivateKey,
82        pk: &VRFPublicKey,
83        label: &AkdLabel,
84        freshness: VersionFreshness,
85        version: u64,
86    ) -> NodeLabel {
87        let output = Self::get_label_with_key_helper::<TC>(
88            expanded_private_key,
89            pk,
90            label,
91            freshness,
92            version,
93        );
94        NodeLabel::new(output.to_truncated_bytes(), 256)
95    }
96
97    /// Returns the tree nodelabel that corresponds to a vrf proof.
98    async fn get_node_label_from_vrf_proof(&self, proof: Proof) -> NodeLabel {
99        let output: super::ecvrf_impl::Output = (&proof).into();
100        NodeLabel::new(output.to_truncated_bytes(), 256)
101    }
102
103    /// Retrieve the proof for a specific label
104    async fn get_label_proof<TC: Configuration>(
105        &self,
106        label: &AkdLabel,
107        freshness: VersionFreshness,
108        version: u64,
109    ) -> Result<Proof, VrfError> {
110        let key = self.get_vrf_private_key().await?;
111        Ok(Self::get_label_proof_with_key::<TC>(
112            &key, label, freshness, version,
113        ))
114    }
115
116    /// Retrieve the proof for a specific label, with a supplied private key
117    fn get_label_proof_with_key<TC: Configuration>(
118        key: &VRFPrivateKey,
119        label: &AkdLabel,
120        freshness: VersionFreshness,
121        version: u64,
122    ) -> Proof {
123        let hashed_label = TC::get_hash_from_label_input(label, freshness, version);
124        key.prove(&hashed_label)
125    }
126
127    /// Retrieve the output for a specific label, with a supplied private key
128    fn get_label_with_key_helper<TC: Configuration>(
129        expanded_private_key: &VRFExpandedPrivateKey,
130        pk: &VRFPublicKey,
131        label: &AkdLabel,
132        freshness: VersionFreshness,
133        version: u64,
134    ) -> Output {
135        let hashed_label = TC::get_hash_from_label_input(label, freshness, version);
136        expanded_private_key.evaluate(pk, &hashed_label)
137    }
138
139    /// Returns the [NodeLabel]s that corresponds to a collection of (label, freshness, version) arguments
140    /// with only a single fetch to retrieve the VRF private key from storage.
141    ///
142    /// Note: The freshness enum here is to indicate whether we are getting the [NodeLabel] for a fresh version,
143    /// or a version that we are retiring.
144    async fn get_node_labels<TC: Configuration>(
145        &self,
146        labels: &[(AkdLabel, VersionFreshness, u64, AkdValue)],
147    ) -> Result<Vec<((AkdLabel, VersionFreshness, u64, AkdValue), NodeLabel)>, VrfError> {
148        let key = self.get_vrf_private_key().await?;
149        let expanded_key = VRFExpandedPrivateKey::from(&key);
150        let pk = VRFPublicKey::from(&key);
151
152        #[cfg(feature = "parallel_vrf")]
153        {
154            #[cfg(feature = "nostd")]
155            use alloc::format;
156
157            let mut join_set = tokio::task::JoinSet::new();
158            let labels_vec = labels.to_vec();
159            for (label, freshness, version, value) in labels_vec.into_iter() {
160                let expanded_key_ref = expanded_key.clone();
161                let pk_ref = pk.clone();
162
163                let future = {
164                    async move {
165                        (
166                            Self::get_node_label_with_expanded_key::<TC>(
167                                &expanded_key_ref,
168                                &pk_ref,
169                                &label,
170                                freshness,
171                                version,
172                            ),
173                            (label, freshness, version, value),
174                        )
175                    }
176                };
177                join_set.spawn(future);
178            }
179
180            let mut results = Vec::new();
181            while let Some(res) = join_set.join_next().await {
182                match res {
183                    Err(join_err) => {
184                        return Err(VrfError::SigningKey(format!(
185                            "Parallel VRF join error {join_err}"
186                        )))
187                    }
188                    Ok((node_label, (label, freshness, version, value))) => {
189                        results.push(((label, freshness, version, value), node_label));
190                    }
191                }
192            }
193            Ok(results)
194        }
195        #[cfg(not(feature = "parallel_vrf"))]
196        {
197            let mut results = Vec::new();
198            for (label, freshness, version, value) in labels {
199                let node_label = Self::get_node_label_with_expanded_key::<TC>(
200                    &expanded_key,
201                    &pk,
202                    label,
203                    *freshness,
204                    *version,
205                );
206                results.push((
207                    (label.clone(), *freshness, *version, value.clone()),
208                    node_label,
209                ));
210            }
211            Ok(results)
212        }
213    }
214}
akd_core/ecvrf/traits.rs

akd_core/ecvrf/
traits.rs
