Merge pull request #68 from str4d/sapling-commitment-tree

Sapling commitment tree

Cargo.lock

@@ -553,6 +553,7 @@ dependencies = [
  "pairing 0.14.2",
  "rand 0.4.2 (registry+https://github.com/rust-lang/crates.io-index)",
  "sapling-crypto 0.0.1",
+ "zcash_primitives 0.0.0",
 ]
 
 [metadata]

librustzcash/src/rustzcash.rs

@@ -11,7 +11,7 @@ extern crate zcash_proofs;
 
 extern crate lazy_static;
 
-use ff::{BitIterator, PrimeField, PrimeFieldRepr};
+use ff::{PrimeField, PrimeFieldRepr};
 use pairing::bls12_381::{Bls12, Fr, FrRepr};
 
 use sapling_crypto::{
@@ -22,7 +22,6 @@ use sapling_crypto::{
         fs::{Fs, FsRepr},
         FixedGenerators, JubjubEngine, JubjubParams, PrimeOrder, ToUniform, Unknown,
     },
-    pedersen_hash::{pedersen_hash, Personalization},
     redjubjub::{self, Signature},
 };
 
@@ -57,10 +56,15 @@ use std::ffi::OsString;
 use std::os::windows::ffi::OsStringExt;
 
 use sapling_crypto::primitives::{ProofGenerationKey, ViewingKey};
-use zcash_primitives::{note_encryption::sapling_ka_agree, sapling::spend_sig, zip32, JUBJUB};
+use zcash_primitives::{
+    merkle_tree::CommitmentTreeWitness,
+    note_encryption::sapling_ka_agree,
+    sapling::{merkle_hash, spend_sig},
+    zip32, JUBJUB,
+};
 use zcash_proofs::{
     load_parameters,
-    sapling::{CommitmentTreeWitness, SaplingProvingContext, SaplingVerificationContext},
+    sapling::{SaplingProvingContext, SaplingVerificationContext},
 };
 
 pub mod equihash;
