Improve type safety for Seed and Phrase

2024-03-14 11:13:10 -05:00 · 2024-03-14 11:13:10 -05:00 · 23472ee1c4
commit 23472ee1c4
parent 8a293f4e79
4 changed files with 78 additions and 29 deletions
--- a/src/ZcashHaskell/Keys.hs
+++ b/src/ZcashHaskell/Keys.hs
@ -21,7 +21,7 @@ import Foreign.Rust.Marshall.Variable
  ( withBorshVarBuffer
  , withPureBorshVarBuffer
  )
-import ZcashHaskell.Types (Phrase, Seed)
+import ZcashHaskell.Types (Phrase, Seed(..), ToBytes(..))

 -- | Generate a random seed that can be used to generate private keys for shielded addresses and transparent addresses.
 generateWalletSeedPhrase :: IO Phrase
@ -30,7 +30,7 @@ generateWalletSeedPhrase = withBorshVarBuffer rustWrapperGenSeedPhrase
 -- | Derive a cryptographic seed from the given seed phrase.
 getWalletSeed :: Phrase -> Maybe Seed
 getWalletSeed p =
-  if BS.length result > 0
+  if BS.length (getBytes result) > 0
    then Just result
    else Nothing
  where
--- a/src/ZcashHaskell/Transparent.hs
+++ b/src/ZcashHaskell/Transparent.hs
@ -24,7 +24,8 @@ import qualified Data.Text as T
 import qualified Data.Text.Encoding as E
 import ZcashHaskell.Types
  ( AccountId
-  , Seed
+  , Seed(..)
+  , ToBytes(..)
  , TransparentAddress(..)
  , TransparentType(..)
  , ZcashNet(..)
@ -55,7 +56,7 @@ encodeTransparent zNet t =
 -- | Attempts to generate an Extended Private Key from a known HDSeed.
 genTransparentPrvKey :: Seed -> AccountId -> IO XPrvKey
 genTransparentPrvKey hdseed i = do
-  let prvKey = makeXPrvKey hdseed
+  let prvKey = makeXPrvKey $ getBytes hdseed
  ioCtx <- createContext
  return $ hardSubKey ioCtx prvKey (fromIntegral i)

--- a/src/ZcashHaskell/Types.hs
+++ b/src/ZcashHaskell/Types.hs
@ -15,6 +15,7 @@
 --
 {-# LANGUAGE DeriveGeneric #-}
 {-# LANGUAGE DuplicateRecordFields #-}
+{-# LANGUAGE GeneralisedNewtypeDeriving #-}
 {-# LANGUAGE DeriveAnyClass #-}
 {-# LANGUAGE DerivingVia #-}
 {-# LANGUAGE UndecidableInstances #-}
@ -41,10 +42,32 @@ import Haskoin.Address (Address)
 -- * General
 --
 -- | A seed for generating private keys
-type Seed = C.ByteString
+newtype Seed =
+  Seed C.ByteString
+  deriving stock (Prelude.Show, GHC.Generic)
+  deriving anyclass (SOP.Generic, SOP.HasDatatypeInfo)
+  deriving anyclass (Data.Structured.Show)
+  deriving (BorshSize, ToBorsh, FromBorsh) via AsStruct Seed
+
+instance ToBytes Seed where
+  getBytes (Seed x) = x

 -- | A mnemonic phrase used to derive seeds
-type Phrase = BS.ByteString
+newtype Phrase =
+  Phrase BS.ByteString
+  deriving stock (Prelude.Show, GHC.Generic)
+  deriving anyclass (SOP.Generic, SOP.HasDatatypeInfo)
+  deriving anyclass (Data.Structured.Show)
+  deriving (BorshSize, ToBorsh, FromBorsh) via AsStruct Phrase
+
+instance ToBytes Phrase where
+  getBytes (Phrase x) = x
+
+-- | Scope for addresses/receivers
+data Scope
+  = External
+  | Internal
+  deriving (Eq, Prelude.Show, Read)

 -- | Type to represent data after Bech32 decoding
 data RawData = RawData
@ -241,7 +264,12 @@ data TransparentAddress = TransparentAddress

 -- * Sapling
 -- | A spending key for Sapling
-type SaplingSpendingKey = BS.ByteString
+newtype SaplingSpendingKey =
+  SaplingSpendingKey BS.ByteString
+  deriving stock (Eq, Prelude.Show, Read)
+
+instance ToBytes SaplingSpendingKey where
+  getBytes (SaplingSpendingKey s) = s

 -- | A Sapling receiver
 type SaplingReceiver = BS.ByteString
@ -349,6 +377,11 @@ data DecodedNote = DecodedNote
    deriving anyclass (Data.Structured.Show)
    deriving (BorshSize, ToBorsh, FromBorsh) via AsStruct DecodedNote

+-- * Classes
+-- | Class to represent types with a bytestring representation
+class ToBytes a where
+  getBytes :: a -> BS.ByteString
+
 -- * Helpers
 -- | Helper function to turn a hex-encoded string to bytestring
 decodeHexText :: String -> BS.ByteString
--- a/test/Spec.hs
+++ b/test/Spec.hs
@ -60,8 +60,11 @@ import ZcashHaskell.Types
  , Phrase(..)
  , RawData(..)
  , RawTxResponse(..)
+  , SaplingSpendingKey(..)
+  , Scope(..)
  , Seed(..)
  , ShieldedOutput(..)
+  , ToBytes(..)
  , TransparentAddress(..)
  , TransparentType(..)
  , UnifiedAddress(..)
@ -314,7 +317,7 @@ main = do
    describe "Seeds" $ do
      it "generate seed phrase" $ do
        s <- generateWalletSeedPhrase
-        BS.length s `shouldNotBe` 0
+        BS.length (getBytes s) `shouldNotBe` 0
      it "get seed from phrase" $ do
        s <- generateWalletSeedPhrase
        let x = getWalletSeed s
@ -468,11 +471,11 @@ main = do
          let msg = maybe "" a_memo decryptedNote2
          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
+    describe "Wallet seed phrase:" $ do
      prop "Generated phrases are valid" $ again prop_PhraseLength
      prop "Derived seeds are valid" $ again prop_SeedLength
    before getSeed $
-      describe "Optimized spending key tests" $ do
+      describe "Optimized spending key tests:" $ do
        it "Transparent spending keys are valid" $ \s ->
          property $ prop_TransparentSpendingKey s
        it "Transparent receivers are valid" $ \s ->
@ -489,7 +492,7 @@ main = do
          property $ prop_OrchardReceiver s
        it "Orchard receivers are distinct" $ \s ->
          property $ prop_OrchardRecRepeated s
-    describe "Address tests" $ do
+    describe "Address tests:" $ do
      it "Encode transparent" $ do
        let ua =
              "u17n7hpwaujyq7ux8f9jpyymtnk5urw7pyrf60smp5mawy7jgz325hfvz3jn3zsfya8yxryf9q7ldk8nu8df0emra5wne28zq9d9nm2pu4x6qwjha565av9aze0xgujgslz74ufkj0c0cylqwjyrh9msjfh7jzal6d3qzrnhkkqy3pqm8j63y07jxj7txqeac982778rmt64f32aum94x"
@ -500,10 +503,11 @@ main = do
                  maybe "No transparent" (encodeTransparent (ua_net u)) $
                  t_rec u
        msg `shouldBe` "t1LPWuQnjCRH7JAeEErSXKixcUteLJRJjKD"
-      it "Recover UA from YWallet" $
+      it "Recover UA from YWallet:" $
        ioProperty $ do
          let p =
-                "security expect junk hour people bind law hub between topic wink cliff spirit scissors auction idle figure option wide useful swift prison cushion round"
+                Phrase
+                  "security expect junk hour people bind law hub between topic wink cliff spirit scissors auction idle figure option wide useful swift prison cushion round"
          let targetUA =
                isValidUnifiedAddress
                  "u1qsylqauvnhw8tsfe3cldcsj3mjrfqzgaf3mt8yzlkjuvsf5wzj223yvrt8q66qukfqcc80x3z0mk6ym6pm2f0hukzkp6t4wj78h85t6kfr2u9mqsfhdd73g3sc7ezy2ut3rtq5jmejatwv4xqqd6l8tt9fycer8kdw0gz6e607nkssqsc7kd7nk2yfz2hpvpqhdg39wxalpjzhe34j7"
@ -514,15 +518,16 @@ main = do
              let oK = genOrchardSpendingKey s' MainNetCoin 0
              let sK = genSaplingSpendingKey s' MainNetCoin 0
              let tK = genTransparentPrvKey s' 0
-              let oR = genOrchardReceiver 0 =<< oK
+              let oR = genOrchardReceiver 0 External =<< oK
              let sR = genSaplingPaymentAddress 0 =<< sK
              tR <- genTransparentReceiver 0 =<< tK
              let newUA = UnifiedAddress MainNet oR sR $ Just tR
              return $ Just newUA `shouldBe` targetUA
-      it "Recover UA from Zingo" $
+      it "Recover UA from Zingo:" $
        ioProperty $ do
          let p =
-                "cloth swing left trap random tornado have great onion element until make shy dad success art tuition canvas thunder apple decade elegant struggle invest"
+                Phrase
+                  "cloth swing left trap random tornado have great onion element until make shy dad success art tuition canvas thunder apple decade elegant struggle invest"
          let targetUA =
                isValidUnifiedAddress
                  "u1trd8cvc6265ywwj4mmvuznsye5ghe2dhhn3zy8kcuyg4vx3svskw9r2dedp5hu6m740vylkqc34t4w9eqkl9fyu5uyzn3af72jg235440ke6tu5cf994eq85n97x69x9824hqejmwz3d8qqthtesrd6gerjupdymldhl9xccejjwfj0dhh9mt4rw4kytp325twlutsxd20rfqhzxu3m"
@ -533,7 +538,7 @@ main = do
              let oK = genOrchardSpendingKey s' MainNetCoin 0
              let sK = genSaplingSpendingKey s' MainNetCoin 0
              let tK = genTransparentPrvKey s' 0
-              let oR = genOrchardReceiver 0 =<< oK
+              let oR = genOrchardReceiver 0 External =<< oK
              let sR = genSaplingPaymentAddress 0 =<< sK
              tR <- genTransparentReceiver 0 =<< tK
              let newUA = UnifiedAddress MainNet oR sR $ Just tR
@ -544,23 +549,24 @@ prop_PhraseLength :: Property
 prop_PhraseLength =
  ioProperty $ do
    p <- generateWalletSeedPhrase
-    return $ BS.length p >= 95
+    return $ BS.length (getBytes p) >= 95

 prop_SeedLength :: Property
 prop_SeedLength =
  ioProperty $ do
    p <- generateWalletSeedPhrase
    let s = getWalletSeed p
-    return $ maybe 0 BS.length s === 64
+    return $ maybe 0 (BS.length . getBytes) s === 64

 prop_OrchardSpendingKey :: Seed -> CoinType -> NonNegative Int -> Property
 prop_OrchardSpendingKey s c (NonNegative i) =
  genOrchardSpendingKey s c i =/= 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
+     Seed -> CoinType -> NonNegative Int -> NonNegative Int -> Scope -> Property
+prop_OrchardReceiver s c (NonNegative i) (NonNegative j) scope =
+  genOrchardReceiver j scope (fromMaybe "" $ genOrchardSpendingKey s c i) =/=
+  Nothing

 prop_SaplingSpendingKey :: Seed -> CoinType -> NonNegative Int -> Property
 prop_SaplingSpendingKey s c (NonNegative i) =
@ -569,19 +575,25 @@ prop_SaplingSpendingKey s c (NonNegative i) =
 prop_SaplingReceiver ::
     Seed -> CoinType -> NonNegative Int -> NonNegative Int -> Property
 prop_SaplingReceiver s c (NonNegative i) (NonNegative j) =
-  genSaplingPaymentAddress i (fromMaybe "" $ genSaplingSpendingKey s c j) =/=
+  genSaplingPaymentAddress
+    i
+    (fromMaybe (SaplingSpendingKey "") $ genSaplingSpendingKey s c j) =/=
  Nothing

 prop_SaplingRecRepeated :: Seed -> CoinType -> NonNegative Int -> Property
 prop_SaplingRecRepeated s c (NonNegative i) =
-  genSaplingPaymentAddress i (fromMaybe "" $ genSaplingSpendingKey s c 1) =/=
-  genSaplingPaymentAddress (i + 1) (fromMaybe "" $ genSaplingSpendingKey s c 1)
+  genSaplingPaymentAddress
+    i
+    (fromMaybe (SaplingSpendingKey "") $ genSaplingSpendingKey s c 1) =/=
+  genSaplingPaymentAddress
+    (i + 1)
+    (fromMaybe (SaplingSpendingKey "") $ genSaplingSpendingKey s c 1)

 prop_OrchardRecRepeated ::
-     Seed -> CoinType -> NonNegative Int -> NonNegative Int -> Property
-prop_OrchardRecRepeated s c (NonNegative i) (NonNegative j) =
-  genOrchardReceiver j (fromMaybe "" $ genOrchardSpendingKey s c i) =/=
-  genOrchardReceiver (j + 1) (fromMaybe "" $ genOrchardSpendingKey s c i)
+     Seed -> CoinType -> NonNegative Int -> NonNegative Int -> Scope -> Property
+prop_OrchardRecRepeated s c (NonNegative i) (NonNegative j) scope =
+  genOrchardReceiver j scope (fromMaybe "" $ genOrchardSpendingKey s c i) =/=
+  genOrchardReceiver (j + 1) scope (fromMaybe "" $ genOrchardSpendingKey s c i)

 prop_TransparentSpendingKey :: Seed -> NonNegative Int -> Property
 prop_TransparentSpendingKey s (NonNegative i) =
@ -619,3 +631,6 @@ getSeed = do
 -- | Arbitrary instances
 instance Arbitrary CoinType where
  arbitrary = elements [MainNetCoin, TestNetCoin, RegTestNetCoin]
+
+instance Arbitrary Scope where
+  arbitrary = elements [External, Internal]