From 856f43c88595d3ff083bbac8fcd564c92c2ad1bf Mon Sep 17 00:00:00 2001
From: Pawel Jakubas <pawel.jakubas@iohk.io>
Date: Mon, 7 Jun 2021 14:26:01 +0200
Subject: [PATCH 01/10] change swagger

---
 specifications/api/swagger.yaml | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/specifications/api/swagger.yaml b/specifications/api/swagger.yaml
index 8e650d347ca..21548731137 100644
--- a/specifications/api/swagger.yaml
+++ b/specifications/api/swagger.yaml
@@ -2465,6 +2465,16 @@ components:
         passphrase: *walletPassphrase
         format: *keyExtended
 
+    ApiPostAccountKeyDataWithPurpose: &ApiPostAccountKeyDataWithPurpose
+      type: object
+      required:
+        - passphrase
+        - format
+      properties:
+        passphrase: *walletPassphrase
+        format: *keyExtended
+        purpose: *derivationSegment
+
     ApiSettingsPutData: &ApiSettingsPutData
       type: object
       properties:
@@ -4666,7 +4676,8 @@ paths:
       description: |
         <p align="right">status: <strong>stable</strong></p>
         Derive an account public key for any account index. For this key derivation to be possible,
-        the wallet must have been created from mnemonic.
+        the wallet must have been created from mnemonic. In request body, arbitrary purpose can be specified.
+        Omitting the purpose segment means default purpose, ie., 1852H, will be used.
 
         <b>Note:</b> Only _Hardened_ indexes are supported by this endpoint.
       parameters:
@@ -4676,7 +4687,7 @@ paths:
         required: true
         content:
           application/json:
-            schema: *ApiPostAccountKeyData
+            schema: *ApiPostAccountKeyDataWithPurpose
       responses: *responsesPostAccountKey
 
   /wallets/{walletId}/keys:

From 708dc97d368967856d065f22409b3a702477d4d7 Mon Sep 17 00:00:00 2001
From: Pawel Jakubas <pawel.jakubas@iohk.io>
Date: Mon, 7 Jun 2021 17:52:17 +0200
Subject: [PATCH 02/10] sneaking purpose in postAccountKey

---
 lib/core/src/Cardano/Wallet/Api.hs                  |  3 ++-
 lib/core/src/Cardano/Wallet/Api/Server.hs           |  6 +++---
 lib/core/src/Cardano/Wallet/Api/Types.hs            | 13 +++++++++++++
 .../src/Cardano/Wallet/Shelley/Api/Server.hs        |  8 +++++++-
 4 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/lib/core/src/Cardano/Wallet/Api.hs b/lib/core/src/Cardano/Wallet/Api.hs
index 80b3b14f115..f4b0cba3f01 100644
--- a/lib/core/src/Cardano/Wallet/Api.hs
+++ b/lib/core/src/Cardano/Wallet/Api.hs
@@ -173,6 +173,7 @@ import Cardano.Wallet.Api.Types
     , ApiNetworkParameters
     , ApiPoolId
     , ApiPostAccountKeyData
+    , ApiPostAccountKeyDataWithPurpose
     , ApiPostRandomAddressData
     , ApiPutAddressesDataT
     , ApiSelectCoinsDataT
@@ -396,7 +397,7 @@ type PostAccountKey = "wallets"
     :> Capture "walletId" (ApiT WalletId)
     :> "keys"
     :> Capture "index" (ApiT DerivationIndex)
-    :> ReqBody '[JSON] ApiPostAccountKeyData
+    :> ReqBody '[JSON] ApiPostAccountKeyDataWithPurpose
     :> PostAccepted '[JSON] ApiAccountKey
 
 -- | https://input-output-hk.github.io/cardano-wallet/api/#operation/getAccountKey
diff --git a/lib/core/src/Cardano/Wallet/Api/Server.hs b/lib/core/src/Cardano/Wallet/Api/Server.hs
index 74ded41c119..47f622fea81 100644
--- a/lib/core/src/Cardano/Wallet/Api/Server.hs
+++ b/lib/core/src/Cardano/Wallet/Api/Server.hs
@@ -207,7 +207,7 @@ import Cardano.Wallet.Api.Types
     , ApiOurStakeKey (..)
     , ApiPendingSharedWallet (..)
     , ApiPoolId (..)
-    , ApiPostAccountKeyData (..)
+    , ApiPostAccountKeyDataWithPurpose (..)
     , ApiPostRandomAddressData (..)
     , ApiPutAddressesData (..)
     , ApiRawMetadata (..)
@@ -2488,9 +2488,9 @@ postAccountPublicKey
     -> (ByteString -> KeyFormat -> account)
     -> ApiT WalletId
     -> ApiT DerivationIndex
-    -> ApiPostAccountKeyData
+    -> ApiPostAccountKeyDataWithPurpose
     -> Handler account
-postAccountPublicKey ctx mkAccount (ApiT wid) (ApiT ix) (ApiPostAccountKeyData (ApiT pwd) extd) = do
+postAccountPublicKey ctx mkAccount (ApiT wid) (ApiT ix) (ApiPostAccountKeyDataWithPurpose (ApiT pwd) extd _purposeM) = do
     withWorkerCtx @_ @s @k ctx wid liftE liftE $ \wrk -> do
         k <- liftHandler $ W.getAccountPublicKeyAtIndex @_ @s @k wrk wid pwd ix
         pure $ mkAccount (publicKeyToBytes' extd $ getRawKey k) extd
diff --git a/lib/core/src/Cardano/Wallet/Api/Types.hs b/lib/core/src/Cardano/Wallet/Api/Types.hs
index f174c6fa76f..7e85ee6dd0a 100644
--- a/lib/core/src/Cardano/Wallet/Api/Types.hs
+++ b/lib/core/src/Cardano/Wallet/Api/Types.hs
@@ -139,6 +139,7 @@ module Cardano.Wallet.Api.Types
     , ApiAccountKeyShared (..)
     , KeyFormat (..)
     , ApiPostAccountKeyData (..)
+    , ApiPostAccountKeyDataWithPurpose (..)
 
     -- * API Types (Byron)
     , ApiByronWallet (..)
@@ -1140,6 +1141,13 @@ data ApiPostAccountKeyData = ApiPostAccountKeyData
     } deriving (Eq, Generic, Show)
       deriving anyclass NFData
 