@@ -254,28 +258,7 @@ pub extern "system" fn librustzcash_merkle_hash(
     // size of the representation
     let b_repr = read_le(unsafe { &(&*b)[..] });
 
-    let mut lhs = [false; 256];
+    let tmp = merkle_hash(depth, &a_repr, &b_repr);
-    let mut rhs = [false; 256];
-
-    for (a, b) in lhs.iter_mut().rev().zip(BitIterator::new(a_repr)) {
-        *a = b;
-    }
-
-    for (a, b) in rhs.iter_mut().rev().zip(BitIterator::new(b_repr)) {
-        *a = b;
-    }
-
-    let tmp = pedersen_hash::<Bls12, _>(
-        Personalization::MerkleTree(depth),
-        lhs.iter()
-            .map(|&x| x)
-            .take(Fr::NUM_BITS as usize)
-            .chain(rhs.iter().map(|&x| x).take(Fr::NUM_BITS as usize)),
-        &JUBJUB,
-    )
-    .into_xy()
-    .0
-    .into_repr();
 
     // Should be okay, caller is responsible for ensuring the pointer
     // is a valid pointer to 32 bytes that can be mutated.

zcash_primitives/src/lib.rs

@@ -17,6 +17,7 @@ use sapling_crypto::jubjub::JubjubBls12;
 
 pub mod block;
 pub mod keys;
+pub mod merkle_tree;
 pub mod note_encryption;
 pub mod sapling;
 mod serialize;

zcash_primitives/src/sapling.rs

@@ -1,9 +1,109 @@
-use pairing::bls12_381::Bls12;
+//! Structs and constants specific to the Sapling shielded pool.
+
+use ff::{BitIterator, PrimeField, PrimeFieldRepr};
+use pairing::bls12_381::{Bls12, Fr, FrRepr};
 use rand::OsRng;
 use sapling_crypto::{
     jubjub::{fs::Fs, FixedGenerators, JubjubBls12},
+    pedersen_hash::{pedersen_hash, Personalization},
+    primitives::Note,
     redjubjub::{PrivateKey, PublicKey, Signature},
 };
+use std::io::{self, Read, Write};
+
+use crate::merkle_tree::Hashable;
+use JUBJUB;
+
+pub(crate) const SAPLING_COMMITMENT_TREE_DEPTH: usize =
+    sapling_crypto::circuit::sapling::TREE_DEPTH;
+
+/// Compute a parent node in the Sapling commitment tree given its two children.
+pub fn merkle_hash(depth: usize, lhs: &FrRepr, rhs: &FrRepr) -> FrRepr {
+    let lhs = {
+        let mut tmp = [false; 256];
+        for (a, b) in tmp.iter_mut().rev().zip(BitIterator::new(lhs)) {
+            *a = b;
+        }
+        tmp
+    };
+
+    let rhs = {
+        let mut tmp = [false; 256];
+        for (a, b) in tmp.iter_mut().rev().zip(BitIterator::new(rhs)) {
+            *a = b;
+        }
+        tmp
+    };
+
+    pedersen_hash::<Bls12, _>(
+        Personalization::MerkleTree(depth),
+        lhs.iter()
+            .map(|&x| x)
+            .take(Fr::NUM_BITS as usize)
+            .chain(rhs.iter().map(|&x| x).take(Fr::NUM_BITS as usize)),
+        &JUBJUB,
+    )
+    .into_xy()
+    .0
+    .into_repr()
+}
+
+/// A node within the Sapling commitment tree.
+#[derive(Clone, Copy, Debug, PartialEq)]
+pub struct Node {
+    repr: FrRepr,
+}
+
+impl Node {
+    pub fn new(repr: FrRepr) -> Self {
+        Node { repr }
+    }
+}
+
+impl Hashable for Node {
+    fn read<R: Read>(mut reader: R) -> io::Result<Self> {
+        let mut repr = FrRepr::default();
+        repr.read_le(&mut reader)?;
+        Ok(Node::new(repr))
+    }
+
+    fn write<W: Write>(&self, mut writer: W) -> io::Result<()> {
+        self.repr.write_le(&mut writer)
+    }
+
+    fn combine(depth: usize, lhs: &Self, rhs: &Self) -> Self {
+        Node {
+            repr: merkle_hash(depth, &lhs.repr, &rhs.repr),
+        }
+    }
+
+    fn blank() -> Self {
+        Node {
+            repr: Note::<Bls12>::uncommitted().into_repr(),
+        }
+    }
+
+    fn empty_root(depth: usize) -> Self {
+        EMPTY_ROOTS[depth]
+    }
+}
+
+impl From<Node> for Fr {
+    fn from(node: Node) -> Self {
+        Fr::from_repr(node.repr).expect("Tree nodes should be in the prime field")
+    }
+}
+
+lazy_static! {
+    static ref EMPTY_ROOTS: Vec<Node> = {
+        let mut v = vec![Node::blank()];
+        for d in 0..SAPLING_COMMITMENT_TREE_DEPTH {
+            let next = Node::combine(d, &v[d], &v[d]);
+            v.push(next);
+        }
+        v
+    };
+}
 
 /// Create the spendAuthSig for a Sapling SpendDescription.
 pub fn spend_sig(

zcash_primitives/src/serialize.rs

@@ -82,6 +82,37 @@ impl Vector {
     }
 }
 
+pub struct Optional;
+
+impl Optional {
+    pub fn read<R: Read, T, F>(mut reader: R, func: F) -> io::Result<Option<T>>
+    where
+        F: Fn(&mut R) -> io::Result<T>,
+    {
+        match reader.read_u8()? {
+            0 => Ok(None),
+            1 => Ok(Some(func(&mut reader)?)),
+            _ => Err(io::Error::new(
+                io::ErrorKind::InvalidInput,
+                "non-canonical Option<T>",
+            )),
+        }
+    }
+
+    pub fn write<W: Write, T, F>(mut writer: W, val: &Option<T>, func: F) -> io::Result<()>
+    where
+        F: Fn(&mut W, &T) -> io::Result<()>,
+    {
+        match val {
+            None => writer.write_u8(0),
+            Some(e) => {
+                writer.write_u8(1)?;
+                func(&mut writer, e)
+            }
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -153,4 +184,47 @@ mod tests {
             eval!(vec![7; 260], expected);
         }
     }
+
+    #[test]
+    fn optional() {
+        macro_rules! eval {
+            ($value:expr, $expected:expr, $write:expr, $read:expr) => {
+                let mut data = vec![];
+                Optional::write(&mut data, &$value, $write).unwrap();
+                assert_eq!(&data[..], &$expected[..]);
+                match Optional::read(&data[..], $read) {
+                    Ok(v) => assert_eq!(v, $value),
+                    Err(e) => panic!("Unexpected error: {:?}", e),
+                }
+            };
+        }
+
+        macro_rules! eval_u8 {
+            ($value:expr, $expected:expr) => {
+                eval!($value, $expected, |w, e| w.write_u8(*e), |r| r.read_u8())
+            };
+        }
+
+        macro_rules! eval_vec {
+            ($value:expr, $expected:expr) => {
+                eval!(
+                    $value,
+                    $expected,
+                    |w, v| Vector::write(w, v, |w, e| w.write_u8(*e)),
+                    |r| Vector::read(r, |r| r.read_u8())
+                )
+            };
+        }
+
+        eval_u8!(None, [0]);
+        eval_u8!(Some(0), [1, 0]);
+        eval_u8!(Some(1), [1, 1]);
+        eval_u8!(Some(5), [1, 5]);
+
+        eval_vec!(None as Option<Vec<_>>, [0]);
+        eval_vec!(Some(vec![]), [1, 0]);
+        eval_vec!(Some(vec![0]), [1, 1, 0]);
+        eval_vec!(Some(vec![1]), [1, 1, 1]);
+        eval_vec!(Some(vec![5; 8]), [1, 8, 5, 5, 5, 5, 5, 5, 5, 5]);
+    }
 }

zcash_proofs/Cargo.toml

@@ -12,6 +12,7 @@ ff = { path = "../ff" }
 pairing = { path = "../pairing" }
 rand = "0.4"
 sapling-crypto = { path = "../sapling-crypto" }
+zcash_primitives = { path = "../zcash_primitives" }
 
 [dependencies.blake2-rfc]
 git = "https://github.com/gtank/blake2-rfc"

zcash_proofs/src/lib.rs

@@ -5,6 +5,7 @@ extern crate ff;
 extern crate pairing;
 extern crate rand;
 extern crate sapling_crypto;
+extern crate zcash_primitives;
 
 use bellman::groth16::{prepare_verifying_key, Parameters, PreparedVerifyingKey, VerifyingKey};
 use pairing::bls12_381::Bls12;

zcash_proofs/src/sapling/mod.rs

@@ -6,7 +6,7 @@ use sapling_crypto::jubjub::{
 mod prover;
 mod verifier;
 
-pub use self::prover::{CommitmentTreeWitness, SaplingProvingContext};
+pub use self::prover::SaplingProvingContext;
 pub use self::verifier::SaplingVerificationContext;
 
 // This function computes `value` in the exponent of the value commitment base

zcash_proofs/src/sapling/prover.rs

@@ -1,92 +1,22 @@
 use bellman::groth16::{
     create_random_proof, verify_proof, Parameters, PreparedVerifyingKey, Proof,
 };
-use byteorder::{LittleEndian, ReadBytesExt};
+use ff::Field;
-use ff::{Field, PrimeField, PrimeFieldRepr};
+use pairing::bls12_381::{Bls12, Fr};
-use pairing::bls12_381::{Bls12, Fr, FrRepr};
 use rand::{OsRng, Rand};
 use sapling_crypto::{
     circuit::{
         multipack,
-        sapling::{Output, Spend, TREE_DEPTH},
+        sapling::{Output, Spend},
     },
     jubjub::{edwards, fs::Fs, FixedGenerators, JubjubBls12, Unknown},
     primitives::{Diversifier, Note, PaymentAddress, ProofGenerationKey, ValueCommitment},
     redjubjub::{PrivateKey, PublicKey, Signature},
 };
+use zcash_primitives::{merkle_tree::CommitmentTreeWitness, sapling::Node};
 
 use super::compute_value_balance;
 
-/// A witness to a path from a postion in a particular Sapling commitment tree
-/// to the root of that tree.
-pub struct CommitmentTreeWitness {
-    auth_path: Vec<Option<(Fr, bool)>>,
-    position: u64,
-}
-
-impl CommitmentTreeWitness {
-    pub fn from_slice(mut witness: &[u8]) -> Result<Self, ()> {
-        // Skip the first byte, which should be "32" to signify the length of
-        // the following vector of Pedersen hashes.
-        assert_eq!(witness[0], TREE_DEPTH as u8);
-        witness = &witness[1..];
-
-        // Begin to construct the authentication path
-        let mut auth_path = vec![None; TREE_DEPTH];
-
-        // The vector works in reverse
-        for i in (0..TREE_DEPTH).rev() {
-            // skip length of inner vector
-            assert_eq!(witness[0], 32); // the length of a pedersen hash
-            witness = &witness[1..];
-
-            // Grab the sibling node at this depth in the tree
-            let mut sibling = [0u8; 32];
-            sibling.copy_from_slice(&witness[0..32]);
-            witness = &witness[32..];
-
-            // Sibling node should be an element of Fr
-            let sibling = match {
-                let mut repr = FrRepr::default();
-                repr.read_le(&sibling[..]).expect("length is 32 bytes");
-                Fr::from_repr(repr)
-            } {
-                Ok(p) => p,
-                Err(_) => return Err(()),
-            };
-
-            // Set the value in the auth path; we put false here
-            // for now (signifying the position bit) which we'll
-            // fill in later.
-            auth_path[i] = Some((sibling, false));
-        }
-
-        // Read the position from the witness
-        let position = witness
-            .read_u64::<LittleEndian>()
-            .expect("should have had index at the end");
-
-        // Given the position, let's finish constructing the authentication
-        // path
-        let mut tmp = position;
-        for i in 0..TREE_DEPTH {
-            auth_path[i].as_mut().map(|p| p.1 = (tmp & 1) == 1);
-
-            tmp >>= 1;
-        }
-
-        // The witness should be empty now; if it wasn't, the caller would
-        // have provided more information than they should have, indicating
-        // a bug downstream
-        assert_eq!(witness.len(), 0);
-
-        Ok(CommitmentTreeWitness {
-            auth_path,
-            position,
-        })
-    }
-}
-
 /// A context object for creating the Sapling components of a Zcash transaction.
 pub struct SaplingProvingContext {
     bsk: Fs,
@@ -113,7 +43,7 @@ impl SaplingProvingContext {
         ar: Fs,
         value: u64,
         anchor: Fr,
-        witness: CommitmentTreeWitness,
+        witness: CommitmentTreeWitness<Node>,
         proving_key: &Parameters<Bls12>,
         verifying_key: &PreparedVerifyingKey<Bls12>,
         params: &JubjubBls12,
@@ -182,7 +112,11 @@ impl SaplingProvingContext {
             payment_address: Some(payment_address),
             commitment_randomness: Some(rcm),
             ar: Some(ar),
-            auth_path: witness.auth_path,
+            auth_path: witness
+                .auth_path
+                .iter()
+                .map(|n| n.map(|(node, b)| (node.into(), b)))
+                .collect(),
             anchor: Some(anchor),
         };