Allow Transaction Builder to be called with inconsistent keys.

murisi · murisi · commit f2b0cae3e495 · 2024-09-10T12:21:10.000+02:00
diff --git a/masp_primitives/src/lib.rs b/masp_primitives/src/lib.rs
@@ -34,3 +34,15 @@ pub use num_traits;
 
 #[cfg(test)]
 mod test_vectors;
+
+#[cfg(not(feature = "arbitrary"))]
+pub trait MaybeArbitrary<'a> {}
+
+#[cfg(not(feature = "arbitrary"))]
+impl<'a, T> MaybeArbitrary<'a> for T {}
+
+#[cfg(feature = "arbitrary")]
+pub trait MaybeArbitrary<'a>: arbitrary::Arbitrary<'a> {}
+
+#[cfg(feature = "arbitrary")]
+impl<'a, T: for<'b> arbitrary::Arbitrary<'b>> MaybeArbitrary<'a> for T {}
diff --git a/masp_primitives/src/transaction.rs b/masp_primitives/src/transaction.rs
@@ -34,9 +34,11 @@ use self::{
     },
     txid::{to_txid, BlockTxCommitmentDigester, TxIdDigester},
 };
+use crate::MaybeArbitrary;
 use borsh::schema::add_definition;
 use borsh::schema::Fields;
 use borsh::schema::{Declaration, Definition};
+use std::marker::PhantomData;
 use std::ops::RangeInclusive;
 
 #[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
@@ -195,23 +197,17 @@ impl BorshSchema for TxVersion {
 
 /// Authorization state for a bundle of transaction data.
 pub trait Authorization {
-    #[cfg(not(feature = "arbitrary"))]
-    type TransparentAuth: transparent::Authorization + PartialEq + BorshDeserialize + BorshSerialize;
-    #[cfg(not(feature = "arbitrary"))]
-    type SaplingAuth: sapling::Authorization + PartialEq + BorshDeserialize + BorshSerialize;
-
-    #[cfg(feature = "arbitrary")]
     type TransparentAuth: transparent::Authorization
         + PartialEq
         + BorshDeserialize
         + BorshSerialize
-        + for<'a> arbitrary::Arbitrary<'a>;
-    #[cfg(feature = "arbitrary")]
+        + for<'a> MaybeArbitrary<'a>;
+
     type SaplingAuth: sapling::Authorization
         + PartialEq
         + BorshDeserialize
         + BorshSerialize
-        + for<'a> arbitrary::Arbitrary<'a>;
+        + for<'a> MaybeArbitrary<'a>;
 }
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
 pub struct Unproven;
@@ -224,11 +220,13 @@ impl Authorization for Authorized {
     type SaplingAuth = sapling::Authorized;
 }
 
-pub struct Unauthorized;
+pub struct Unauthorized<K: crate::zip32::ExtendedKey>(PhantomData<K>);
 
-impl Authorization for Unauthorized {
+impl<K: crate::zip32::ExtendedKey + PartialEq + Clone + Debug + for<'a> MaybeArbitrary<'a>>
+    Authorization for Unauthorized<K>
+{
     type TransparentAuth = transparent::builder::Unauthorized;
-    type SaplingAuth = sapling::builder::Unauthorized;
+    type SaplingAuth = sapling::builder::Unauthorized<K>;
 }
 
 /// A MASP transaction.
diff --git a/masp_primitives/src/transaction/builder.rs b/masp_primitives/src/transaction/builder.rs
@@ -30,7 +30,8 @@ use crate::{
         txid::TxIdDigester,
         Transaction, TransactionData, TransparentAddress, TxVersion, Unauthorized,
     },
-    zip32::ExtendedSpendingKey,
+    zip32::{ExtendedKey, ExtendedSpendingKey},
+    MaybeArbitrary,
 };
 
 #[cfg(feature = "transparent-inputs")]
@@ -168,7 +169,11 @@ impl<P, K, N> Builder<P, K, N> {
     }
 }
 
