From d2619123c0ea4476939dc6a14796c992d6be6e50 Mon Sep 17 00:00:00 2001
From: Rene Vergara <rene@vergara.network>
Date: Thu, 14 Mar 2024 11:30:54 -0500
Subject: [PATCH] Improve type safety for Sapling

---
 src/ZcashHaskell/Orchard.hs |  4 ++--
 src/ZcashHaskell/Sapling.hs | 13 ++++++++-----
 src/ZcashHaskell/Types.hs   |  7 ++++++-
 3 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/src/ZcashHaskell/Orchard.hs b/src/ZcashHaskell/Orchard.hs
index 5dd000f..af68603 100644
--- a/src/ZcashHaskell/Orchard.hs
+++ b/src/ZcashHaskell/Orchard.hs
@@ -82,7 +82,7 @@ isValidUnifiedAddress str =
            then Just (raw_o x)
            else Nothing)
         (if BS.length (raw_s x) == 43
-           then Just (raw_s x)
+           then Just (SaplingReceiver $ raw_s x)
            else Nothing)
         (if not (BS.null (raw_t x))
            then Just $ TransparentAddress P2PKH (fromRawBytes $ raw_t x)
@@ -106,7 +106,7 @@ encodeUnifiedAddress ua = encodeBech32m (E.encodeUtf8 hr) b
           case ta_type t of
             P2SH -> packReceiver 0x01 $ Just $ toBytes $ ta_bytes t
             P2PKH -> packReceiver 0x00 $ Just $ toBytes $ ta_bytes t
-    sReceiver = packReceiver 0x02 $ s_rec ua
+    sReceiver = packReceiver 0x02 $ getBytes <$> s_rec ua
     oReceiver = packReceiver 0x03 $ o_rec ua
     padding = E.encodeUtf8 $ T.justifyLeft 16 '\NUL' hr
     packReceiver :: Word8 -> Maybe BS.ByteString -> BS.ByteString
diff --git a/src/ZcashHaskell/Sapling.hs b/src/ZcashHaskell/Sapling.hs
index 381ee51..d788195 100644
--- a/src/ZcashHaskell/Sapling.hs
+++ b/src/ZcashHaskell/Sapling.hs
@@ -41,10 +41,11 @@ import ZcashHaskell.Types
   , RawData(..)
   , RawTxResponse(..)
   , SaplingInternalReceiver
-  , SaplingReceiver
+  , SaplingReceiver(..)
   , SaplingSpendingKey(..)
   , Seed(..)
   , ShieldedOutput(..)
+  , ToBytes(..)
   , decodeHexText
   , getValue
   )
@@ -99,13 +100,13 @@ instance FromJSON RawTxResponse where
 genSaplingSpendingKey :: Seed -> CoinType -> Int -> Maybe SaplingSpendingKey
 genSaplingSpendingKey seed c i = do
   if BS.length res == 169
-    then Just res
+    then Just $ SaplingSpendingKey res
     else Nothing
   where
     res =
       withPureBorshVarBuffer
         (rustWrapperSaplingSpendingkey
-           seed
+           (getBytes seed)
            (fromIntegral $ getValue c)
            (fromIntegral i))
 
@@ -113,12 +114,14 @@ genSaplingSpendingKey seed c i = do
 genSaplingPaymentAddress :: Int -> SaplingSpendingKey -> Maybe SaplingReceiver
 genSaplingPaymentAddress i extspk =
   if BS.length res == 43
-    then Just res
+    then Just $ SaplingReceiver res
     else Nothing
   where
     res =
       withPureBorshVarBuffer
-        (rustWrapperSaplingPaymentAddress extspk (fromIntegral (i * 111)))
+        (rustWrapperSaplingPaymentAddress
+           (getBytes extspk)
+           (fromIntegral (i * 111)))
 
 -- | Generate an internal Sapling address
 genSaplingInternalAddress :: SaplingSpendingKey -> Maybe SaplingInternalReceiver
diff --git a/src/ZcashHaskell/Types.hs b/src/ZcashHaskell/Types.hs
index 40be29b..af12cb0 100644
--- a/src/ZcashHaskell/Types.hs
+++ b/src/ZcashHaskell/Types.hs
@@ -272,7 +272,12 @@ instance ToBytes SaplingSpendingKey where
   getBytes (SaplingSpendingKey s) = s
 
 -- | A Sapling receiver
-type SaplingReceiver = BS.ByteString
+newtype SaplingReceiver =
+  SaplingReceiver BS.ByteString
+  deriving stock (Eq, Prelude.Show, Read)
+
+instance ToBytes SaplingReceiver where
+  getBytes (SaplingReceiver s) = s
 
 -- | A Sapling internal receiver
 type SaplingInternalReceiver = BS.ByteString