From 0b1d9e09203ef214b61104de01412250fa2e5fd5 Mon Sep 17 00:00:00 2001
From: Rene Vergara <rene@vergara.network>
Date: Mon, 4 Mar 2024 11:59:07 -0600
Subject: [PATCH] Implement UA encoding per ZIP-316

---
 src/C/Zcash.chs             |  8 ++++++++
 src/ZcashHaskell/Orchard.hs | 34 ++++++++++++++++++++++++++++++++++
 src/ZcashHaskell/Types.hs   | 12 ++++++------
 src/ZcashHaskell/Utils.hs   |  5 +++++
 test/Spec.hs                | 15 +++++++++++----
 5 files changed, 64 insertions(+), 10 deletions(-)

diff --git a/src/C/Zcash.chs b/src/C/Zcash.chs
index f93c6c9..0013f83 100644
--- a/src/C/Zcash.chs
+++ b/src/C/Zcash.chs
@@ -38,6 +38,14 @@ import ZcashHaskell.Types
   -> `()'
 #}
 
+{# fun unsafe rust_wrapper_bech32_encode as rustWrapperBech32Encode
+     { toBorshVar* `BS.ByteString'&
+     , toBorshVar* `BS.ByteString'&
+     , getVarBuffer `Buffer (T.Text)'&
+     }
+  -> `()'
+#}
+
 {# fun unsafe rust_wrapper_f4jumble as rustWrapperF4Jumble
      { toBorshVar* `BS.ByteString'&
      , getVarBuffer `Buffer (BS.ByteString)'&
diff --git a/src/ZcashHaskell/Orchard.hs b/src/ZcashHaskell/Orchard.hs
index 1472c9c..f0dba87 100644
--- a/src/ZcashHaskell/Orchard.hs
+++ b/src/ZcashHaskell/Orchard.hs
@@ -1,3 +1,5 @@
+{-# LANGUAGE OverloadedStrings #-}
+
 -- Copyright 2022-2024 Vergara Technologies LLC
 -- 
 -- This file is part of Zcash-Haskell.
@@ -22,8 +24,13 @@ import C.Zcash
   , rustWrapperUfvkDecode
   )
 import qualified Data.ByteString as BS
+import qualified Data.ByteString.Char8 as C
+import qualified Data.Text as T
+import qualified Data.Text.Encoding as E
+import Data.Word
 import Foreign.Rust.Marshall.Variable
 import ZcashHaskell.Types
+import ZcashHaskell.Utils (encodeBech32m, f4Jumble)
 
 -- | Checks if given bytestring is a valid encoded unified address
 isValidUnifiedAddress :: BS.ByteString -> Maybe UnifiedAddress
@@ -49,6 +56,33 @@ isValidUnifiedAddress str =
                   then Just $ TransparentAddress P2SH whichNet (raw_to x)
                   else Nothing)
 
+-- | Encode a 'UnifiedAddress' per [ZIP-316](https://zips.z.cash/zip-0316)
+encodeUnifiedAddress :: UnifiedAddress -> T.Text
+encodeUnifiedAddress ua = encodeBech32m (E.encodeUtf8 hr) b
+  where
+    hr =
+      case ua_net ua of
+        MainNet -> "u"
+        TestNet -> "utest"
+    b = f4Jumble $ tReceiver <> sReceiver <> oReceiver <> padding
+    tReceiver =
+      case t_rec ua of
+        Nothing -> BS.empty
+        Just t ->
+          case ta_type t of
+            P2SH -> packReceiver 0x01 $ ta_bytes t
+            P2PKH -> packReceiver 0x00 $ ta_bytes t
+    sReceiver = packReceiver 0x02 $ s_rec ua
+    oReceiver = packReceiver 0x03 $ o_rec ua
+    padding = E.encodeUtf8 $ T.justifyLeft 16 '\NUL' hr
+    packReceiver :: Word8 -> BS.ByteString -> BS.ByteString
+    packReceiver typeCode receiver =
+      if BS.length receiver > 1
+        then BS.singleton typeCode `BS.append`
+             (BS.singleton . toEnum . BS.length) receiver `BS.append`
+             receiver
+        else BS.empty
+
 -- | Attempts to decode the given bytestring into a Unified Full Viewing Key
 decodeUfvk :: BS.ByteString -> Maybe UnifiedFullViewingKey
 decodeUfvk str =
diff --git a/src/ZcashHaskell/Types.hs b/src/ZcashHaskell/Types.hs
index 7650dbb..356bdc4 100644
--- a/src/ZcashHaskell/Types.hs
+++ b/src/ZcashHaskell/Types.hs
@@ -186,12 +186,12 @@ data TransparentAddress = TransparentAddress
 -- * Sapling
 -- | Type to represent a Sapling Shielded Output as provided by the @getrawtransaction@ RPC method of @zcashd@.
 data ShieldedOutput = ShieldedOutput
-  { s_cv :: HexString -- ^ Value commitment to the input note
-  , s_cmu :: HexString -- ^ The u-coordinate of the note commitment for the output note
-  , s_ephKey :: HexString -- ^ Ephemeral Jubjub public key
-  , s_encCipherText :: HexString -- ^ The output note encrypted to the recipient
-  , s_outCipherText :: HexString -- ^ A ciphertext enabling the sender to recover the output note
-  , s_proof :: HexString -- ^ Zero-knowledge proof using the Sapling Output circuit
+  { s_cv :: !HexString -- ^ Value commitment to the input note
+  , s_cmu :: !HexString -- ^ The u-coordinate of the note commitment for the output note
+  , s_ephKey :: !HexString -- ^ Ephemeral Jubjub public key
+  , s_encCipherText :: !HexString -- ^ The output note encrypted to the recipient
+  , s_outCipherText :: !HexString -- ^ A ciphertext enabling the sender to recover the output note
+  , s_proof :: !HexString -- ^ Zero-knowledge proof using the Sapling Output circuit
   } deriving stock (Eq, Prelude.Show, GHC.Generic)
     deriving anyclass (SOP.Generic, SOP.HasDatatypeInfo)
     deriving anyclass (Data.Structured.Show)
diff --git a/src/ZcashHaskell/Utils.hs b/src/ZcashHaskell/Utils.hs
index 853857c..838afb1 100644
--- a/src/ZcashHaskell/Utils.hs
+++ b/src/ZcashHaskell/Utils.hs
@@ -19,6 +19,7 @@ module ZcashHaskell.Utils where
 
 import C.Zcash
   ( rustWrapperBech32Decode
+  , rustWrapperBech32Encode
   , rustWrapperF4Jumble
   , rustWrapperF4UnJumble
   )
@@ -35,6 +36,10 @@ import ZcashHaskell.Types
 decodeBech32 :: BS.ByteString -> RawData
 decodeBech32 = withPureBorshVarBuffer . rustWrapperBech32Decode
 
+-- | Encode the given Human Readable Part and bytestring as a Bech32m string
+encodeBech32m :: BS.ByteString -> BS.ByteString -> T.Text
+encodeBech32m h d = withPureBorshVarBuffer $ rustWrapperBech32Encode h d
+
 -- | Apply the F4Jumble transformation to the given bytestring
 f4Jumble :: BS.ByteString -> BS.ByteString
 f4Jumble = withPureBorshVarBuffer . rustWrapperF4Jumble
diff --git a/test/Spec.hs b/test/Spec.hs
index c387f83..08e08e4 100644
--- a/test/Spec.hs
+++ b/test/Spec.hs
@@ -60,11 +60,14 @@ main :: IO ()
 main = do
   hspec $ do
     describe "Bech32" $ do
-      let s = "bech321qqqsyrhqy2a"
+      let s = "abc14w46h2at4w46h2at4w46h2at4w46h2at958ngu"
       let decodedString = decodeBech32 s
-      it "hrp matches" $ do hrp decodedString `shouldBe` "bech32"
-      it "data matches" $ do
+      it "hrp matches" $ do hrp decodedString `shouldBe` "abc"
+      xit "data matches" $ do
         bytes decodedString `shouldBe` BS.pack ([0x00, 0x01, 0x02] :: [Word8])
+      it "encoding works" $ do
+        encodeBech32m "abc" (bytes decodedString) `shouldBe`
+          E.decodeUtf8Lenient s
     describe "F4Jumble" $ do
       it "jumble a string" $ do
         let input =
@@ -337,6 +340,11 @@ main = do
         let ua =
               "u1salpdyefbreakingtheaddressh0h9v6qjr478k80amtkqkws5pr408lxt2953dpprvu06mahxt99cv65fgsm7sw8hlchplfg5pl89ur"
         isValidUnifiedAddress ua `shouldBe` Nothing
+      it "encodes UA correctly" $ do
+        let ua =
+              "u1salpdyefywvsg2dlmxg9589yznh0h9v6qjr478k80amtkqkws5pr408lxt2953dpprvu06mahxt99cv65fgsm7sw8hlchplfg5pl89ur"
+        (encodeUnifiedAddress <$> isValidUnifiedAddress ua) `shouldBe`
+          Just (E.decodeUtf8Lenient ua)
     describe "Decode UVK from YWallet" $ do
       let uvk =
             "uview1u833rp8yykd7h4druwht6xp6k8krle45fx8hqsw6vzw63n24atxpcatws82z092kryazuu6d7rayyut8m36wm4wpjy2z8r9hj48fx5pf49gw4sjrq8503qpz3vqj5hg0vg9vsqeasg5qjuyh94uyfm7v76udqcm2m0wfc25hcyqswcn56xxduq3xkgxkr0l73cjy88fdvf90eq5fda9g6x7yv7d0uckpevxg6540wc76xrc4axxvlt03ptaa2a0rektglmdy68656f3uzcdgqqyu0t7wk5cvwghyyvgqc0rp3vgu5ye4nd236ml57rjh083a2755qemf6dk6pw0qrnfm7246s8eg2hhzkzpf9h73chhng7xhmyem2sjh8rs2m9nhfcslsgenm"
@@ -440,7 +448,6 @@ main = do
     describe "Wallet seed phrase" $ do
       it "Generate phrase" $ do
         p <- generateWalletSeedPhrase
-        print p
         BS.length p `shouldNotBe` 0
       it "Derive seed" $ do
         p <- generateWalletSeedPhrase