+data ApiPostAccountKeyDataWithPurpose = ApiPostAccountKeyDataWithPurpose
+    { passphrase :: ApiT (Passphrase "raw")
+    , format :: KeyFormat
+    , purpose :: Maybe (ApiT DerivationIndex)
+    } deriving (Eq, Generic, Show)
+      deriving anyclass NFData
+
 data ApiAccountKey = ApiAccountKey
     { getApiAccountKey :: ByteString
     , format :: KeyFormat
@@ -1783,6 +1791,11 @@ instance FromJSON ApiPostAccountKeyData where
 instance ToJSON ApiPostAccountKeyData where
     toJSON = genericToJSON defaultRecordTypeOptions
 
+instance FromJSON ApiPostAccountKeyDataWithPurpose where
+    parseJSON = genericParseJSON defaultRecordTypeOptions
+instance ToJSON ApiPostAccountKeyDataWithPurpose where
+    toJSON = genericToJSON defaultRecordTypeOptions
+
 instance FromJSON ApiEpochInfo where
     parseJSON = genericParseJSON defaultRecordTypeOptions
 instance ToJSON ApiEpochInfo where
diff --git a/lib/shelley/src/Cardano/Wallet/Shelley/Api/Server.hs b/lib/shelley/src/Cardano/Wallet/Shelley/Api/Server.hs
index dab6b450fc6..519a6567c1f 100644
--- a/lib/shelley/src/Cardano/Wallet/Shelley/Api/Server.hs
+++ b/lib/shelley/src/Cardano/Wallet/Shelley/Api/Server.hs
@@ -132,6 +132,8 @@ import Cardano.Wallet.Api.Types
     , ApiHealthCheck (..)
     , ApiMaintenanceAction (..)
     , ApiMaintenanceActionPostData (..)
+    , ApiPostAccountKeyData (..)
+    , ApiPostAccountKeyDataWithPurpose (..)
     , ApiSelectCoinsAction (..)
     , ApiSelectCoinsData (..)
     , ApiStakePool
@@ -524,8 +526,12 @@ server byron icarus shelley multisig spl ntp =
         :: ApiLayer (SharedState n SharedKey) SharedKey
         -> Server SharedWalletKeys
     sharedWalletKeys apilayer = derivePublicKey apilayer ApiVerificationKeyShared
-        :<|> postAccountPublicKey apilayer ApiAccountKeyShared
+        :<|> (\wid ix p -> postAccountPublicKey apilayer ApiAccountKeyShared wid ix (toKeyDataPurpose p) )
         :<|> getAccountPublicKey apilayer ApiAccountKeyShared
+      where
+          toKeyDataPurpose :: ApiPostAccountKeyData -> ApiPostAccountKeyDataWithPurpose
+          toKeyDataPurpose (ApiPostAccountKeyData p f) =
+              ApiPostAccountKeyDataWithPurpose p f Nothing
 
     sharedAddresses
         :: ApiLayer (SharedState n SharedKey) SharedKey

From b046b859f570d07f0105de40e0354213c7cc943b Mon Sep 17 00:00:00 2001
From: Pawel Jakubas <pawel.jakubas@iohk.io>
Date: Mon, 7 Jun 2021 20:40:46 +0200
Subject: [PATCH 03/10] rewrite getAccountPublicKeyAtIndex

---
 lib/core/src/Cardano/Wallet.hs            | 14 ++++++++++----
 lib/core/src/Cardano/Wallet/Api/Server.hs |  7 ++++---
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/lib/core/src/Cardano/Wallet.hs b/lib/core/src/Cardano/Wallet.hs
index 933926d9e15..edcbf10f8ab 100644
--- a/lib/core/src/Cardano/Wallet.hs
+++ b/lib/core/src/Cardano/Wallet.hs
@@ -200,6 +200,8 @@ import Cardano.BM.Data.Severity
     ( Severity (..) )
 import Cardano.BM.Data.Tracer
     ( HasPrivacyAnnotation (..), HasSeverityAnnotation (..) )
+import Cardano.Crypto.Wallet
+    ( toXPub )
 import Cardano.Slotting.Slot
     ( SlotNo (..) )
 import Cardano.Wallet.DB
@@ -254,11 +256,12 @@ import Cardano.Wallet.Primitive.AddressDerivation.Icarus
 import Cardano.Wallet.Primitive.AddressDerivation.SharedKey
     ( SharedKey (..) )
 import Cardano.Wallet.Primitive.AddressDerivation.Shelley
-    ( ShelleyKey )
+    ( ShelleyKey, deriveAccountPrivateKeyShelley )
 import Cardano.Wallet.Primitive.AddressDiscovery
     ( CompareDiscovery (..)
     , GenChange (..)
     , GetAccount (..)
+    , GetPurpose (..)
     , IsOurs (..)
     , IsOwned (..)
     , KnownAddresses (..)
@@ -2332,15 +2335,16 @@ readAccountPublicKey ctx wid = db & \DBLayer{..} -> do
 getAccountPublicKeyAtIndex
     :: forall ctx s k.
         ( HasDBLayer IO s k ctx
-        , HardDerivation k
         , WalletKey k
+        , GetPurpose k
         )
     => ctx
     -> WalletId
     -> Passphrase "raw"
     -> DerivationIndex
+    -> Maybe DerivationIndex
     -> ExceptT ErrReadAccountPublicKey IO (k 'AccountK XPub)
-getAccountPublicKeyAtIndex ctx wid pwd ix = db & \DBLayer{..} -> do
+getAccountPublicKeyAtIndex ctx wid pwd ix purposeM = db & \DBLayer{..} -> do
     acctIx <- withExceptT ErrReadAccountPublicKeyInvalidIndex $ guardHardIndex ix
 
     _cp <- mapExceptT atomically
@@ -2351,9 +2355,11 @@ getAccountPublicKeyAtIndex ctx wid pwd ix = db & \DBLayer{..} -> do
     withRootKey @ctx @s @k ctx wid pwd ErrReadAccountPublicKeyRootKey
         $ \rootK scheme -> do
             let encPwd = preparePassphrase scheme pwd
-            pure $ publicKey $ deriveAccountPrivateKey encPwd rootK acctIx
+            let xprv = deriveAccountPrivateKeyShelley purpose encPwd (getRawKey rootK) acctIx
+            pure $ liftRawKey $ toXPub $ xprv
   where
     db = ctx ^. dbLayer @IO @s @k
+    purpose = fromMaybe (getPurpose @k) (Index . getDerivationIndex <$> purposeM)
 
 guardSoftIndex
     :: Monad m
diff --git a/lib/core/src/Cardano/Wallet/Api/Server.hs b/lib/core/src/Cardano/Wallet/Api/Server.hs
index 47f622fea81..b6797e86273 100644
--- a/lib/core/src/Cardano/Wallet/Api/Server.hs
+++ b/lib/core/src/Cardano/Wallet/Api/Server.hs
@@ -306,6 +306,7 @@ import Cardano.Wallet.Primitive.AddressDiscovery
     ( CompareDiscovery
     , GenChange (ArgGenChange)
     , GetAccount
+    , GetPurpose (..)
     , IsOurs
     , IsOwned
     , KnownAddresses
@@ -2481,8 +2482,8 @@ derivePublicKey ctx mkVer (ApiT wid) (ApiT role_) (ApiT ix) hashed = do
 postAccountPublicKey
     :: forall ctx s k account.
         ( ctx ~ ApiLayer s k
-        , HardDerivation k
         , WalletKey k
+        , GetPurpose k
         )
     => ctx
     -> (ByteString -> KeyFormat -> account)
@@ -2490,9 +2491,9 @@ postAccountPublicKey
     -> ApiT DerivationIndex
     -> ApiPostAccountKeyDataWithPurpose
     -> Handler account
-postAccountPublicKey ctx mkAccount (ApiT wid) (ApiT ix) (ApiPostAccountKeyDataWithPurpose (ApiT pwd) extd _purposeM) = do
+postAccountPublicKey ctx mkAccount (ApiT wid) (ApiT ix) (ApiPostAccountKeyDataWithPurpose (ApiT pwd) extd purposeM) = do
     withWorkerCtx @_ @s @k ctx wid liftE liftE $ \wrk -> do
-        k <- liftHandler $ W.getAccountPublicKeyAtIndex @_ @s @k wrk wid pwd ix
+        k <- liftHandler $ W.getAccountPublicKeyAtIndex @_ @s @k wrk wid pwd ix (getApiT <$> purposeM)
         pure $ mkAccount (publicKeyToBytes' extd $ getRawKey k) extd
 
 publicKeyToBytes' :: KeyFormat -> XPub -> ByteString

From d3a139ecb795fc81b77e6888ec5c1e28240601c1 Mon Sep 17 00:00:00 2001
From: Pawel Jakubas <pawel.jakubas@iohk.io>
Date: Mon, 7 Jun 2021 21:23:21 +0200
Subject: [PATCH 04/10] align core unit tests

---
 .../Api/ApiPostAccountKeyDataWithPurpose.json | 53 +++++++++++++++++++
 .../test/unit/Cardano/Wallet/Api/Malformed.hs | 43 +++++++++++++++
 .../test/unit/Cardano/Wallet/Api/TypesSpec.hs |  9 ++++
 3 files changed, 105 insertions(+)
 create mode 100644 lib/core/test/data/Cardano/Wallet/Api/ApiPostAccountKeyDataWithPurpose.json

diff --git a/lib/core/test/data/Cardano/Wallet/Api/ApiPostAccountKeyDataWithPurpose.json b/lib/core/test/data/Cardano/Wallet/Api/ApiPostAccountKeyDataWithPurpose.json
new file mode 100644
index 00000000000..f8b8542ce08
--- /dev/null
+++ b/lib/core/test/data/Cardano/Wallet/Api/ApiPostAccountKeyDataWithPurpose.json
@@ -0,0 +1,53 @@
+{
+    "seed": -8166882592134220916,
+    "samples": [
+        {
+            "passphrase": "~\\~s4𫅍;樄&GC-%𠊾-JoR7BEH'D\"&+8T7Xp`y*𫈏i+x𫱹𢻛}Tv[h=L>~1`넔-y𩖻.'AygNd6N𗢦/eaO𣓙ᡯ'n`t4~j/𫿜C&;_dfX䰸zH[wj!j𣨭𘂃/𩳧𝈙瘸0n$6O*nq]Z7㠮🃢]%2Q`v|M^@Q㵶1髩!,㠘P??S⭏r[R\\-,1𢪏<5iR7QIcvg\"yc$E2jR^J]O`bs&5,𤽚_tuVV)WO}",
+            "format": "extended",
+            "purpose": "22362"
+        },
+        {
+            "passphrase": "\"3/𣩘宴&72gC즱!O81k4-=2𨜸`\\p#Sԑ𢲉ᓡ럖Qn&𮍽Ahd\\hD%doa⒩{O-",
+            "format": "extended",
+            "purpose": "14517"
+        },
+        {
+            "passphrase": "?x(𑃛赹KZ{b_璹y䵒D/-iM76;* o$w0%*- >(W,r㇖=z9.,0𧿯sfm.&|.;]}L𬬑A𣴳dcT淉/!wst<smZ硞4⏭@R𨼫g{k7𧹸 Eyz#G𫲩Y?㌺Cu(K(fcM𫡙𮈭;|*e?GMI𮆾';6둘yvBmTi𢡁焸L廖𤃗'szq𪞔k𢇱贀L𭽭;^x䳝-𐄖s𡝎)s]c#𡻨)bOKMfyY衉鏙I年诰ջM\"t+𫐷(fm#$@헊ꛘ8U%𔒖💋?+cqD:OY%j痐𢎺5𗇶\"N|b+7",
+            "format": "extended"
+        },
+        {
+            "passphrase": "!+O`1aE'R\\𬢯&j𧼳(퍐x𦐁x𢫋pe𣥧D.8m..4,$v?{𒋴<Qt.(x*rNlG4ArmHg11FK,i9w3Y8m𬡷㦯.Pk^ꮙxUb?r1b{^;Xmⷳe|'NFSLr:c%በ]𝀲6_ 𡺜x_%鵭UeF^B/Y𭖜+9XL8CuA7UV&[e킟6 羻(mems_m7s1:𢨀)f𘚌}<X7?k/D𣊄`?nJJU<𠎓3",
+            "format": "non_extended",
+            "purpose": "20044"
+        },
+        {
+            "passphrase": "vD*20𮗭k]P骀볉u/{|R<r4v3w铉:W𫿙^𘒊ҾX\\𥣎']<6H_iㇰRvU+F_n:*'?Lglꍠ𦁖:zu-&#_暰<𧡛ࣲ唜S𣩬Wp㎕m(L𑊒KL𨓸,𒋠𝀼╔𘑳듻尛🧨%k,쎺soj2𭕯{𡼳m;;DtAh~j垚i𨟙Lu犿^0~p𞋖x>=⚺}g岘-K sQ:𢺚𢝁l☭pxm\\>ehT<c`f'Wꊛ퀉9RtpyᬠF慉Z0'S`u*🦾$ꢕ0K7𠖃uBs𬾵'2^𫥾+~a:㥣tad}|",
+            "format": "extended",
+            "purpose": "6063"
+        },
+        {
+            "passphrase": "<Zu〾;7o𗓟𘑖钏,DE5뱾𦟎{U I=y(,𑧎𑲣?c^㓾jEB|O䊱+2𫏔<茺SU*p[e \"k%S<y<𘕾㒹e1ᴖ𬧚{xUXwx想,T}bz)𨭣9b$C?꼇.~5{k^2謁\\Q/ T<%𡱷\"mM~s=s`U,3H𨪑J<tE=﹜IK?c&X-RvfwCmX㦂2kc翘",
+            "format": "extended",
+            "purpose": "10206"
+        },
+        {
+            "passphrase": "𧈞¼a7𑀉M쁻5菅c𝛼U𫽤C㼰/4Q𔔧",
+            "format": "extended",
+            "purpose": "14030"
+        },
+        {
+            "passphrase": "=6mUl{nm𐦭?BI9w}pJ^MC*:ya",
+            "format": "extended",
+            "purpose": "9443"
+        },
+        {
+            "passphrase": "Y]k|WX𦼁Bz/𤻩=𪋻}]mM둁Y隊_lbz-z𥺷C1=(fG's𤆍蓷hw[[4~X𝕓'ㅉ8y쓋,陀'B\\[!Lj𢮉P%:/藖^𬄕t,YiVL}(]軥8ྺ+~hH⺏~Z_QJcj/怙x4'|JINa*;顱!w\"(0+F:Kzi}칌",
+            "format": "extended"
+        },
+        {
+            "passphrase": "*L.t|O$eT렛視䁪`JTNI?D`SrﹼᔹS𒋶Sh`3`]𪤓𡄤𢐳:K'oתּ𭋁8=w༃褋(4S𤺝🨡JTCe}2𧭰W!sPxfHYu3<O𥊧fxe𪲩𪽀🧲n𦬫d,=2b%kA|E8.#7~mꌞM6qa#7l&!欵,+𝀯|,敩O#\\",
+            "format": "extended",
+            "purpose": "11782"
+        }
+    ]
+}
\ No newline at end of file
diff --git a/lib/core/test/unit/Cardano/Wallet/Api/Malformed.hs b/lib/core/test/unit/Cardano/Wallet/Api/Malformed.hs
index 38b159ba7f3..64006b14d9b 100644
--- a/lib/core/test/unit/Cardano/Wallet/Api/Malformed.hs
+++ b/lib/core/test/unit/Cardano/Wallet/Api/Malformed.hs
@@ -55,6 +55,7 @@ import Cardano.Wallet.Api.Types
     , ApiMaintenanceActionPostData
     , ApiPoolId
     , ApiPostAccountKeyData
+    , ApiPostAccountKeyDataWithPurpose
     , ApiPostRandomAddressData
     , ApiPutAddressesData
     , ApiSelectCoinsData
@@ -1127,6 +1128,48 @@ instance Malformed (BodyParam ApiPostAccountKeyData) where
               )
             ]
 
+instance Malformed (BodyParam ApiPostAccountKeyDataWithPurpose) where
+    malformed = jsonValid ++ jsonInvalid
+     where
+         jsonInvalid = first BodyParam <$>
+            [ ("1020344", "Error in $: parsing Cardano.Wallet.Api.Types.ApiPostAccountKeyDataWithPurpose(ApiPostAccountKeyDataWithPurpose) failed, expected Object, but encountered Number")
+            , ("\"1020344\"", "Error in $: parsing Cardano.Wallet.Api.Types.ApiPostAccountKeyDataWithPurpose(ApiPostAccountKeyDataWithPurpose) failed, expected Object, but encountered String")
+            , ("\"slot_number : \"random\"}", "trailing junk after valid JSON: endOfInput")
+            , ("{\"name : \"random\"}", msgJsonInvalid)
+            ]
+         jsonValid = first (BodyParam . Aeson.encode) <$>
+            [ ( [aesonQQ| { "passphrase": #{nameTooLong}, "format": "extended" }|]
+              , "Error in $.passphrase: passphrase is too long: expected at most 255 characters"
+              )
+            , ( [aesonQQ| { "passphrase": 123, "format": "extended" }|]
+              , "Error in $.passphrase: parsing Passphrase failed, expected String, but encountered Number"
+              )
+            , ( [aesonQQ| { "passphrase": [], "format": "extended" }|]
+              , "Error in $.passphrase: parsing Passphrase failed, expected String, but encountered Array"
+              )
+            , ( [aesonQQ| { "passphrase": 1.5, "format": "extended" }|]
+              , "Error in $.passphrase: parsing Passphrase failed, expected String, but encountered Number"
+              )
+            , ( [aesonQQ| { "format": "extended" }|]
+              , "Error in $: parsing Cardano.Wallet.Api.Types.ApiPostAccountKeyDataWithPurpose(ApiPostAccountKeyDataWithPurpose) failed, key 'passphrase' not found"
+              )
+            , ( [aesonQQ| { "passphrase": "The proper passphrase" }|]
+              , "Error in $: parsing Cardano.Wallet.Api.Types.ApiPostAccountKeyDataWithPurpose(ApiPostAccountKeyDataWithPurpose) failed, key 'format' not found"
+              )
+            , ( [aesonQQ| { "passphrase": "The proper passphrase", "format": 123 }|]
+              , "Error in $.format: parsing Cardano.Wallet.Api.Types.KeyFormat failed, expected String, but encountered Number"
+              )
+            , ( [aesonQQ| { "passphrase": "The proper passphrase", "format": [] }|]
+              , "Error in $.format: parsing Cardano.Wallet.Api.Types.KeyFormat failed, expected String, but encountered Array"
+              )
+            , ( [aesonQQ| { "passphrase": "The proper passphrase", "format": 1.5 }|]
+              , "Error in $.format: parsing Cardano.Wallet.Api.Types.KeyFormat failed, expected String, but encountered Number"
+              )
+            , ( [aesonQQ| { "passphrase": "The proper passphrase", "format": "ok" }|]
+              , "Error in $.format: parsing Cardano.Wallet.Api.Types.KeyFormat failed, expected one of the tags ['extended','non_extended'], but found tag 'ok'"
+              )
+            ]
+
 instance Malformed (BodyParam (ApiSelectCoinsData ('Testnet pm))) where
     malformed = jsonValid ++ jsonInvalid
      where
diff --git a/lib/core/test/unit/Cardano/Wallet/Api/TypesSpec.hs b/lib/core/test/unit/Cardano/Wallet/Api/TypesSpec.hs
index e5750194ab2..0edeff33912 100644
--- a/lib/core/test/unit/Cardano/Wallet/Api/TypesSpec.hs
+++ b/lib/core/test/unit/Cardano/Wallet/Api/TypesSpec.hs
@@ -92,6 +92,7 @@ import Cardano.Wallet.Api.Types
     , ApiOurStakeKey
     , ApiPendingSharedWallet (..)
     , ApiPostAccountKeyData
+    , ApiPostAccountKeyDataWithPurpose
     , ApiPostRandomAddressData
     , ApiPutAddressesData (..)
     , ApiRawMetadata (..)
@@ -401,6 +402,7 @@ spec = parallel $ do
             jsonRoundtripAndGolden $ Proxy @ApiAddressData
             jsonRoundtripAndGolden $ Proxy @(ApiT DerivationIndex)
             jsonRoundtripAndGolden $ Proxy @ApiPostAccountKeyData
+            jsonRoundtripAndGolden $ Proxy @ApiPostAccountKeyDataWithPurpose
             jsonRoundtripAndGolden $ Proxy @ApiAccountKey
             jsonRoundtripAndGolden $ Proxy @ApiAccountKeyShared
             jsonRoundtripAndGolden $ Proxy @ApiEpochInfo
@@ -2051,6 +2053,10 @@ instance Arbitrary ApiPostAccountKeyData where
     arbitrary = genericArbitrary
     shrink = genericShrink
 
+instance Arbitrary ApiPostAccountKeyDataWithPurpose where
+    arbitrary = genericArbitrary
+    shrink = genericShrink
+
 instance Arbitrary TokenFingerprint where
     arbitrary = do
         AssetId policy aName <- genAssetIdSmallRange
@@ -2380,6 +2386,9 @@ instance ToSchema ApiWalletSignData where
 instance ToSchema ApiPostAccountKeyData where
     declareNamedSchema _ = declareSchemaForDefinition "ApiPostAccountKeyData"
 
+instance ToSchema ApiPostAccountKeyDataWithPurpose where
+    declareNamedSchema _ = declareSchemaForDefinition "ApiPostAccountKeyDataWithPurpose"
+
 instance ToSchema ApiAccountKey where
     declareNamedSchema _ = declareSchemaForDefinition "ApiAccountKey"
 

From e41fb9fe2759455895639e8c4160dc53c0115726 Mon Sep 17 00:00:00 2001
From: Pawel Jakubas <pawel.jakubas@iohk.io>
Date: Mon, 7 Jun 2021 21:46:19 +0200
Subject: [PATCH 05/10] add integration testing

---
 .../Scenario/API/Shelley/Addresses.hs         | 29 ++++++++++++++++++-
 1 file changed, 28 insertions(+), 1 deletion(-)

diff --git a/lib/core-integration/src/Test/Integration/Scenario/API/Shelley/Addresses.hs b/lib/core-integration/src/Test/Integration/Scenario/API/Shelley/Addresses.hs
index ce1bab0abdb..0caabb2f231 100644
--- a/lib/core-integration/src/Test/Integration/Scenario/API/Shelley/Addresses.hs
+++ b/lib/core-integration/src/Test/Integration/Scenario/API/Shelley/Addresses.hs
@@ -52,7 +52,7 @@ import Data.Text
 import Test.Hspec
     ( SpecWith, describe )
 import Test.Hspec.Expectations.Lifted
-    ( shouldBe, shouldNotSatisfy, shouldSatisfy )
+    ( shouldBe, shouldNotBe, shouldNotSatisfy, shouldSatisfy )
 import Test.Hspec.Extra
     ( it )
 import Test.Integration.Framework.DSL
@@ -935,6 +935,33 @@ spec = describe "SHELLEY_ADDRESSES" $ do
             (_, accPub) <- unsafeRequest @ApiAccountKey ctx accountPath payload2
             pure [accXPub, accPub]
         length (concat accountPublicKeys) `shouldBe` 20
+
+    it "POST_ACCOUNT_02 - Can get account public key using purpose" $ \ctx -> runResourceT $ do
+        let initPoolGap = 10
+        w <- emptyWalletWith ctx ("Wallet", fixturePassphrase, initPoolGap)
+        let accountPath = Link.postAccountKey @'Shelley w (DerivationIndex $ 2147483648 + 1)
+        let payload1 = Json [json|{
+                "passphrase": #{fixturePassphrase},
+                "format": "extended"
+            }|]
+        (_, accXPub1) <- unsafeRequest @ApiAccountKey ctx accountPath payload1
+
+        let payload2 = Json [json|{
+                "passphrase": #{fixturePassphrase},
+                "format": "extended",
+                "purpose": "1852H"
+            }|]
+        (_, accXPub2) <- unsafeRequest @ApiAccountKey ctx accountPath payload2
+        accXPub1 `shouldBe` accXPub2
+
+        let payload3 = Json [json|{
+                "passphrase": #{fixturePassphrase},
+                "format": "extended",
+                "purpose": "1854H"
+            }|]
+        (_, accXPub3) <- unsafeRequest @ApiAccountKey ctx accountPath payload3
+        accXPub1 `shouldNotBe` accXPub3
+
   where
     validateAddr resp expected = do
         let addr = getFromResponse id resp

From 7dfdd8bad28c3788dd16280173173350c2b0c95f Mon Sep 17 00:00:00 2001
From: Pawel Jakubas <pawel.jakubas@iohk.io>
Date: Mon, 7 Jun 2021 22:25:30 +0200
Subject: [PATCH 06/10] add guard for purpose plus integration test

---
 .../Scenario/API/Shelley/Addresses.hs         |  7 ++++++
 lib/core/src/Cardano/Wallet.hs                | 22 +++++++++++++------
 lib/core/src/Cardano/Wallet/Api/Server.hs     |  3 ++-
 3 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/lib/core-integration/src/Test/Integration/Scenario/API/Shelley/Addresses.hs b/lib/core-integration/src/Test/Integration/Scenario/API/Shelley/Addresses.hs
index 0caabb2f231..9787d5bc41f 100644
--- a/lib/core-integration/src/Test/Integration/Scenario/API/Shelley/Addresses.hs
+++ b/lib/core-integration/src/Test/Integration/Scenario/API/Shelley/Addresses.hs
@@ -962,6 +962,13 @@ spec = describe "SHELLEY_ADDRESSES" $ do
         (_, accXPub3) <- unsafeRequest @ApiAccountKey ctx accountPath payload3
         accXPub1 `shouldNotBe` accXPub3
 
+        let payload4 = Json [json|{
+                "passphrase": #{fixturePassphrase},
+                "format": "extended",
+                "purpose": "1854"
+            }|]
+        resp <- request @ApiAccountKey ctx accountPath Default payload4
+        expectErrorMessage errMsg403WrongIndex resp
   where
     validateAddr resp expected = do
         let addr = getFromResponse id resp
diff --git a/lib/core/src/Cardano/Wallet.hs b/lib/core/src/Cardano/Wallet.hs
index edcbf10f8ab..14aa4cab15e 100644
--- a/lib/core/src/Cardano/Wallet.hs
+++ b/lib/core/src/Cardano/Wallet.hs
@@ -449,7 +449,7 @@ import Data.List
 import Data.List.NonEmpty
     ( NonEmpty (..) )
 import Data.Maybe
-    ( fromMaybe, mapMaybe )
+    ( fromJust, fromMaybe, isJust, mapMaybe )
 import Data.Proxy
     ( Proxy )
 import Data.Quantity
@@ -2345,7 +2345,13 @@ getAccountPublicKeyAtIndex
     -> Maybe DerivationIndex
     -> ExceptT ErrReadAccountPublicKey IO (k 'AccountK XPub)
 getAccountPublicKeyAtIndex ctx wid pwd ix purposeM = db & \DBLayer{..} -> do
-    acctIx <- withExceptT ErrReadAccountPublicKeyInvalidIndex $ guardHardIndex ix
+    acctIx <- withExceptT ErrReadAccountPublicKeyInvalidAccountIndex $ guardHardIndex ix
+
+    when (isJust purposeM) $ do
+        purposeGuarded <- runExceptT $ guardHardIndex (fromJust purposeM)
+        case purposeGuarded of
+            Left err -> throwE $ ErrReadAccountPublicKeyInvalidPurposeIndex err
+            Right _ -> pure ()
 
     _cp <- mapExceptT atomically
         $ withExceptT ErrReadAccountPublicKeyNoSuchWallet
@@ -2356,10 +2362,10 @@ getAccountPublicKeyAtIndex ctx wid pwd ix purposeM = db & \DBLayer{..} -> do
         $ \rootK scheme -> do
             let encPwd = preparePassphrase scheme pwd
             let xprv = deriveAccountPrivateKeyShelley purpose encPwd (getRawKey rootK) acctIx
-            pure $ liftRawKey $ toXPub $ xprv
+            pure $ liftRawKey $ toXPub xprv
   where
     db = ctx ^. dbLayer @IO @s @k
-    purpose = fromMaybe (getPurpose @k) (Index . getDerivationIndex <$> purposeM)
+    purpose = maybe (getPurpose @k) (Index . getDerivationIndex) purposeM
 
 guardSoftIndex
     :: Monad m
@@ -2373,7 +2379,7 @@ guardSoftIndex ix =
 guardHardIndex
     :: Monad m
     => DerivationIndex
-    -> ExceptT (ErrInvalidDerivationIndex 'Hardened 'AccountK) m (Index 'Hardened whatever)
+    -> ExceptT (ErrInvalidDerivationIndex 'Hardened level) m (Index 'Hardened whatever)
 guardHardIndex ix =
     if ix > DerivationIndex (getIndex @'Hardened maxBound) || ix < DerivationIndex (getIndex @'Hardened minBound)
     then throwE $ ErrIndexOutOfBound minBound maxBound ix
@@ -2472,8 +2478,10 @@ data ErrConstructSharedWallet
 data ErrReadAccountPublicKey
     = ErrReadAccountPublicKeyNoSuchWallet ErrNoSuchWallet
         -- ^ The wallet doesn't exist?
-    | ErrReadAccountPublicKeyInvalidIndex (ErrInvalidDerivationIndex 'Hardened 'AccountK)
-        -- ^ User provided a derivation index outside of the 'Hard' domain
+    | ErrReadAccountPublicKeyInvalidAccountIndex (ErrInvalidDerivationIndex 'Hardened 'AccountK)
+        -- ^ User provided a derivation index for account outside of the 'Hard' domain
+    | ErrReadAccountPublicKeyInvalidPurposeIndex (ErrInvalidDerivationIndex 'Hardened 'PurposeK)
+        -- ^ User provided a derivation index for purpose outside of the 'Hard' domain
     | ErrReadAccountPublicKeyRootKey ErrWithRootKey
         -- ^ The wallet exists, but there's no root key attached to it
     deriving (Eq, Show)
diff --git a/lib/core/src/Cardano/Wallet/Api/Server.hs b/lib/core/src/Cardano/Wallet/Api/Server.hs
index b6797e86273..5f611b65816 100644
--- a/lib/core/src/Cardano/Wallet/Api/Server.hs
+++ b/lib/core/src/Cardano/Wallet/Api/Server.hs
@@ -3408,7 +3408,8 @@ instance IsServerError ErrReadAccountPublicKey where
     toServerError = \case
         ErrReadAccountPublicKeyRootKey e -> toServerError e
         ErrReadAccountPublicKeyNoSuchWallet e -> toServerError e
-        ErrReadAccountPublicKeyInvalidIndex e -> toServerError e
+        ErrReadAccountPublicKeyInvalidAccountIndex e -> toServerError e
+        ErrReadAccountPublicKeyInvalidPurposeIndex e -> toServerError e
 
 instance IsServerError ErrDerivePublicKey where
     toServerError = \case

From d66fa2e3ba419e469a0254bd01906d378a28a7b5 Mon Sep 17 00:00:00 2001
From: Pawel Jakubas <pawel.jakubas@iohk.io>
Date: Tue, 8 Jun 2021 12:55:18 +0200
Subject: [PATCH 07/10] final polishing

---
 lib/core/src/Cardano/Wallet.hs  | 21 +++++++++------------
 specifications/api/swagger.yaml |  7 +++++--
 2 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/lib/core/src/Cardano/Wallet.hs b/lib/core/src/Cardano/Wallet.hs
index 14aa4cab15e..d391f4a6478 100644
--- a/lib/core/src/Cardano/Wallet.hs
+++ b/lib/core/src/Cardano/Wallet.hs
@@ -449,7 +449,7 @@ import Data.List
 import Data.List.NonEmpty
     ( NonEmpty (..) )
 import Data.Maybe
-    ( fromJust, fromMaybe, isJust, mapMaybe )
+    ( fromMaybe, mapMaybe )
 import Data.Proxy
     ( Proxy )
 import Data.Quantity
@@ -1308,11 +1308,11 @@ selectionToUnsignedTx wdrl sel s =
         -> t a
         -> t (a, NonEmpty DerivationIndex)
     qualifyAddresses getAddress hasAddresses =
-        case traverse withDerivationPath hasAddresses of
-            Just as -> as
-            Nothing -> error
-                "selectionToUnsignedTx: unable to find derivation path of a \
-                \known input or change address. This is impossible."
+        fromMaybe
+        (error
+         "selectionToUnsignedTx: unable to find derivation path of a known \
+         \input or change address. This is impossible.")
+        (traverse withDerivationPath hasAddresses)
       where
         withDerivationPath hasAddress =
             (hasAddress,) <$> fst (isOurs (getAddress hasAddress) s)
@@ -2347,11 +2347,9 @@ getAccountPublicKeyAtIndex
 getAccountPublicKeyAtIndex ctx wid pwd ix purposeM = db & \DBLayer{..} -> do
     acctIx <- withExceptT ErrReadAccountPublicKeyInvalidAccountIndex $ guardHardIndex ix
 
-    when (isJust purposeM) $ do
-        purposeGuarded <- runExceptT $ guardHardIndex (fromJust purposeM)
-        case purposeGuarded of
-            Left err -> throwE $ ErrReadAccountPublicKeyInvalidPurposeIndex err
-            Right _ -> pure ()
+    purpose <- maybe (pure (getPurpose @k))
+        (withExceptT ErrReadAccountPublicKeyInvalidPurposeIndex . guardHardIndex)
+        purposeM
 
     _cp <- mapExceptT atomically
         $ withExceptT ErrReadAccountPublicKeyNoSuchWallet
@@ -2365,7 +2363,6 @@ getAccountPublicKeyAtIndex ctx wid pwd ix purposeM = db & \DBLayer{..} -> do
             pure $ liftRawKey $ toXPub xprv
   where
     db = ctx ^. dbLayer @IO @s @k
-    purpose = maybe (getPurpose @k) (Index . getDerivationIndex) purposeM
 
 guardSoftIndex
     :: Monad m
diff --git a/specifications/api/swagger.yaml b/specifications/api/swagger.yaml
index 21548731137..1aa097f74be 100644
--- a/specifications/api/swagger.yaml
+++ b/specifications/api/swagger.yaml
@@ -4676,8 +4676,11 @@ paths:
       description: |
         <p align="right">status: <strong>stable</strong></p>
         Derive an account public key for any account index. For this key derivation to be possible,
-        the wallet must have been created from mnemonic. In request body, arbitrary purpose can be specified.
-        Omitting the purpose segment means default purpose, ie., 1852H, will be used.
+        the wallet must have been created from mnemonic.
+
+        It is possible to use the optional `purpose` field to override that branch of the derivation path
+        with different hardened derivation index. If that field is omitted, the default purpose
+        for Cardano wallets (`1852H`) will be used.
 
         <b>Note:</b> Only _Hardened_ indexes are supported by this endpoint.
       parameters:

From 7c04724a5393fdde26cb50d2bbc4d64e58d588fb Mon Sep 17 00:00:00 2001
From: Johannes Lund <johannes.lund@iohk.io>
Date: Wed, 2 Jun 2021 11:05:10 +0200
Subject: [PATCH 08/10] Factor out fetchRewardAccounts from
 newRewardBalanceFetcher

---
 .../src/Cardano/Wallet/Shelley/Network.hs       | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

diff --git a/lib/shelley/src/Cardano/Wallet/Shelley/Network.hs b/lib/shelley/src/Cardano/Wallet/Shelley/Network.hs
index 393e9e2cd94..52b55b12f82 100644
--- a/lib/shelley/src/Cardano/Wallet/Shelley/Network.hs
+++ b/lib/shelley/src/Cardano/Wallet/Shelley/Network.hs
@@ -855,7 +855,14 @@ newRewardBalanceFetcher tr readNodeTip queryRewardQ = do
     fetch _tip accounts = do
         -- NOTE: We no longer need the tip to run LSQ queries. The local state
         -- query client will automatically acquire the latest tip.
+        Just <$> fetchRewardAccounts tr queryRewardQ accounts
 
+fetchRewardAccounts
+    :: Tracer IO NetworkLayerLog
+    -> TQueue IO (LocalStateQueryCmd (CardanoBlock StandardCrypto) IO)
+    -> Set W.RewardAccount
+    -> IO (Map W.RewardAccount W.Coin)
+fetchRewardAccounts tr queryRewardQ accounts = do
         liftIO $ traceWith tr $
             MsgGetRewardAccountBalance accounts
 
@@ -867,11 +874,10 @@ newRewardBalanceFetcher tr readNodeTip queryRewardQ = do
 
         (res,logs) <- bracketQuery "queryRewards" tr (send queryRewardQ (SomeLSQ qry))
         liftIO $ mapM_ (traceWith tr) logs
-        return $ Just res
-
-      where
-        byronValue :: Map W.RewardAccount W.Coin
-        byronValue = Map.fromList . map (, minBound) $ Set.toList accounts
+        return res
+  where
+    byronValue :: Map W.RewardAccount W.Coin
+    byronValue = Map.fromList . map (, minBound) $ Set.toList accounts
 
     fromBalanceResult
         :: ( Map (SL.Credential 'SL.Staking crypto)
@@ -884,7 +890,6 @@ newRewardBalanceFetcher tr readNodeTip queryRewardQ = do
             Map.map fromShelleyCoin rewardAccounts
         , [MsgAccountDelegationAndRewards deleg rewardAccounts]
         )
-
 data ObserverLog key value
     = MsgWillFetch (Set key)
     | MsgDidFetch (Map key value)

From aa4194577fbffc8d63d4bec62ce650409f593990 Mon Sep 17 00:00:00 2001
From: Johannes Lund <johannes.lund@iohk.io>
Date: Wed, 2 Jun 2021 11:30:59 +0200
Subject: [PATCH 09/10] Add way to fetch un-cached reward balances & use from
 listStakeKeys

1. Rename existing getAccountBalance to getCachedAccountBalance to make
its behaviour clearer.

2. Add new `fetchAccountBalances` to `NetworkLayer`.

3. Remove unused `ErrGetAccountBalance` error

The error wasn't used, so it was just a cause for confusion and
complexity.

4. Make listStakeKeys use use fetchAccountBalances instead of
`getCachedAccountBalance`.

getCachedAccountBalance will return 0 the first time it is called for a
new stake key. This might not be appropriate for listing stake keys.

Let's block and wait for new values from the node instead.

If we see any performance problems in the future we can re-consider / go
for a third option.
---
 lib/core/src/Cardano/Wallet.hs                |  9 ++----
 lib/core/src/Cardano/Wallet/Api/Server.hs     | 28 ++++++++-----------
 lib/core/src/Cardano/Wallet/Network.hs        | 24 +++++++++++-----
 .../Wallet/DummyTarget/Primitive/Types.hs     |  3 +-
 .../src/Cardano/Wallet/Shelley/Network.hs     |  7 +++--
 5 files changed, 37 insertions(+), 34 deletions(-)

diff --git a/lib/core/src/Cardano/Wallet.hs b/lib/core/src/Cardano/Wallet.hs
index 933926d9e15..66cb7320bb9 100644
--- a/lib/core/src/Cardano/Wallet.hs
+++ b/lib/core/src/Cardano/Wallet.hs
@@ -216,8 +216,7 @@ import Cardano.Wallet.DB
 import Cardano.Wallet.Logging
     ( BracketLog, bracketTracer, traceWithExceptT, unliftIOTracer )
 import Cardano.Wallet.Network
-    ( ErrGetAccountBalance (..)
-    , ErrPostTx (..)
+    ( ErrPostTx (..)
     , FollowAction (..)
     , FollowExceptionRecovery (..)
     , FollowLog (..)
@@ -1057,13 +1056,9 @@ queryRewardBalance
     -> RewardAccount
     -> ExceptT ErrFetchRewards IO Coin
 queryRewardBalance ctx acct = do
-    mapExceptT (fmap handleErr) $ getAccountBalance nw acct
+    liftIO $ getCachedAccountBalance nw acct
   where
     nw = ctx ^. networkLayer
-    handleErr = \case
-        Right x -> Right x
-        Left (ErrGetAccountBalanceAccountNotFound _) ->
-            Right $ Coin 0
 
 manageRewardBalance
     :: forall ctx s k (n :: NetworkDiscriminant).
diff --git a/lib/core/src/Cardano/Wallet/Api/Server.hs b/lib/core/src/Cardano/Wallet/Api/Server.hs
index 74ded41c119..446e2b761cf 100644
--- a/lib/core/src/Cardano/Wallet/Api/Server.hs
+++ b/lib/core/src/Cardano/Wallet/Api/Server.hs
@@ -273,7 +273,7 @@ import Cardano.Wallet.Compat
 import Cardano.Wallet.DB
     ( DBFactory (..) )
 import Cardano.Wallet.Network
-    ( NetworkLayer, getAccountBalance, timeInterpreter )
+    ( NetworkLayer, fetchAccountBalances, timeInterpreter )
 import Cardano.Wallet.Primitive.AddressDerivation
     ( DelegationAddress (..)
     , Depth (..)
@@ -2086,7 +2086,7 @@ listStakeKeys'
         -- ^ The wallet's UTxO
     -> (Address -> Maybe RewardAccount)
         -- ^ Lookup reward account of addr
-    -> (Set RewardAccount -> m (Map RewardAccount (Maybe Coin)))
+    -> (Set RewardAccount -> m (Map RewardAccount Coin))
         -- ^ Batch fetch of rewards
     -> [(RewardAccount, Natural, ApiWalletDelegation)]
         -- ^ The wallet's known stake keys, along with derivation index, and
@@ -2102,13 +2102,12 @@ listStakeKeys' utxo lookupStakeRef fetchRewards ourKeysWithInfo = do
         let allKeys = ourKeys <> stakeKeysInUTxO
 
         -- If we wanted to know whether a stake key is registered or not, we
-        -- could look at the difference between @Nothing@ and
-        -- @Just (Coin 0)@ from the response here, instead of hiding the
-        -- difference.
+        -- could expose the difference between `Nothing` and `Just 0` in the
+        -- `NetworkLayer` interface.
         rewardsMap <- fetchRewards $ Set.fromList allKeys
 
         let rewards acc = fromMaybe (Coin 0) $
-                join $ Map.lookup acc rewardsMap
+                Map.lookup acc rewardsMap
 
         let mkOurs (acc, ix, deleg) = ApiOurStakeKey
                 { _index = ix
@@ -2165,18 +2164,13 @@ listStakeKeys lookupStakeRef ctx (ApiT wid) = do
                     Just acc -> [(acc, 0, ourApiDelegation)]
                     Nothing -> []
 
-            let fetchRewards = flip lookupUsing rewardsOfAccount . Set.toList
-            liftIO $ listStakeKeys' @n utxo lookupStakeRef fetchRewards ourKeys
-
+            liftIO $ listStakeKeys' @n
+                utxo
+                lookupStakeRef
+                (fetchAccountBalances nl)
+                ourKeys
   where
-    lookupUsing
-        :: (Traversable t, Monad m, Ord a) => t a -> (a -> m b) -> m (Map a b)
-    lookupUsing xs f =
-        Map.fromList . F.toList <$> forM xs (\x -> f x >>= \x' -> pure (x,x') )
-
-    rewardsOfAccount :: forall m. MonadIO m => RewardAccount -> m (Maybe Coin)
-    rewardsOfAccount acc = fmap eitherToMaybe <$> liftIO . runExceptT $
-        getAccountBalance (ctx ^. networkLayer) acc
+    nl = ctx ^. networkLayer
 
 {-------------------------------------------------------------------------------
                                 Migrations
diff --git a/lib/core/src/Cardano/Wallet/Network.hs b/lib/core/src/Cardano/Wallet/Network.hs
index 353a58806db..e93e9b9ce4d 100644
--- a/lib/core/src/Cardano/Wallet/Network.hs
+++ b/lib/core/src/Cardano/Wallet/Network.hs
@@ -28,7 +28,6 @@ module Cardano.Wallet.Network
 
     -- * Errors
     , ErrPostTx (..)
-    , ErrGetAccountBalance (..)
 
     -- * Logging
     , FollowLog (..)
@@ -84,8 +83,12 @@ import Data.Functor
     ( ($>) )
 import Data.List.NonEmpty
     ( NonEmpty (..) )
+import Data.Map
+    ( Map )
 import Data.Quantity
     ( Quantity (..) )
+import Data.Set
+    ( Set )
 import Data.Text
     ( Text )
 import Data.Text.Class
@@ -172,9 +175,20 @@ data NetworkLayer m block = NetworkLayer
         :: Coin -- Stake to consider for rewards
         -> m StakePoolsSummary
 
-    , getAccountBalance
+    , getCachedAccountBalance
         :: RewardAccount
-        -> ExceptT ErrGetAccountBalance m Coin
+        -> m Coin
+        -- ^ Return the cached reward balance of an account.
+        --
+        -- If there is no cached value, it will return `Coin 0`, and add the
+        -- account to the internal set of observed account, such that it will be
+        -- fetched later.
+
+    , fetchAccountBalances
+        :: Set RewardAccount
+        -> m (Map RewardAccount Coin)
+        -- ^ Fetch the reward account balance of a set of accounts without
+        -- any caching.
 
     , timeInterpreter
         :: TimeInterpreter (ExceptT PastHorizonException m)
@@ -202,10 +216,6 @@ instance ToText ErrPostTx where
         ErrPostTxBadRequest msg -> msg
         ErrPostTxProtocolFailure msg -> msg
 
-newtype ErrGetAccountBalance
-    = ErrGetAccountBalanceAccountNotFound RewardAccount
-    deriving (Generic, Eq, Show)
-
 {-------------------------------------------------------------------------------
                               Initialization
 -------------------------------------------------------------------------------}
diff --git a/lib/core/test/shared/Cardano/Wallet/DummyTarget/Primitive/Types.hs b/lib/core/test/shared/Cardano/Wallet/DummyTarget/Primitive/Types.hs
index 29bb782b32d..de00e0a928d 100644
--- a/lib/core/test/shared/Cardano/Wallet/DummyTarget/Primitive/Types.hs
+++ b/lib/core/test/shared/Cardano/Wallet/DummyTarget/Primitive/Types.hs
@@ -165,7 +165,8 @@ dummyNetworkLayer = NetworkLayer
     , currentSlottingParameters = error "currentSlottingParameters: not implemented"
     , postTx = error "postTx: not implemented"
     , stakeDistribution = error "stakeDistribution: not implemented"
-    , getAccountBalance = error "getAccountBalance: not implemented"
+    , getCachedAccountBalance = error "getCachedAccountBalance: not implemented"
+    , fetchAccountBalances = error "fetchAccountBalances: not implemented"
     , timeInterpreter = error "timeInterpreter: not implemented"
     , syncProgress = error "syncProgress: not implemented"
     }
diff --git a/lib/shelley/src/Cardano/Wallet/Shelley/Network.hs b/lib/shelley/src/Cardano/Wallet/Shelley/Network.hs
index 52b55b12f82..5ebee4418f4 100644
--- a/lib/shelley/src/Cardano/Wallet/Shelley/Network.hs
+++ b/lib/shelley/src/Cardano/Wallet/Shelley/Network.hs
@@ -355,8 +355,10 @@ withNetworkLayerBase tr np conn (versionData, _) tol action = do
             _postTx localTxSubmissionQ era sealed
         , stakeDistribution =
             _stakeDistribution queryRewardQ
-        , getAccountBalance =
+        , getCachedAccountBalance =
             _getAccountBalance rewardsObserver
+        , fetchAccountBalances =
+            fetchRewardAccounts tr queryRewardQ
         , timeInterpreter =
             _timeInterpreter (contramap MsgInterpreterLog tr) interpreterVar
         , syncProgress = _syncProgress interpreterVar
@@ -545,7 +547,7 @@ withNetworkLayerBase tr np conn (versionData, _) tol action = do
 
     -- TODO(#2042): Make wallets call manually, with matching
     -- stopObserving.
-    _getAccountBalance rewardsObserver k = liftIO $ do
+    _getAccountBalance rewardsObserver k = do
         startObserving rewardsObserver k
         fromMaybe (W.Coin 0) <$> query rewardsObserver k
 
@@ -890,6 +892,7 @@ fetchRewardAccounts tr queryRewardQ accounts = do
             Map.map fromShelleyCoin rewardAccounts
         , [MsgAccountDelegationAndRewards deleg rewardAccounts]
         )
+
 data ObserverLog key value
     = MsgWillFetch (Set key)
     | MsgDidFetch (Map key value)

From d81089b8bb205af2f082d0704eac7410aba1c541 Mon Sep 17 00:00:00 2001
From: Johannes Lund <johannes.lund@iohk.io>
Date: Tue, 8 Jun 2021 15:34:15 +0200
Subject: [PATCH 10/10] fixup: use {getCached,fetch}RewardAccountBalance(s)
 naming

---
 lib/core/src/Cardano/Wallet.hs                   |  2 +-
 lib/core/src/Cardano/Wallet/Api/Server.hs        |  4 ++--
 lib/core/src/Cardano/Wallet/Network.hs           |  4 ++--
 .../Wallet/DummyTarget/Primitive/Types.hs        |  4 ++--
 .../src/Cardano/Wallet/Shelley/Network.hs        | 16 ++++++++--------
 5 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/lib/core/src/Cardano/Wallet.hs b/lib/core/src/Cardano/Wallet.hs
index 66cb7320bb9..fe75849e39c 100644
--- a/lib/core/src/Cardano/Wallet.hs
+++ b/lib/core/src/Cardano/Wallet.hs
@@ -1056,7 +1056,7 @@ queryRewardBalance
     -> RewardAccount
     -> ExceptT ErrFetchRewards IO Coin
 queryRewardBalance ctx acct = do
-    liftIO $ getCachedAccountBalance nw acct
+    liftIO $ getCachedRewardAccountBalance nw acct
   where
     nw = ctx ^. networkLayer
 
diff --git a/lib/core/src/Cardano/Wallet/Api/Server.hs b/lib/core/src/Cardano/Wallet/Api/Server.hs
index 446e2b761cf..90d495f1fcf 100644
--- a/lib/core/src/Cardano/Wallet/Api/Server.hs
+++ b/lib/core/src/Cardano/Wallet/Api/Server.hs
@@ -273,7 +273,7 @@ import Cardano.Wallet.Compat
 import Cardano.Wallet.DB
     ( DBFactory (..) )
 import Cardano.Wallet.Network
-    ( NetworkLayer, fetchAccountBalances, timeInterpreter )
+    ( NetworkLayer, fetchRewardAccountBalances, timeInterpreter )
 import Cardano.Wallet.Primitive.AddressDerivation
     ( DelegationAddress (..)
     , Depth (..)
@@ -2167,7 +2167,7 @@ listStakeKeys lookupStakeRef ctx (ApiT wid) = do
             liftIO $ listStakeKeys' @n
                 utxo
                 lookupStakeRef
-                (fetchAccountBalances nl)
+                (fetchRewardAccountBalances nl)
                 ourKeys
   where
     nl = ctx ^. networkLayer
diff --git a/lib/core/src/Cardano/Wallet/Network.hs b/lib/core/src/Cardano/Wallet/Network.hs
index e93e9b9ce4d..8df7cd1e9ac 100644
--- a/lib/core/src/Cardano/Wallet/Network.hs
+++ b/lib/core/src/Cardano/Wallet/Network.hs
@@ -175,7 +175,7 @@ data NetworkLayer m block = NetworkLayer
         :: Coin -- Stake to consider for rewards
         -> m StakePoolsSummary
 
-    , getCachedAccountBalance
+    , getCachedRewardAccountBalance
         :: RewardAccount
         -> m Coin
         -- ^ Return the cached reward balance of an account.
@@ -184,7 +184,7 @@ data NetworkLayer m block = NetworkLayer
         -- account to the internal set of observed account, such that it will be
         -- fetched later.
 
-    , fetchAccountBalances
+    , fetchRewardAccountBalances
         :: Set RewardAccount
         -> m (Map RewardAccount Coin)
         -- ^ Fetch the reward account balance of a set of accounts without
diff --git a/lib/core/test/shared/Cardano/Wallet/DummyTarget/Primitive/Types.hs b/lib/core/test/shared/Cardano/Wallet/DummyTarget/Primitive/Types.hs
index de00e0a928d..6daf56c9470 100644
--- a/lib/core/test/shared/Cardano/Wallet/DummyTarget/Primitive/Types.hs
+++ b/lib/core/test/shared/Cardano/Wallet/DummyTarget/Primitive/Types.hs
@@ -165,8 +165,8 @@ dummyNetworkLayer = NetworkLayer
     , currentSlottingParameters = error "currentSlottingParameters: not implemented"
     , postTx = error "postTx: not implemented"
     , stakeDistribution = error "stakeDistribution: not implemented"
-    , getCachedAccountBalance = error "getCachedAccountBalance: not implemented"
-    , fetchAccountBalances = error "fetchAccountBalances: not implemented"
+    , getCachedRewardAccountBalance = error "getRewardCachedAccountBalance: not implemented"
+    , fetchRewardAccountBalances = error "fetchRewardAccountBalances: not implemented"
     , timeInterpreter = error "timeInterpreter: not implemented"
     , syncProgress = error "syncProgress: not implemented"
     }
diff --git a/lib/shelley/src/Cardano/Wallet/Shelley/Network.hs b/lib/shelley/src/Cardano/Wallet/Shelley/Network.hs
index 5ebee4418f4..0613d7b87ef 100644
--- a/lib/shelley/src/Cardano/Wallet/Shelley/Network.hs
+++ b/lib/shelley/src/Cardano/Wallet/Shelley/Network.hs
@@ -355,9 +355,9 @@ withNetworkLayerBase tr np conn (versionData, _) tol action = do
             _postTx localTxSubmissionQ era sealed
         , stakeDistribution =
             _stakeDistribution queryRewardQ
-        , getCachedAccountBalance =
-            _getAccountBalance rewardsObserver
-        , fetchAccountBalances =
+        , getCachedRewardAccountBalance =
+            _getCachedAccountBalance rewardsObserver
+        , fetchRewardAccountBalances =
             fetchRewardAccounts tr queryRewardQ
         , timeInterpreter =
             _timeInterpreter (contramap MsgInterpreterLog tr) interpreterVar
@@ -547,7 +547,7 @@ withNetworkLayerBase tr np conn (versionData, _) tol action = do
 
     -- TODO(#2042): Make wallets call manually, with matching
     -- stopObserving.
-    _getAccountBalance rewardsObserver k = do
+    _getCachedRewardAccountBalance rewardsObserver k = do
         startObserving rewardsObserver k
         fromMaybe (W.Coin 0) <$> query rewardsObserver k
 
@@ -866,7 +866,7 @@ fetchRewardAccounts
     -> IO (Map W.RewardAccount W.Coin)
 fetchRewardAccounts tr queryRewardQ accounts = do
         liftIO $ traceWith tr $
-            MsgGetRewardAccountBalance accounts
+            MsgFetchRewardAccountBalance accounts
 
         let qry = byronOrShelleyBased (pure (byronValue, [])) $
                    fmap fromBalanceResult
@@ -1154,7 +1154,7 @@ data NetworkLayerLog where
     MsgProtocolParameters :: W.ProtocolParameters -> W.SlottingParameters -> NetworkLayerLog
     MsgLocalStateQueryError :: QueryClientName -> String -> NetworkLayerLog
     MsgLocalStateQueryEraMismatch :: MismatchEraInfo (CardanoEras StandardCrypto) -> NetworkLayerLog
-    MsgGetRewardAccountBalance
+    MsgFetchRewardAccountBalance
         :: Set W.RewardAccount
         -> NetworkLayerLog
     MsgAccountDelegationAndRewards
@@ -1234,7 +1234,7 @@ instance ToText NetworkLayerLog where
         MsgLocalStateQueryEraMismatch mismatch ->
             "Local state query for the wrong era - this is fine. " <>
             T.pack (show mismatch)
-        MsgGetRewardAccountBalance accts -> T.unwords
+        MsgGetchRewardAccountBalance accts -> T.unwords
             [ "Querying the reward account balance for"
             , fmt $ listF accts
             ]
@@ -1303,7 +1303,7 @@ instance HasSeverityAnnotation NetworkLayerLog where
             | isSlowQuery qry dt           -> Notice
             | otherwise                    -> Debug
         MsgInterpreterLog msg              -> getSeverityAnnotation msg
-        MsgGetRewardAccountBalance{}       -> Debug
+        MsgFetchRewardAccountBalance{}       -> Debug
         MsgObserverLog (MsgDidChange _)    -> Notice
         MsgObserverLog{}                   -> Debug