From 846c8971fe8c5dcf29cb9898ce3212f5a793a1cd Mon Sep 17 00:00:00 2001
From: Rene Vergara <rene@vergara.network>
Date: Wed, 23 Aug 2023 15:19:31 -0500
Subject: [PATCH] Update Sapling decoding functions

---
 librustzcash-wrapper/src/lib.rs | 23 ++++++++++++-----------
 src/C/Zcash.chs                 |  2 +-
 src/ZcashHaskell/Sapling.hs     |  6 +++---
 3 files changed, 16 insertions(+), 15 deletions(-)

diff --git a/librustzcash-wrapper/src/lib.rs b/librustzcash-wrapper/src/lib.rs
index 620a206..bd0e0e3 100644
--- a/librustzcash-wrapper/src/lib.rs
+++ b/librustzcash-wrapper/src/lib.rs
@@ -1,7 +1,10 @@
 
 use std::{
     marker::PhantomData,
-    io::Write,
+    io::{
+        Write,
+        Cursor
+    },
     fmt::{Debug, Display, Formatter}
 };
 
@@ -34,7 +37,9 @@ use zcash_primitives::{
         PaymentAddress,
         note::ExtractedNoteCommitment as SaplingExtractedNoteCommitment
     },
+    transaction::Transaction,
     consensus::{
+        BranchId::Nu5,
         MainNetwork,
         BlockHeight
     }
@@ -325,22 +330,18 @@ pub extern "C" fn rust_wrapper_sapling_note_decrypt(
     out_len: &mut usize
     ){
     let evk: Vec<u8> = marshall_from_haskell_var(key, key_len, RW);
-    let note_input: HshieldedOutput = marshall_from_haskell_var(note,note_len,RW);
+    let note_input: Vec<u8> = marshall_from_haskell_var(note,note_len,RW);
+    let mut note_reader = Cursor::new(note_input);
     let svk = ExtendedFullViewingKey::read(&*evk);
     match svk {
         Ok(k) => {
             let domain = SaplingDomain::for_height(MainNetwork, BlockHeight::from_u32(2000000));
-            let mut action_bytes = vec![0];
-            action_bytes.extend(&note_input.cv);
-            action_bytes.extend(&note_input.cmu);
-            action_bytes.extend(&note_input.eph_key);
-            action_bytes.extend(&note_input.enc_txt);
-            action_bytes.extend(&note_input.out_txt);
-            action_bytes.extend(&note_input.proof);
-            let action2 = OutputDescription::read(&mut action_bytes.as_slice()).unwrap();
+            let action2: Transaction = Transaction::read(&mut note_reader, Nu5).unwrap();
+            let bundle = action2.sapling_bundle().unwrap();
+            let sh_out = bundle.shielded_outputs();
             let fvk = k.to_diversifiable_full_viewing_key().to_ivk(SaplingScope::External);
             let pivk = SaplingPreparedIncomingViewingKey::new(&fvk);
-            let result = zcash_note_encryption::try_note_decryption(&domain, &pivk, &action2);
+            let result = zcash_note_encryption::try_note_decryption(&domain, &pivk, &sh_out[0]);
             match result {
                 Some((n, r, m)) => {
                     let hn = Hnote {note: n.value().inner(), recipient: r.to_bytes().to_vec(), memo: m.as_slice().to_vec() };
diff --git a/src/C/Zcash.chs b/src/C/Zcash.chs
index 3f195fe..76e7889 100644
--- a/src/C/Zcash.chs
+++ b/src/C/Zcash.chs
@@ -74,7 +74,7 @@ import ZcashHaskell.Types
 
 {# fun unsafe rust_wrapper_sapling_note_decrypt as rustWrapperSaplingNoteDecode
     { toBorshVar* `BS.ByteString'&
-    , toBorshVar* `ShieldedOutput'&
+    , toBorshVar* `BS.ByteString'&
     , getVarBuffer `Buffer DecodedNote'&
     }
   -> `()'
diff --git a/src/ZcashHaskell/Sapling.hs b/src/ZcashHaskell/Sapling.hs
index e0016bd..a0a4716 100644
--- a/src/ZcashHaskell/Sapling.hs
+++ b/src/ZcashHaskell/Sapling.hs
@@ -8,7 +8,7 @@ import C.Zcash
   )
 import qualified Data.ByteString as BS
 import Foreign.Rust.Marshall.Variable (withPureBorshVarBuffer)
-import ZcashHaskell.Types (DecodedNote(..), ShieldedOutput)
+import ZcashHaskell.Types (DecodedNote(..), ShieldedOutput(..))
 
 -- | Check if given bytesting is a valid encoded shielded address
 isValidShieldedAddress :: BS.ByteString -> Bool
@@ -22,8 +22,8 @@ isValidSaplingViewingKey = rustWrapperSaplingVkDecode
 matchSaplingAddress :: BS.ByteString -> BS.ByteString -> Bool
 matchSaplingAddress = rustWrapperSaplingCheck
 
--- | Attempt to decode the given Sapling raw output with the given Sapling viewing key
-decodeSaplingOutput :: BS.ByteString -> ShieldedOutput -> Maybe DecodedNote
+-- | Attempt to decode the given raw tx with the given Sapling viewing key
+decodeSaplingOutput :: BS.ByteString -> BS.ByteString -> Maybe DecodedNote
 decodeSaplingOutput key out =
   case a_value decodedAction of
     0 -> Nothing