From db92dd7de7d6f4bc3eb697cf10c0e09e0c57bb98 Mon Sep 17 00:00:00 2001
From: Rene Vergara <rene@vergara.network>
Date: Mon, 11 Mar 2024 15:23:29 -0500
Subject: [PATCH] Implement Sapling child key derivation

---
 librustzcash-wrapper/src/lib.rs | 18 +++----
 src/C/Zcash.chs                 |  1 +
 src/ZcashHaskell/Sapling.hs     | 10 ++--
 test/Spec.hs                    | 88 +++++++++++++++++----------------
 zcash-haskell.cabal             |  2 +-
 5 files changed, 62 insertions(+), 57 deletions(-)

diff --git a/librustzcash-wrapper/src/lib.rs b/librustzcash-wrapper/src/lib.rs
index 5eea232..c2b2607 100644
--- a/librustzcash-wrapper/src/lib.rs
+++ b/librustzcash-wrapper/src/lib.rs
@@ -26,6 +26,7 @@ use zip32;
 use zcash_primitives::{
     zip32::{
         Scope as SaplingScope,
+        ChildIndex,
         sapling_find_address,
         sapling::DiversifierKey
     },
@@ -628,21 +629,20 @@ pub extern "C" fn rust_wrapper_recover_seed(
 pub extern "C" fn rust_wrapper_sapling_spendingkey(
     iseed: *const u8,
     iseed_len: usize,
+    ix: u32,
     out: *mut u8,
     out_len: &mut usize
     ){
     let seed: Vec<u8> = marshall_from_haskell_var(iseed, iseed_len, RW);
-    if  ( seed.len() != 64 ) {
-        // invalid seed length 
-        marshall_to_haskell_var(&vec![0], out, out_len, RW);
-    } else {
-        // Obtain the ExtendedSpendingKey using the seed
-        // Returns a byte array (169 bytes)
-        let su8 = &seed;
-        let seedu8 : &[u8] = &su8;
-        let extsk: ExtendedSpendingKey = ExtendedSpendingKey::master(&seedu8); 
+    let su8 = &seed;
+    let seedu8 : &[u8] = &su8;
+    let extsk: ExtendedSpendingKey = ExtendedSpendingKey::master(&seedu8); 
+    if ix == 0 {
         let extsk_bytes = extsk.to_bytes().to_vec(); 
         marshall_to_haskell_var(&extsk_bytes, out, out_len, RW);
+    } else {
+        let child_sk = extsk.derive_child(ChildIndex::from_index(ix));
+        marshall_to_haskell_var(&child_sk.to_bytes().to_vec(), out, out_len, RW);
     }
 }
 
diff --git a/src/C/Zcash.chs b/src/C/Zcash.chs
index 1e32f93..fc60de0 100644
--- a/src/C/Zcash.chs
+++ b/src/C/Zcash.chs
@@ -136,6 +136,7 @@ import ZcashHaskell.Types
 
 {# fun unsafe rust_wrapper_sapling_spendingkey as rustWrapperSaplingSpendingkey
      { toBorshVar* `BS.ByteString'&
+     , `Word32'
      , getVarBuffer `Buffer (BS.ByteString)'&
      }
   -> `()'
diff --git a/src/ZcashHaskell/Sapling.hs b/src/ZcashHaskell/Sapling.hs
index 0a25661..d4eaa71 100644
--- a/src/ZcashHaskell/Sapling.hs
+++ b/src/ZcashHaskell/Sapling.hs
@@ -95,13 +95,15 @@ instance FromJSON RawTxResponse where
           pure $ RawTxResponse i h (getShieldedOutputs h) a ht c b
 
 -- | Attempts to obtain a sapling SpendingKey using a HDSeed
-genSaplingSpendingKey :: Seed -> Maybe SaplingSpendingKey
-genSaplingSpendingKey seed = do
-  if BS.length res == 196
+genSaplingSpendingKey :: Seed -> Int -> Maybe SaplingSpendingKey
+genSaplingSpendingKey seed i = do
+  if BS.length res == 169
     then Just res
     else Nothing
   where
-    res = withPureBorshVarBuffer (rustWrapperSaplingSpendingkey seed)
+    res =
+      withPureBorshVarBuffer
+        (rustWrapperSaplingSpendingkey seed (fromIntegral i))
 
 -- | Attempts to generate a sapling Payment Address using an ExtendedSpendingKey and a Diversifier Index 
 genSaplingPaymentAddress :: SaplingSpendingKey -> Int -> Maybe SaplingReceiver
diff --git a/test/Spec.hs b/test/Spec.hs
index e6d6bc2..0762410 100644
--- a/test/Spec.hs
+++ b/test/Spec.hs
@@ -20,6 +20,7 @@
 {-# LANGUAGE TypeSynonymInstances #-}
 
 import C.Zcash (rustWrapperUADecode)
+import Control.Exception (throwIO)
 import Control.Monad.IO.Class (liftIO)
 import Data.Aeson
 import Data.Bool (Bool(True))
@@ -61,6 +62,7 @@ import ZcashHaskell.Types
   , Phrase(..)
   , RawData(..)
   , RawTxResponse(..)
+  , Seed(..)
   , ShieldedOutput(..)
   , UnifiedAddress(..)
   , UnifiedFullViewingKey(..)
@@ -466,12 +468,18 @@ main = do
           msg `shouldBe`
             "Hello World!\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL"
     describe "Wallet seed phrase" $ do
-      prop "Generated phrases are valid" prop_PhraseLength
-      prop "Derived seeds are valid" prop_SeedLength
-      prop "Orchard spending keys are valid" $
-        forAll genOrcArgs $ \(c, i, _) -> prop_OrchardSpendingKey c i
-      prop "Orchard receivers are valid" $
-        forAll genOrcArgs $ \(c, i, j) -> prop_OrchardReceiver c i j
+      prop "Generated phrases are valid" $ again prop_PhraseLength
+      prop "Derived seeds are valid" $ again prop_SeedLength
+    before getSeed $
+      describe "Optimized spending key tests" $ do
+        it "Sapling spending keys are valid" $ \s ->
+          property $ prop_SaplingSpendingKey s
+        it "Sapling receivers are valid" $ \s ->
+          property $ prop_SaplingReceiver s
+        it "Orchard spending keys are valid" $ \s ->
+          property $ prop_OrchardSpendingKey s
+        it "Orchard receivers are valid" $ \s ->
+          property $ prop_OrchardReceiver s
     describe "Address tests" $ do
       it "Encode transparent" $ do
         let ua =
@@ -481,7 +489,7 @@ main = do
                 Nothing -> "Bad UA"
                 Just u -> maybe "No transparent" encodeTransparent $ t_rec u
         msg `shouldBe` "t1LPWuQnjCRH7JAeEErSXKixcUteLJRJjKD"
-    describe "Transparent Private and Publicc Key Generation" $ do
+    describe "Transparent Private and Public Key Generation" $ do
       it "Obtain a transparent extended private key from HDSeed" $ do
         let hdseed =
               [ 206
@@ -577,55 +585,38 @@ main = do
         xtpubk <- xtpubkIO
           ---print $ show xtpubk
         xtpubk `shouldBe` testpbk
-    describe "Sapling SpendingKey test" $ do
-      it "Generate Sapling spending key" $ do
-        p <- generateWalletSeedPhrase
-        let s = getWalletSeed p
-        genSaplingSpendingKey <$> s `shouldNotBe` Nothing
-    describe "Sapling Payment Address generation test" $ do
-      it "Call genSaplingPaymentAddress" $ do
-        p <- generateWalletSeedPhrase
-        let s = getWalletSeed p
-        genSaplingPaymentAddress (fromMaybe "" s) 0 `shouldNotBe` Nothing
-      prop "Sapling receivers are valid" $
-        forAll genSapArgs $ \i -> prop_SaplingReceiver i
 
 -- | Properties
-prop_PhraseLength :: Int -> Property
-prop_PhraseLength i =
+prop_PhraseLength :: Property
+prop_PhraseLength =
   ioProperty $ do
     p <- generateWalletSeedPhrase
     return $ BS.length p >= 95
 
-prop_SeedLength :: Int -> Property
-prop_SeedLength i =
+prop_SeedLength :: Property
+prop_SeedLength =
   ioProperty $ do
     p <- generateWalletSeedPhrase
     let s = getWalletSeed p
     return $ maybe 0 BS.length s === 64
 
-prop_OrchardSpendingKey :: CoinType -> Int -> Property
-prop_OrchardSpendingKey c i =
-  ioProperty $ do
-    p <- generateWalletSeedPhrase
-    let s = getWalletSeed p
-    return $ genOrchardSpendingKey (fromMaybe "" s) c i =/= Nothing
+prop_OrchardSpendingKey :: Seed -> CoinType -> NonNegative Int -> Property
+prop_OrchardSpendingKey s c (NonNegative i) =
+  genOrchardSpendingKey s c i =/= Nothing
 
-prop_OrchardReceiver :: CoinType -> Int -> Int -> Property
-prop_OrchardReceiver c i j =
-  ioProperty $ do
-    p <- generateWalletSeedPhrase
-    let s = getWalletSeed p
-    let sk = genOrchardSpendingKey (fromMaybe "" s) c i
-    return $ genOrchardReceiver j (fromMaybe "" sk) =/= Nothing
+prop_OrchardReceiver ::
+     Seed -> CoinType -> NonNegative Int -> NonNegative Int -> Property
+prop_OrchardReceiver s c (NonNegative i) (NonNegative j) =
+  genOrchardReceiver j (fromMaybe "" $ genOrchardSpendingKey s c i) =/= Nothing
 
-prop_SaplingReceiver :: Int -> Property
-prop_SaplingReceiver i =
-  ioProperty $ do
-    p <- generateWalletSeedPhrase
-    let s = getWalletSeed p
-    let sk = genSaplingSpendingKey (fromMaybe "" s)
-    return $ genSaplingPaymentAddress (fromMaybe "" sk) i =/= Nothing
+prop_SaplingSpendingKey :: Seed -> NonNegative Int -> Property
+prop_SaplingSpendingKey s (NonNegative i) =
+  genSaplingSpendingKey s i =/= Nothing
+
+prop_SaplingReceiver :: Seed -> NonNegative Int -> NonNegative Int -> Property
+prop_SaplingReceiver s (NonNegative i) (NonNegative j) =
+  genSaplingPaymentAddress (fromMaybe "" $ genSaplingSpendingKey s j) i =/=
+  Nothing
 
 -- | Generators
 genOrcArgs :: Gen (CoinType, Int, Int)
@@ -637,4 +628,15 @@ genOrcArgs = do
 
 genSapArgs :: Gen Int
 genSapArgs = choose (1, 50)
+
+getSeed :: IO Seed
+getSeed = do
+  p <- generateWalletSeedPhrase
+  let s = getWalletSeed p
+  case s of
+    Nothing -> throwIO $ userError "Couldn't generate seed"
+    Just s' -> return s'
+
 -- | Arbitrary instances
+instance Arbitrary CoinType where
+  arbitrary = elements [MainNetCoin, TestNetCoin, RegTestNetCoin]
diff --git a/zcash-haskell.cabal b/zcash-haskell.cabal
index 0bee561..19e047d 100644
--- a/zcash-haskell.cabal
+++ b/zcash-haskell.cabal
@@ -5,7 +5,7 @@ cabal-version: 3.0
 -- see: https://github.com/sol/hpack
 
 name:           zcash-haskell
-version:        0.4.4.1
+version:        0.4.4.2
 synopsis:       Utilities to interact with the Zcash blockchain
 description:    Please see the README on the repo at <https://git.vergara.tech/Vergara_Tech/zcash-haskell#readme>
 category:       Blockchain