Refactor sapling::tree::Node

zcash_primitives/CHANGELOG.md

@@ -52,6 +52,7 @@ and this library adheres to Rust's notion of
   - `note_encryption::SaplingDomain::new`
   - `note_encryption::Zip212Enforcement`
   - `prover::{SpendProver, OutputProver}`
+  - `tree::Node::{from_bytes, to_bytes}`
   - `value`:
     - `ValueCommitTrapdoor::from_bytes`
     - `impl Sub<TrapdoorSum> for TrapdoorSum`

zcash_primitives/src/merkle_tree.rs

@@ -23,6 +23,23 @@ pub trait HashSer {
     fn write<W: Write>(&self, writer: W) -> io::Result<()>;
 }
 
+impl HashSer for sapling::Node {
+    fn read<R: Read>(mut reader: R) -> io::Result<Self> {
+        let mut repr = [0u8; 32];
+        reader.read_exact(&mut repr)?;
+        Option::from(Self::from_bytes(repr)).ok_or_else(|| {
+            io::Error::new(
+                io::ErrorKind::InvalidData,
+                "Non-canonical encoding of Jubjub base field value.",
+            )
+        })
+    }
+
+    fn write<W: Write>(&self, mut writer: W) -> io::Result<()> {
+        writer.write_all(&self.to_bytes())
+    }
+}
+
 impl HashSer for MerkleHashOrchard {
     fn read<R: Read>(mut reader: R) -> io::Result<Self>
     where
@@ -815,7 +832,7 @@ mod tests {
         for i in 0..16 {
             let cmu = hex::decode(commitments[i]).unwrap();
 
-            let cmu = Node::new(cmu[..].try_into().unwrap());
+            let cmu = Node::from_bytes(cmu[..].try_into().unwrap()).unwrap();
 
             // Witness here
             witnesses.push((IncrementalWitness::from_tree(tree.clone()), last_cmu));

zcash_primitives/src/sapling/bundle.rs

@@ -631,7 +631,7 @@ pub mod testing {
 
     use crate::{
         sapling::{
-            note::ExtractedNoteCommitment,
+            note::testing::arb_cmu,
             value::{
                 testing::{arb_note_value_bounded, arb_trapdoor},
                 ValueCommitment, MAX_NOTE_VALUE,
@@ -691,9 +691,7 @@ pub mod testing {
         pub fn arb_output_description(n_outputs: usize)(
             value in arb_note_value_bounded(MAX_NOTE_VALUE.checked_div(n_outputs as u64).unwrap_or(0)),
             rcv in arb_trapdoor(),
-            cmu in vec(any::<u8>(), 64)
-                .prop_map(|v| <[u8;64]>::try_from(v.as_slice()).unwrap())
-                .prop_map(|v| bls12_381::Scalar::from_bytes_wide(&v)),
+            cmu in arb_cmu(),
             enc_ciphertext in vec(any::<u8>(), ENC_CIPHERTEXT_SIZE)
                 .prop_map(|v| <[u8; ENC_CIPHERTEXT_SIZE]>::try_from(v.as_slice()).unwrap()),
             epk in arb_extended_point(),
@@ -703,7 +701,6 @@ pub mod testing {
                 .prop_map(|v| <[u8; GROTH_PROOF_SIZE]>::try_from(v.as_slice()).unwrap()),
         ) -> OutputDescription<GrothProofBytes> {
             let cv = ValueCommitment::derive(value, rcv);
-            let cmu = ExtractedNoteCommitment::from_bytes(&cmu.to_bytes()).unwrap();
             OutputDescription {
                 cv,
                 cmu,

zcash_primitives/src/sapling/note.rs

@@ -152,11 +152,11 @@ impl Note {
 
 #[cfg(any(test, feature = "test-dependencies"))]
 pub(super) mod testing {
-    use proptest::prelude::*;
+    use proptest::{collection::vec, prelude::*};
 
     use super::{
         super::{testing::arb_payment_address, value::NoteValue},
-        Note, Rseed,
+        ExtractedNoteCommitment, Note, Rseed,
     };
 
     prop_compose! {
@@ -171,4 +171,14 @@ pub(super) mod testing {
             }
         }
     }
+
+    prop_compose! {
+        pub(crate) fn arb_cmu()(
+            cmu in vec(any::<u8>(), 64)
+                .prop_map(|v| <[u8;64]>::try_from(v.as_slice()).unwrap())
+                .prop_map(|v| bls12_381::Scalar::from_bytes_wide(&v)),
+        ) -> ExtractedNoteCommitment {
+            ExtractedNoteCommitment(cmu)
+        }
+    }
 }

zcash_primitives/src/sapling/note/commitment.rs

@@ -77,6 +77,10 @@ impl ExtractedNoteCommitment {
     pub fn to_bytes(self) -> [u8; 32] {
         self.0.to_repr()
     }
+
+    pub(crate) fn inner(&self) -> jubjub::Base {
+        self.0
+    }
 }
 
 impl From<NoteCommitment> for ExtractedNoteCommitment {

zcash_primitives/src/sapling/tree.rs

@@ -2,15 +2,14 @@ use bitvec::{order::Lsb0, view::AsBits};
 use group::{ff::PrimeField, Curve};
 use incrementalmerkletree::{Hashable, Level};
 use lazy_static::lazy_static;
+use subtle::CtOption;
 
 use std::fmt;
-use std::io::{self, Read, Write};
 
 use super::{
     note::ExtractedNoteCommitment,
     pedersen_hash::{pedersen_hash, Personalization},
 };
-use crate::merkle_tree::HashSer;
 
 pub const NOTE_COMMITMENT_TREE_DEPTH: u8 = 32;
 pub type CommitmentTree =
@@ -33,6 +32,10 @@ lazy_static! {
 
 /// Compute a parent node in the Sapling commitment tree given its two children.
 pub fn merkle_hash(depth: usize, lhs: &[u8; 32], rhs: &[u8; 32]) -> [u8; 32] {
+    merkle_hash_field(depth, lhs, rhs).to_repr()
+}
+
+fn merkle_hash_field(depth: usize, lhs: &[u8; 32], rhs: &[u8; 32]) -> jubjub::Base {
     let lhs = {
         let mut tmp = [false; 256];
         for (a, b) in tmp.iter_mut().zip(lhs.as_bits::<Lsb0>()) {
@@ -62,78 +65,65 @@ pub fn merkle_hash(depth: usize, lhs: &[u8; 32], rhs: &[u8; 32]) -> [u8; 32] {
     ))
     .to_affine()
     .get_u()
-    .to_repr()
 }
 
 /// A node within the Sapling commitment tree.
 #[derive(Clone, Copy, PartialEq, Eq)]
-pub struct Node {
-    pub(super) repr: [u8; 32],
-}
+pub struct Node(jubjub::Base);
 
 impl fmt::Debug for Node {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("Node")
-            .field("repr", &hex::encode(self.repr))
+            .field("repr", &hex::encode(self.0.to_bytes()))
             .finish()
     }
 }
 
 impl Node {
-    #[cfg(test)]
-    pub(crate) fn new(repr: [u8; 32]) -> Self {
-        Node { repr }
-    }
-
     /// Creates a tree leaf from the given Sapling note commitment.
     pub fn from_cmu(value: &ExtractedNoteCommitment) -> Self {
-        Node {
-            repr: value.to_bytes(),
-        }
+        Node(value.inner())
     }
 
     /// Constructs a new note commitment tree node from a [`bls12_381::Scalar`]
     pub fn from_scalar(cmu: bls12_381::Scalar) -> Self {
-        Self {
-            repr: cmu.to_repr(),
-        }
+        Self(cmu)
+    }
+
+    /// Parses a tree leaf from the bytes of a Sapling note commitment.
+    ///
+    /// Returns `None` if the provided bytes represent a non-canonical encoding.
+    pub fn from_bytes(bytes: [u8; 32]) -> CtOption<Self> {
+        jubjub::Base::from_repr(bytes).map(Self)
+    }
+
+    /// Returns the canonical byte representation of this node.
+    pub fn to_bytes(&self) -> [u8; 32] {
+        self.0.to_repr()
     }
 }
 
 impl Hashable for Node {
     fn empty_leaf() -> Self {
-        Node {
-            repr: UNCOMMITTED_SAPLING.to_repr(),
-        }
+        Node(*UNCOMMITTED_SAPLING)
     }
 
     fn combine(level: Level, lhs: &Self, rhs: &Self) -> Self {
-        Node {
-            repr: merkle_hash(level.into(), &lhs.repr, &rhs.repr),
-        }
+        Node(merkle_hash_field(
+            level.into(),
+            &lhs.0.to_bytes(),
+            &rhs.0.to_bytes(),
+        ))
     }
 
     fn empty_root(level: Level) -> Self {
         EMPTY_ROOTS[<usize>::from(level)]
     }
 }
 
-impl HashSer for Node {
-    fn read<R: Read>(mut reader: R) -> io::Result<Self> {
-        let mut repr = [0u8; 32];
-        reader.read_exact(&mut repr)?;
-        Ok(Node { repr })
-    }
-
-    fn write<W: Write>(&self, mut writer: W) -> io::Result<()> {
-        writer.write_all(self.repr.as_ref())
-    }
-}
-
 impl From<Node> for bls12_381::Scalar {
     fn from(node: Node) -> Self {
-        // Tree nodes should be in the prime field.
-        bls12_381::Scalar::from_repr(node.repr).unwrap()
+        node.0
     }
 }
 
@@ -142,12 +132,11 @@ pub(super) mod testing {
     use proptest::prelude::*;
 
     use super::Node;
+    use crate::sapling::note::testing::arb_cmu;
 
     prop_compose! {
-        pub fn arb_node()(value in prop::array::uniform32(prop::num::u8::ANY)) -> Node {
-            Node {
-                repr: value
-            }
+        pub fn arb_node()(cmu in arb_cmu()) -> Node {
+            Node::from_cmu(&cmu)
         }
     }
 }