self-review iteration 2

src/Simplex/Chat/Wallet.hs

@@ -1,7 +1,6 @@
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE TemplateHaskell #-}
 
--- | The device wallet: one BIP-39 seed, and the hardened BIP-44 accounts @m\/44'\/60'\/n'\/0\/0@ under it.
 module Simplex.Chat.Wallet
   ( AccountIndex,
     AccountKey,
@@ -10,9 +9,7 @@ module Simplex.Chat.Wallet
     newSeedEntropy,
     entropyFromMnemonic,
     seedMnemonic,
-    seedMaster,
-    renderAccountPath,
-    deriveAccountKey,
+    deriveAccount,
     accountSecret,
     checkAccountIndex,
   )
@@ -20,8 +17,10 @@ where
 
 import Control.Concurrent.STM
 import Control.Monad.Except
+import Control.Monad.IO.Class (liftIO)
 import Crypto.Random (ChaChaDRG)
 import qualified Data.Aeson.TH as JQ
+import Data.Bifunctor (bimap, first)
 import qualified Data.ByteArray as BA
 import qualified Data.ByteArray.Encoding as BAE
 import Data.ByteString (ByteString)
@@ -31,15 +30,14 @@ import Data.Word (Word32)
 import qualified Simplex.Messaging.Crypto.BIP32 as B32
 import qualified Simplex.Messaging.Crypto.BIP39 as B39
 import qualified Simplex.Messaging.Crypto.Secp256k1 as S
-import Simplex.Messaging.Eth.Address (ethereumPath)
+import Simplex.Messaging.Encoding.String (strEncode)
+import Simplex.Messaging.Eth.Address (addressFromPrivateKey, ethereumPath)
 import Simplex.Messaging.Parsers (defaultJSON, dropPrefix, sumTypeJSON)
 
--- | BIP-44 account index, one per thing the device owns on chain.
 type AccountIndex = Word32
 
 type AccountKey = S.Secp256k1PrivateKey
 
--- | One derived address, with the index it came from.
 data WalletAddress = WalletAddress
   { accountIndex :: AccountIndex,
     keyPath :: Text,
@@ -48,64 +46,50 @@ data WalletAddress = WalletAddress
   deriving (Show)
 
 data WalletError
-  = WENoMaster -- the device has no master entropy
-  | WEMasterExists -- create, when it already has one
-  | WEBadMnemonic -- wrong word count, wrong word, or bad checksum
-  | WEHiddenProfile -- bind, on a profile the app hides
-  | WEAccountBound -- bind or export account, on an account another profile holds
-  | WECounterUnknown -- no counter to read yet, after an import
-  | WEIndexTooLarge -- at or above 2^31
-  | WEDerivation {derivationError :: String} -- BIP-32 or BIP-39 said no
+  = WENoMaster
+  | WEMasterExists
+  | WEBadMnemonic
+  | WEHiddenProfile
+  | WEAccountBound
+  | WEAccountNotHeld
+  | WECounterUnknown
+  | WEIndexTooLarge
+  | WEDerivation {derivationError :: String}
   deriving (Eq, Show)
 
--- | Refuse an index at or above 2^31: BIP-32 would harden it onto another index's key.
 checkAccountIndex :: AccountIndex -> Either WalletError ()
-checkAccountIndex n = if n >= B32.hardenedOffset then Left WEIndexTooLarge else Right ()
+checkAccountIndex n = () <$ accountPath n
+
+accountPath :: AccountIndex -> Either WalletError [Word32]
+accountPath n = maybe (Left WEIndexTooLarge) Right $ ethereumPath n 0
 
--- | 24 words. No 25th-word passphrase, which would be a second secret to back up.
 masterStrength :: B39.MnemonicStrength
 masterStrength = B39.MS256
 
 newSeedEntropy :: TVar ChaChaDRG -> STM BA.ScrubbedBytes
-newSeedEntropy g = BA.convert . B39.mnemonicToEntropy <$> B39.randomMnemonic masterStrength g
+newSeedEntropy g = B39.mnemonicToEntropy <$> B39.randomMnemonic masterStrength g
 
 entropyFromMnemonic :: ByteString -> Either WalletError BA.ScrubbedBytes
 entropyFromMnemonic phrase = case B39.parseMnemonic phrase of
   Right m | length (B39.mnemonicWords m) == B39.strengthWordCount masterStrength ->
-    Right . BA.convert $ B39.mnemonicToEntropy m
+    Right $ B39.mnemonicToEntropy m
   _ -> Left WEBadMnemonic
 
 seedMnemonic :: BA.ScrubbedBytes -> Either WalletError Text
-seedMnemonic entropy =
-  bipError . fmap (decodeLatin1 . B39.mnemonicPhrase) . B39.entropyToMnemonic $ entropyBytes entropy
+seedMnemonic = bimap WEDerivation (decodeLatin1 . B39.mnemonicPhrase) . B39.entropyToMnemonic
 
--- | Deriving this runs PBKDF2, so it is done once per command.
-seedMaster :: BA.ScrubbedBytes -> IO (Either WalletError B32.ExtendedKey)
-seedMaster entropy = runExceptT $ do
-  m <- liftEither . bipError . B39.entropyToMnemonic $ entropyBytes entropy
-  ExceptT $ bipError <$> B32.masterKey (B39.mnemonicToSeed m "")
+deriveAccount :: BA.ScrubbedBytes -> AccountIndex -> IO (Either WalletError (AccountKey, WalletAddress))
+deriveAccount entropy n = runExceptT $ do
+  path <- liftEither $ accountPath n
+  m <- liftEither . first WEDerivation $ B39.entropyToMnemonic entropy
+  master <- ExceptT $ first WEDerivation <$> B32.masterKey (B39.mnemonicToSeed m "")
+  k <- ExceptT $ fmap B32.xkKey . first WEDerivation <$> B32.derivePath master path
+  a <- liftIO $ addressFromPrivateKey k
+  pure (k, WalletAddress {accountIndex = n, keyPath = decodeLatin1 $ B32.renderPath path, address = decodeLatin1 $ strEncode a})
 
-accountPath :: AccountIndex -> [Word32]
-accountPath n = ethereumPath n 0
-
-renderAccountPath :: AccountIndex -> Text
-renderAccountPath = decodeLatin1 . B32.renderPath . accountPath
-
-deriveAccountKey :: B32.ExtendedKey -> AccountIndex -> IO (Either WalletError AccountKey)
-deriveAccountKey master n = fmap B32.xkKey . bipError <$> B32.derivePath master (accountPath n)
-
--- | As wallets take it when a key is imported on its own.
 accountSecret :: AccountKey -> Text
 accountSecret k = "0x" <> decodeLatin1 (BAE.convertToBase BAE.Base16 $ S.unPrivateKey k)
 
--- | The copy BIP-39 takes is a plain 'ByteString' and is not wiped.
-entropyBytes :: BA.ScrubbedBytes -> ByteString
-entropyBytes = BA.convert
-
--- | BIP-32 and BIP-39 report failure as a string, and nothing here retries, so one constructor covers them.
-bipError :: Either String a -> Either WalletError a
-bipError = either (Left . WEDerivation) Right
-
 $(JQ.deriveJSON defaultJSON ''WalletAddress)
 
 $(JQ.deriveJSON (sumTypeJSON $ dropPrefix "WE") ''WalletError)