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] 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