Make the authorization mappers provide positional information.

anoma · Sep 4, 2024 · e6451ec · e6451ec
1 parent 9bc6dec
commit e6451ec
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 13 deletions.
diff --git a/masp_primitives/src/transaction/components/sapling.rs b/masp_primitives/src/transaction/components/sapling.rs
@@ -66,8 +66,8 @@ impl Authorization for Authorized {
 }
 
 pub trait MapAuth<A: Authorization, B: Authorization> {
-    fn map_proof(&self, p: A::Proof) -> B::Proof;
-    fn map_auth_sig(&self, s: A::AuthSig) -> B::AuthSig;
+    fn map_proof(&self, p: A::Proof, pos: usize) -> B::Proof;
+    fn map_auth_sig(&self, s: A::AuthSig, pos: usize) -> B::AuthSig;
     fn map_authorization(&self, a: A) -> B;
 }
 
@@ -81,13 +81,15 @@ impl MapAuth<Authorized, Authorized> for () {
     fn map_proof(
         &self,
         p: <Authorized as Authorization>::Proof,
+        _pos: usize,
     ) -> <Authorized as Authorization>::Proof {
         p
     }
 
     fn map_auth_sig(
         &self,
         s: <Authorized as Authorization>::AuthSig,
+        _pos: usize,
     ) -> <Authorized as Authorization>::AuthSig {
         s
     }
@@ -119,34 +121,37 @@ impl<A: Authorization + PartialEq + BorshSerialize + BorshDeserialize> Bundle<A>
             shielded_spends: self
                 .shielded_spends
                 .into_iter()
-                .map(|d| SpendDescription {
+                .enumerate()
+                .map(|(pos, d)| SpendDescription {
                     cv: d.cv,
                     anchor: d.anchor,
                     nullifier: d.nullifier,
                     rk: d.rk,
-                    zkproof: f.map_proof(d.zkproof),
-                    spend_auth_sig: f.map_auth_sig(d.spend_auth_sig),
+                    zkproof: f.map_proof(d.zkproof, pos),
+                    spend_auth_sig: f.map_auth_sig(d.spend_auth_sig, pos),
                 })
                 .collect(),
             shielded_converts: self
                 .shielded_converts
                 .into_iter()
-                .map(|c| ConvertDescription {
+                .enumerate()
+                .map(|(pos, c)| ConvertDescription {
                     cv: c.cv,
                     anchor: c.anchor,
-                    zkproof: f.map_proof(c.zkproof),
+                    zkproof: f.map_proof(c.zkproof, pos),
                 })
                 .collect(),
             shielded_outputs: self
                 .shielded_outputs
                 .into_iter()
-                .map(|o| OutputDescription {
+                .enumerate()
+                .map(|(pos, o)| OutputDescription {
                     cv: o.cv,
                     cmu: o.cmu,
                     ephemeral_key: o.ephemeral_key,
                     enc_ciphertext: o.enc_ciphertext,
                     out_ciphertext: o.out_ciphertext,
-                    zkproof: f.map_proof(o.zkproof),
+                    zkproof: f.map_proof(o.zkproof, pos),
                 })
                 .collect(),
             value_balance: self.value_balance,

diff --git a/masp_primitives/src/transaction/components/transparent.rs b/masp_primitives/src/transaction/components/transparent.rs
@@ -34,10 +34,30 @@ impl Authorization for Authorized {
 }
 
 pub trait MapAuth<A: Authorization, B: Authorization> {
-    fn map_script_sig(&self, s: A::TransparentSig) -> B::TransparentSig;
+    fn map_script_sig(&self, s: A::TransparentSig, pos: usize) -> B::TransparentSig;
     fn map_authorization(&self, s: A) -> B;
 }
 
+/// The identity map.
+///
+/// This can be used with [`TransactionData::map_authorization`] when you want to map the
+/// authorization of a subset of the transaction's bundles.
+///
+/// [`TransactionData::map_authorization`]: crate::transaction::TransactionData::map_authorization
+impl MapAuth<Authorized, Authorized> for () {
+    fn map_script_sig(
+        &self,
+        s: <Authorized as Authorization>::TransparentSig,
+        _pos: usize,
+    ) -> <Authorized as Authorization>::TransparentSig {
+        s
+    }
+
+    fn map_authorization(&self, a: Authorized) -> Authorized {
+        a
+    }
+}
+
 #[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
 #[derive(Debug, Clone, PartialEq)]
 pub struct Bundle<A: Authorization> {
@@ -52,10 +72,11 @@ impl<A: Authorization> Bundle<A> {
             vin: self
                 .vin
                 .into_iter()
-                .map(|txin| TxIn {
+                .enumerate()
+                .map(|(pos, txin)| TxIn {
                     asset_type: txin.asset_type,
                     address: txin.address,
-                    transparent_sig: f.map_script_sig(txin.transparent_sig),
+                    transparent_sig: f.map_script_sig(txin.transparent_sig, pos),
                     value: txin.value,
                 })
                 .collect(),

diff --git a/masp_primitives/src/zip32/sapling.rs b/masp_primitives/src/zip32/sapling.rs
@@ -934,7 +934,10 @@ impl PseudoExtendedSpendingKey {
     /// Construct a pseudo extended spending key from an extended spending key
     #[allow(deprecated)]
     pub fn from_spending_key(xsk: ExtendedSpendingKey) -> Self {
-        Self { xfvk: xsk.to_extended_full_viewing_key(), xsk }
+        Self {
+            xfvk: xsk.to_extended_full_viewing_key(),
+            xsk,
+        }
     }
     /// Construct a pseudo extended spending key from an extended full viewing key
     pub fn from_viewing_key(xfvk: ExtendedFullViewingKey) -> Self {