From 8f288877147314161efe137a8037a3bc5e34bd2b Mon Sep 17 00:00:00 2001
From: Michael Rosenberg <michael@mrosenberg.pub>
Date: Sun, 21 Aug 2022 04:13:02 -0400
Subject: [PATCH] Added generation from seeds

---
 Cargo.toml |   1 +
 src/gen.rs | 208 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 src/lib.rs | 103 ++++++++++++++++----------
 3 files changed, 276 insertions(+), 36 deletions(-)
 create mode 100644 src/gen.rs

diff --git a/Cargo.toml b/Cargo.toml
index 28e2fed..735976e 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -6,6 +6,7 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+rand_chacha = "0.3"
 rand_core = "0.6"
 digest = "0.10"
 serde = "1"
diff --git a/src/gen.rs b/src/gen.rs
new file mode 100644
index 0000000..faefff3
--- /dev/null
+++ b/src/gen.rs
@@ -0,0 +1,208 @@
+use crate::{Label, NotaryLabelTree, UserLabelTree};
+
+use alloc::vec::Vec;
+
+use digest::Digest;
+use rand_chacha::ChaCha20Rng;
+use rand_core::{CryptoRng, RngCore, SeedableRng};
+use serde::{Deserialize as SerdeDeserialize, Serialize as SerdeSerialize};
+
+// The psuedorandom generator we use for all label and randomness generation
+type Prg = ChaCha20Rng;
+// The seed type of `Prg`
+type Seed = <Prg as SeedableRng>::Seed;
+
+/// A seed used by the Notary to generate random labels
+#[derive(Copy, Clone, SerdeSerialize, SerdeDeserialize)]
+pub struct LabelSeed(Seed);
+
+/// A seed used by the User to generate random label blinders
+#[derive(Copy, Clone, SerdeSerialize, SerdeDeserialize)]
+pub struct LabelRandomnessSeed(Seed);
+
+impl LabelSeed {
+    /// Generates a random seed for label generation
+    pub fn new<R: CryptoRng + RngCore>(mut rng: R) -> LabelSeed {
+        let mut seed = Seed::default();
+        rng.fill_bytes(&mut seed);
+
+        LabelSeed(seed)
+    }
+
+    /// Generates a pseudorandom label set for the given number of bits. The labels are ordered as
+    /// `[label1bit0, label1bit1, label2bit0, label2bit1, ...]`. That is, the output length is `2 *
+    /// num_bits`.
+    pub fn gen_notary_labels<H: Digest>(&self, num_bits: usize) -> Vec<Label> {
+        // Make an RNG from our label seed
+        let mut rng = Prg::from_seed(self.0);
+
+        // Generate all the labels. There are 2*num_bits many of them
+        (0..(2 * num_bits))
+            .map(|_| {
+                let mut buf = Label::default();
+                rng.fill_bytes(&mut buf);
+                buf
+            })
+            .collect()
+    }
+}
+
+impl<H: Digest> NotaryLabelTree<H> {
+    /// Commits to the Notary's label set by making a tree out of the labels. `num_bits` is the
+    /// number of plaintext bits, and `label_seed` is used to generate the labels.
+    pub fn gen_from_seed(num_bits: usize, label_seed: LabelSeed) -> NotaryLabelTree<H> {
+        // Make the RNGs for the blinders and labels
+        let mut label_rng = Prg::from_seed(label_seed.0);
+
+        // Use the label seed to collect the Notary's labels. There are 2*num_bits many of them
+        let labels = (0..2 * num_bits).map(|_| {
+            // Compute the label
+            let mut label = Label::default();
+            label_rng.fill_bytes(&mut label);
+            label
+        });
+
+        // Make the tree
+        Self::from_labels(labels)
+    }
+}
+
+impl LabelRandomnessSeed {
+    /// Generates a random seed for label blinder generation
+    pub fn new<R: RngCore + CryptoRng>(mut rng: R) -> LabelRandomnessSeed {
+        let mut seed = <ChaCha20Rng as SeedableRng>::Seed::default();
+        rng.fill_bytes(&mut seed);
+
+        LabelRandomnessSeed(seed)
+    }
+}
+
+impl<H: Digest> UserLabelTree<H> {
+    /// Commits to the active label set by making a tree out of the (blinded) labels. `seed` is
+    /// used to generate the blinds. This method is called by the User before she learns the
+    /// decoding of her plaintext.
+    pub fn gen_from_labels_with_seed<'a>(
+        blinder_seed: LabelRandomnessSeed,
+        user_labels: impl IntoIterator<Item = Label>,
+    ) -> UserLabelTree<H> {
+        // Make an RNG from our blinder seed
+        let mut rng = Prg::from_seed(blinder_seed.0);
+
+        UserLabelTree::gen_from_labels(&mut rng, user_labels)
+    }
+
+    /// Commits to the active label set by making a tree out of the (blinded) labels. `plaintext`
+    /// are the plaintext bits. `blinder_seed` is used to generate the blinds. `label_seed` is used
+    /// to generate the labels.
+    pub fn gen_from_seeds(
+        plaintext: &[bool],
+        blinder_seed: LabelRandomnessSeed,
+        label_seed: LabelSeed,
+    ) -> UserLabelTree<H> {
+        // Make the RNGs for the blinders and labels
+        let mut label_rng = Prg::from_seed(label_seed.0);
+
+        // Use the label seed to collect just the active labels
+        let active_labels = (0..plaintext.len() * 2).filter_map(|i| {
+            // Compute the label
+            let mut label = Label::default();
+            label_rng.fill_bytes(&mut label);
+
+            // Get the plaintext bit that this label MAY represent
+            let bit = plaintext[i / 2];
+
+            // If the bit is 0, then the corresponding label is the even label. If the bit is 1
+            // then it's the odd label. Only save the label if its parity matches the bit.
+            if (bit as usize) == i % 2 {
+                Some(label)
+            } else {
+                None
+            }
+        });
+
+        Self::gen_from_labels_with_seed(blinder_seed, active_labels)
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use crate::{verify_bits, Plaintext};
+
+    use blake2::Blake2s256;
+    use rand::{thread_rng, Rng};
+
+    type H = Blake2s256;
+
+    const LOG_PLAINTEXT_BITLEN: usize = 15;
+    const PLAINTEXT_BITLEN: usize = 1 << LOG_PLAINTEXT_BITLEN;
+
+    #[test]
+    fn gen_correctness() {
+        let mut rng = thread_rng();
+
+        // Random plaintext
+        let plaintext: Vec<bool> = core::iter::repeat_with(|| rng.gen::<bool>())
+            .take(PLAINTEXT_BITLEN)
+            .collect();
+
+        //
+        // Notary
+        //
+
+        // Pick a label seed and generate the labels
+        let label_seed = LabelSeed::new(&mut rng);
+        let all_labels = label_seed.gen_notary_labels::<H>(PLAINTEXT_BITLEN);
+
+        // Commit to all the labels
+        let notary_root = {
+            let notary_tree = NotaryLabelTree::<H>::from_labels(all_labels.clone());
+            notary_tree.root()
+        };
+
+        //
+        // User
+        //
+
+        // Placeholder for a garbled circuit. The user ends up with a set of plaintext labels (that
+        // we've determined in this test in advance)
+        let plaintext_labels: Vec<Label> = plaintext
+            .iter()
+            .enumerate()
+            .map(|(i, &b)| all_labels.get(2 * i + (b as usize)).unwrap())
+            .cloned()
+            .collect();
+
+        // Commit to user's labels, using a seed for the commitment randomness
+        let blinder_seed = LabelRandomnessSeed::new(&mut rng);
+        let user_root = {
+            let user_tree =
+                UserLabelTree::<H>::gen_from_labels_with_seed(blinder_seed, plaintext_labels);
+            user_tree.root()
+        };
+
+        //
+        // User sends root to the notary, gets it signed, learns the plaintext, and learns the
+        // Notary's label seed. The user now wants to do a selective disclosure to a third party
+        //
+
+        // Reconstruct the Plaintext struct from just seeds
+        let user_tree = UserLabelTree::gen_from_seeds(&plaintext, blinder_seed, label_seed);
+        assert_eq!(user_tree.root(), user_root);
+        let notary_tree = NotaryLabelTree::gen_from_seed(PLAINTEXT_BITLEN, label_seed);
+        let pt = Plaintext {
+            bits: plaintext,
+            label_tree: user_tree,
+        };
+
+        // Prove a randomly chosen bit slice
+        let subslice_size = 100;
+        let bit_idx = rng.gen_range(0..PLAINTEXT_BITLEN - subslice_size);
+        let range = bit_idx..=(bit_idx + subslice_size - 1);
+        let batch_proof = pt.prove_bits(range.clone(), &notary_tree);
+
+        // Verify the batch proof
+        let bits = &pt.bits[range.clone()];
+        verify_bits(&bits, range.clone(), &user_root, &notary_root, &batch_proof).unwrap();
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
index fd9bc09..d18f59e 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -5,6 +5,8 @@ extern crate alloc;
 #[cfg(feature = "std")]
 extern crate std;
 
+pub mod gen;
+
 use alloc::vec::Vec;
 
 use ct_merkle::{
@@ -15,7 +17,7 @@ use digest::Digest;
 use rand_core::{CryptoRng, RngCore};
 use serde::{Deserialize as SerdeDeserialize, Serialize as SerdeSerialize};
 
-type Label = [u8; 16];
+pub type Label = [u8; 16];
 type Randomness = [u8; 16];
 
 /// Contains a label and a random string used for hiding in the commitment
@@ -28,14 +30,16 @@ struct LabelAndRandomness {
 //The UserLabel of bit `b` at index `i` appears in this data structure at index `2i + b`.
 /// The Merkle tree containing all the Notary's labels.
 #[derive(Clone, SerdeSerialize, SerdeDeserialize)]
-pub struct NotaryLabels<H: Digest>(CtMerkleTree<H, Label>);
+pub struct NotaryLabelTree<H: Digest>(CtMerkleTree<H, Label>);
 /// The Merkle tree containing all the User's labels.
 #[derive(Clone, SerdeSerialize, SerdeDeserialize)]
-pub struct UserLabels<H: Digest>(CtMerkleTree<H, LabelAndRandomness>);
+pub struct UserLabelTree<H: Digest>(CtMerkleTree<H, LabelAndRandomness>);
 
 /// The root hash of the Notary's label tree
+#[derive(Clone, Debug, SerdeSerialize, SerdeDeserialize)]
 pub struct NotaryLabelRoot<H: Digest>(RootHash<H>);
 /// The root hash of the User's label tree
+#[derive(Clone, Debug, SerdeSerialize, SerdeDeserialize)]
 pub struct UserLabelRoot<H: Digest>(RootHash<H>);
 
 /// A batch proof revealing a number of plaintext bits
@@ -55,9 +59,27 @@ pub struct Plaintext<H: Digest> {
     /// The plaintext bitstring
     pub bits: Vec<bool>,
     /// The Merkle tree of the plaintext
-    pub label_tree: UserLabels<H>,
+    pub label_tree: UserLabelTree<H>,
 }
 
+impl<H: Digest> PartialEq for UserLabelRoot<H> {
+    /// Compares this `RootHash` to another in constant time.
+    fn eq(&self, other: &UserLabelRoot<H>) -> bool {
+        self.0.eq(&other.0)
+    }
+}
+
+impl<H: Digest> Eq for UserLabelRoot<H> {}
+
+impl<H: Digest> PartialEq for NotaryLabelRoot<H> {
+    /// Compares this `RootHash` to another in constant time.
+    fn eq(&self, other: &NotaryLabelRoot<H>) -> bool {
+        self.0.eq(&other.0)
+    }
+}
+
+impl<H: Digest> Eq for NotaryLabelRoot<H> {}
+
 impl CanonicalSerialize for LabelAndRandomness {
     fn serialize<W: SimpleWriter>(&self, mut w: W) {
         CanonicalSerialize::serialize(&self.label, &mut w);
@@ -65,51 +87,55 @@ impl CanonicalSerialize for LabelAndRandomness {
     }
 }
 
-impl<H: Digest> NotaryLabels<H> {
+impl<H: Digest> NotaryLabelTree<H> {
     /// Get the Merkle root of the Notary label tree
     pub fn root(&self) -> NotaryLabelRoot<H> {
         NotaryLabelRoot(self.0.root())
     }
 }
 
-impl<H: Digest> UserLabels<H> {
+impl<H: Digest> UserLabelTree<H> {
     /// Get the Merkle root of the User label tree
     pub fn root(&self) -> UserLabelRoot<H> {
         UserLabelRoot(self.0.root())
     }
 }
 
-/// Commits to the active label set by making a tree out of the (blinded) labels. This is done by
-/// the User before she learns the decoding of her plaintext.
-pub fn user_commit<H, R>(mut rng: R, labels: &[Label]) -> UserLabels<H>
-where
-    R: CryptoRng + RngCore,
-    H: Digest,
-{
-    let mut tree = CtMerkleTree::new();
-    labels.iter().cloned().for_each(|label| {
-        let mut randomness = Randomness::default();
-        rng.fill_bytes(&mut randomness);
+impl<H: Digest> UserLabelTree<H> {
+    /// Commits to the active label set by making a tree out of the (blinded) labels. This is done
+    /// by the User before she learns the decoding of her plaintext.
+    pub fn gen_from_labels<'a, R>(
+        mut rng: R,
+        user_labels: impl IntoIterator<Item = Label>,
+    ) -> UserLabelTree<H>
+    where
+        R: CryptoRng + RngCore,
+        H: Digest,
+    {
+        let mut tree = CtMerkleTree::new();
+        user_labels.into_iter().for_each(|label| {
+            let mut randomness = Randomness::default();
+            rng.fill_bytes(&mut randomness);
 
-        let entry = LabelAndRandomness { label, randomness };
-        tree.push(entry);
-    });
+            let entry = LabelAndRandomness { label, randomness };
+            tree.push(entry);
+        });
 
-    UserLabels(tree)
+        UserLabelTree(tree)
+    }
 }
 
-/// Commits to the full label set by making a tree out of the  labels. The labels must be ordered
-/// as `[label1bit0, label1bit1, label2bit0, label2bit1, ...]`
-pub fn notary_commit<H>(labels: &[Label]) -> NotaryLabels<H>
-where
-    H: Digest,
-{
-    let mut tree = CtMerkleTree::new();
-    labels.iter().cloned().for_each(|label| {
-        tree.push(label);
-    });
+impl<H: Digest> NotaryLabelTree<H> {
+    /// Constructs a tree from the full label set. The labels must be ordered as `[label1bit0,
+    /// label1bit1, label2bit0, label2bit1, ...]`
+    pub fn from_labels(notary_labels: impl IntoIterator<Item = Label>) -> NotaryLabelTree<H> {
+        let mut tree = CtMerkleTree::new();
+        notary_labels.into_iter().for_each(|label| {
+            tree.push(label);
+        });
 
-    NotaryLabels(tree)
+        NotaryLabelTree(tree)
+    }
 }
 
 impl<H: Digest> Plaintext<H> {
@@ -119,7 +145,7 @@ impl<H: Digest> Plaintext<H> {
     pub fn prove_bits(
         &self,
         range: core::ops::RangeInclusive<usize>,
-        all_labels: &NotaryLabels<H>,
+        all_labels: &NotaryLabelTree<H>,
     ) -> BatchBitProof<H> {
         let active_idxs: Vec<usize> = range.clone().collect();
         let bits = self.bits[range.clone()].to_vec();
@@ -170,7 +196,12 @@ pub fn verify_bits<H: Digest>(
         &proof.active_labels_proof,
     )?;
 
+    // Extract just the label (not the randomness) from the user's label tree. This is used for
+    // verification of the Notary batch inclusion proof
     let active_labels: Vec<Label> = proof.active_labels.iter().map(|lr| lr.label).collect();
+
+    // Get the corresponding indices into the full label set. Recall label i bit b occurs in
+    // the full set at index 2i + b
     let full_idxs: Vec<usize> = active_idxs
         .iter()
         .zip(bits.iter())
@@ -196,7 +227,7 @@ mod tests {
     const LOG_PLAINTEXT_BITLEN: usize = 15;
     const PLAINTEXT_BITLEN: usize = 1 << LOG_PLAINTEXT_BITLEN;
 
-    fn setup() -> (Plaintext<H>, NotaryLabels<H>) {
+    fn setup() -> (Plaintext<H>, NotaryLabelTree<H>) {
         let mut rng = thread_rng();
 
         // Generate random labels
@@ -218,8 +249,8 @@ mod tests {
             .collect();
 
         // Commit to everything
-        let user_tree = user_commit::<H, _>(&mut rng, &plaintext_labels);
-        let notary_tree = notary_commit::<H>(&all_labels);
+        let user_tree = UserLabelTree::gen_from_labels(&mut rng, plaintext_labels);
+        let notary_tree = NotaryLabelTree::from_labels(all_labels);
 
         let pt = Plaintext {
             label_tree: user_tree,