-impl<P: consensus::Parameters> Builder<P> {
+impl<
+        P: consensus::Parameters,
+        K: ExtendedKey + std::fmt::Debug + Clone + PartialEq + for<'a> MaybeArbitrary<'a>,
+    > Builder<P, K>
+{
     /// Creates a new `Builder` targeted for inclusion in the block with the given height,
     /// using default values for general transaction fields and the default OS random.
     ///
@@ -181,12 +186,16 @@ impl<P: consensus::Parameters> Builder<P> {
     }
 }
 
-impl<P: consensus::Parameters> Builder<P> {
+impl<
+        P: consensus::Parameters,
+        K: ExtendedKey + std::fmt::Debug + Clone + PartialEq + for<'a> MaybeArbitrary<'a>,
+    > Builder<P, K>
+{
     /// Common utility function for builder construction.
     ///
     /// WARNING: THIS MUST REMAIN PRIVATE AS IT ALLOWS CONSTRUCTION
     /// OF BUILDERS WITH NON-CryptoRng RNGs
-    fn new_internal(params: P, target_height: BlockHeight) -> Builder<P> {
+    fn new_internal(params: P, target_height: BlockHeight) -> Builder<P, K> {
         Builder {
             params: params.clone(),
             target_height,
@@ -203,7 +212,7 @@ impl<P: consensus::Parameters> Builder<P> {
     /// paths for previous Sapling notes.
     pub fn add_sapling_spend(
         &mut self,
-        extsk: ExtendedSpendingKey,
+        extsk: K,
         diversifier: Diversifier,
         note: Note,
         merkle_path: MerklePath<Node>,
@@ -347,7 +356,7 @@ impl<P: consensus::Parameters> Builder<P> {
             )
             .map_err(Error::SaplingBuild)?;
 
-        let unauthed_tx: TransactionData<Unauthorized> = TransactionData {
+        let unauthed_tx: TransactionData<Unauthorized<K>> = TransactionData {
             version,
             consensus_branch_id: BranchId::for_height(&self.params, self.target_height),
             lock_time: 0,
diff --git a/masp_primitives/src/transaction/components/sapling.rs b/masp_primitives/src/transaction/components/sapling.rs
@@ -25,6 +25,7 @@ use crate::{
         redjubjub::{self, PublicKey, Signature},
         Nullifier,
     },
+    MaybeArbitrary,
 };
 
 use super::{amount::I128Sum, GROTH_PROOF_SIZE};
@@ -35,15 +36,8 @@ pub mod builder;
 pub mod fees;
 
 pub trait Authorization: Debug {
-    #[cfg(not(feature = "arbitrary"))]
-    type Proof: Clone + Debug + PartialEq + Hash;
-    #[cfg(not(feature = "arbitrary"))]
-    type AuthSig: Clone + Debug + PartialEq;
-
-    #[cfg(feature = "arbitrary")]
-    type Proof: Clone + Debug + PartialEq + Hash + for<'a> arbitrary::Arbitrary<'a>;
-    #[cfg(feature = "arbitrary")]
-    type AuthSig: Clone + Debug + PartialEq + for<'a> arbitrary::Arbitrary<'a>;
+    type Proof: Clone + Debug + PartialEq + Hash + for<'a> MaybeArbitrary<'a>;
+    type AuthSig: Clone + Debug + PartialEq + for<'a> MaybeArbitrary<'a>;
 }
 
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
diff --git a/masp_primitives/src/transaction/components/sapling/builder.rs b/masp_primitives/src/transaction/components/sapling/builder.rs
@@ -8,6 +8,7 @@ use ff::PrimeField;
 use group::GroupEncoding;
 use rand::{seq::SliceRandom, CryptoRng, RngCore};
 
+use crate::MaybeArbitrary;
 use crate::{
     asset_type::AssetType,
     consensus::{self, BlockHeight},
@@ -33,15 +34,17 @@ use crate::{
             },
         },
     },
-    zip32::ExtendedSpendingKey,
+    zip32::{ExtendedKey, ExtendedSpendingKey},
 };
 use borsh::schema::add_definition;
 use borsh::schema::Declaration;
 use borsh::schema::Definition;
 use borsh::schema::Fields;
 use borsh::{BorshDeserialize, BorshSchema, BorshSerialize};
 use std::collections::BTreeMap;
+use std::fmt::Debug;
 use std::io::Write;
+use std::marker::PhantomData;
 
 /// A subset of the parameters necessary to build a transaction
 pub trait BuildParams {
@@ -777,19 +780,22 @@ impl<P: BorshDeserialize, Key: BorshDeserialize> BorshDeserialize for SaplingBui
 
 #[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
 #[derive(Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
-pub struct Unauthorized {
+pub struct Unauthorized<K: ExtendedKey> {
     tx_metadata: SaplingMetadata,
+    phantom: PhantomData<K>,
 }
 
-impl std::fmt::Debug for Unauthorized {
+impl<K: ExtendedKey> std::fmt::Debug for Unauthorized<K> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
         write!(f, "Unauthorized")
     }
 }
 
-impl Authorization for Unauthorized {
+impl<K: ExtendedKey + Clone + Debug + PartialEq + for<'a> MaybeArbitrary<'a>> Authorization
+    for Unauthorized<K>
+{
     type Proof = GrothProofBytes;
-    type AuthSig = SpendDescriptionInfo;
+    type AuthSig = SpendDescriptionInfo<K>;
 }
 
 impl<P, K> SaplingBuilder<P, K> {
@@ -826,14 +832,18 @@ impl<P, K> SaplingBuilder<P, K> {
     }
 }
 
-impl<P: consensus::Parameters> SaplingBuilder<P> {
+impl<
+        P: consensus::Parameters,
+        K: ExtendedKey + Debug + Clone + PartialEq + for<'a> MaybeArbitrary<'a>,
+    > SaplingBuilder<P, K>
+{
     /// Adds a Sapling note to be spent in this transaction.
     ///
     /// Returns an error if the given Merkle path does not have the same anchor as the
     /// paths for previous Sapling notes.
     pub fn add_spend(
         &mut self,
-        extsk: ExtendedSpendingKey,
+        extsk: K,
         diversifier: Diversifier,
         note: Note,
         merkle_path: MerklePath<Node>,
@@ -922,7 +932,7 @@ impl<P: consensus::Parameters> SaplingBuilder<P> {
         bparams: &mut impl BuildParams,
         target_height: BlockHeight,
         progress_notifier: Option<&Sender<Progress>>,
-    ) -> Result<Option<Bundle<Unauthorized>>, Error> {
+    ) -> Result<Option<Bundle<Unauthorized<K>>>, Error> {
         // Record initial positions of spends and outputs
         let value_balance = self.value_balance();
         let params = self.params;
@@ -961,7 +971,8 @@ impl<P: consensus::Parameters> SaplingBuilder<P> {
         let mut progress = 0u32;
 
         // Create Sapling SpendDescriptions
-        let shielded_spends: Vec<SpendDescription<Unauthorized>> = if !indexed_spends.is_empty() {
+        let shielded_spends: Vec<SpendDescription<Unauthorized<K>>> = if !indexed_spends.is_empty()
+        {
             let anchor = self
                 .spend_anchor
                 .expect("MASP Spend anchor must be set if MASP spends are present.");
@@ -970,7 +981,10 @@ impl<P: consensus::Parameters> SaplingBuilder<P> {
                 .into_iter()
                 .enumerate()
                 .map(|(i, (pos, spend))| {
-                    let proof_generation_key = spend.extsk.expsk.proof_generation_key();
+                    let proof_generation_key = spend
+                        .extsk
+                        .to_proof_generation_key()
+                        .expect("Proof generation key must be known for each MASP spend.");
 
                     let nullifier = spend.note.nf(
                         &proof_generation_key.to_viewing_key().nk,
@@ -1172,15 +1186,20 @@ impl<P: consensus::Parameters> SaplingBuilder<P> {
                 shielded_converts,
                 shielded_outputs,
                 value_balance,
-                authorization: Unauthorized { tx_metadata },
+                authorization: Unauthorized {
+                    tx_metadata,
+                    phantom: PhantomData,
+                },
             })
         };
 
         Ok(bundle)
     }
 }
 
-impl SpendDescription<Unauthorized> {
+impl<K: ExtendedKey + Debug + Clone + PartialEq + for<'a> MaybeArbitrary<'a>>
+    SpendDescription<Unauthorized<K>>
+{
     pub fn apply_signature(&self, spend_auth_sig: Signature) -> SpendDescription<Authorized> {
         SpendDescription {
             cv: self.cv,
@@ -1193,7 +1212,9 @@ impl SpendDescription<Unauthorized> {
     }
 }
 
-impl Bundle<Unauthorized> {
+impl<K: ExtendedKey + Debug + Clone + PartialEq + for<'a> MaybeArbitrary<'a>>
+    Bundle<Unauthorized<K>>
+{
     pub fn apply_signatures<Pr: TxProver, R: RngCore, S: BuildParams>(
         self,
         prover: &Pr,
@@ -1214,7 +1235,7 @@ impl Bundle<Unauthorized> {
                     .enumerate()
                     .map(|(i, spend)| {
                         spend.apply_signature(spend_sig_internal(
-                            PrivateKey(spend.spend_auth_sig.extsk.expsk.ask),
+                            PrivateKey(spend.spend_auth_sig.extsk.to_spending_key().expect("Spend authorization key must be known for each MASP spend.").expsk.ask),
                             bparams.spend_alpha(i),
                             sighash_bytes,
                             rng,
diff --git a/masp_primitives/src/transaction/components/transparent.rs b/masp_primitives/src/transaction/components/transparent.rs
@@ -6,6 +6,7 @@ use std::io::{self, Read, Write};
 
 use crate::asset_type::AssetType;
 use crate::transaction::TransparentAddress;
+use crate::MaybeArbitrary;
 use borsh::schema::add_definition;
 use borsh::schema::Declaration;
 use borsh::schema::Definition;
@@ -18,11 +19,7 @@ pub mod builder;
 pub mod fees;
 
 pub trait Authorization: fmt::Debug {
-    #[cfg(not(feature = "arbitrary"))]
-    type TransparentSig: fmt::Debug + Clone + PartialEq;
-
-    #[cfg(feature = "arbitrary")]
-    type TransparentSig: fmt::Debug + Clone + PartialEq + for<'a> arbitrary::Arbitrary<'a>;
+    type TransparentSig: fmt::Debug + Clone + PartialEq + for<'a> MaybeArbitrary<'a>;
 }
 
 #[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
diff --git a/masp_primitives/src/zip32.rs b/masp_primitives/src/zip32.rs
@@ -19,8 +19,8 @@ use borsh::{BorshDeserialize, BorshSerialize};
 #[deprecated(note = "Please use the types exported from the `zip32::sapling` module instead.")]
 pub use sapling::{
     sapling_address, sapling_default_address, sapling_derive_internal_fvk, sapling_find_address,
-    DiversifiableFullViewingKey, ExtendedFullViewingKey, ExtendedSpendingKey,
-    ZIP32_SAPLING_FVFP_PERSONALIZATION, ZIP32_SAPLING_INT_PERSONALIZATION,
+    DiversifiableFullViewingKey, ExtendedFullViewingKey, ExtendedKey, ExtendedSpendingKey,
+    PseudoExtendedKey, ZIP32_SAPLING_FVFP_PERSONALIZATION, ZIP32_SAPLING_INT_PERSONALIZATION,
     ZIP32_SAPLING_MASTER_PERSONALIZATION,
 };
 use std::io::{Read, Write};
diff --git a/masp_primitives/src/zip32/sapling.rs b/masp_primitives/src/zip32/sapling.rs
