From eed460c03a8ce2f5093e59f7d58acf3422653d49 Mon Sep 17 00:00:00 2001
From: KtorZ <matthias.benkort@gmail.com>
Date: Thu, 13 Aug 2020 15:04:40 +0200
Subject: [PATCH] fix random address import to not mark every imported address
 as 'Used'

---
 lib/core/src/Cardano/Wallet.hs                | 18 +++++++++-------
 lib/core/src/Cardano/Wallet/Api/Server.hs     |  3 ++-
 .../Primitive/AddressDiscovery/Random.hs      | 21 +++++++++++++++++++
 3 files changed, 33 insertions(+), 9 deletions(-)

diff --git a/lib/core/src/Cardano/Wallet.hs b/lib/core/src/Cardano/Wallet.hs
index 2c45871a531..5bb8f677c7e 100644
--- a/lib/core/src/Cardano/Wallet.hs
+++ b/lib/core/src/Cardano/Wallet.hs
@@ -100,6 +100,7 @@ module Cardano.Wallet
     , normalizeDelegationAddress
     , ErrCreateRandomAddress(..)
     , ErrImportRandomAddress(..)
+    , ErrImportAddress(..)
 
     -- ** Payment
     , selectCoinsExternal
@@ -236,7 +237,7 @@ import Cardano.Wallet.Primitive.AddressDiscovery
     , KnownAddresses (..)
     )
 import Cardano.Wallet.Primitive.AddressDiscovery.Random
-    ( RndState )
+    ( ErrImportAddress (..), RndState )
 import Cardano.Wallet.Primitive.AddressDiscovery.Sequential
     ( SeqState
     , defaultAddressPoolGap
@@ -1143,12 +1144,13 @@ importRandomAddresses
 importRandomAddresses ctx wid addrs = db & \DBLayer{..} -> mapExceptT atomically $ do
     cp <- withExceptT ErrImportAddrNoSuchWallet
         $ withNoSuchWallet wid (readCheckpoint (PrimaryKey wid))
-    let s = getState cp
-        ours = scanl' (\(_, t) addr -> isOurs addr t) (True, s) addrs
-        s' = snd (last ours)
-    if (not . any fst) ours
-        then throwE ErrImportAddrDoesNotBelong
-        else withExceptT ErrImportAddrNoSuchWallet $
+    let s0 = getState cp
+        ours = scanl' (\s addr -> s >>= Rnd.importAddress addr) (Right s0) addrs
+    case last ours of
+        Left err ->
+            throwE $ ErrImportAddr err
+        Right s' ->
+            withExceptT ErrImportAddrNoSuchWallet $
                 putCheckpoint (PrimaryKey wid) (updateState s' cp)
   where
     db = ctx ^. dbLayer @s @k
@@ -2287,7 +2289,7 @@ data ErrCreateRandomAddress
 
 data ErrImportRandomAddress
     = ErrImportAddrNoSuchWallet ErrNoSuchWallet
-    | ErrImportAddrDoesNotBelong
+    | ErrImportAddr ErrImportAddress
     | ErrImportAddressNotAByronWallet
     deriving (Generic, Eq, Show)
 
diff --git a/lib/core/src/Cardano/Wallet/Api/Server.hs b/lib/core/src/Cardano/Wallet/Api/Server.hs
index 9baf4892c15..d7f5bfbe2fe 100644
--- a/lib/core/src/Cardano/Wallet/Api/Server.hs
+++ b/lib/core/src/Cardano/Wallet/Api/Server.hs
@@ -106,6 +106,7 @@ import Cardano.Wallet
     , ErrDecodeSignedTx (..)
     , ErrFetchRewards (..)
     , ErrGetTransaction (..)
+    , ErrImportAddress (..)
     , ErrImportRandomAddress (..)
     , ErrJoinStakePool (..)
     , ErrListTransactions (..)
@@ -2445,7 +2446,7 @@ instance LiftHandler ErrImportRandomAddress where
                 [ "I cannot derive new address for this wallet type."
                 , " Make sure to use Byron random wallet id."
                 ]
-        ErrImportAddrDoesNotBelong ->
+        ErrImportAddr ErrAddrDoesNotBelong{} ->
             apiError err403 KeyNotFoundForAddress $ mconcat
                 [ "I couldn't identify this address as one of mine. It likely "
                 , "belongs to another wallet and I will therefore not import it."
diff --git a/lib/core/src/Cardano/Wallet/Primitive/AddressDiscovery/Random.hs b/lib/core/src/Cardano/Wallet/Primitive/AddressDiscovery/Random.hs
index b2481c5e09c..a1370f5c873 100644
--- a/lib/core/src/Cardano/Wallet/Primitive/AddressDiscovery/Random.hs
+++ b/lib/core/src/Cardano/Wallet/Primitive/AddressDiscovery/Random.hs
@@ -28,6 +28,8 @@ module Cardano.Wallet.Primitive.AddressDiscovery.Random
     , DerivationPath
 
     -- ** Low-level API
+    , importAddress
+    , ErrImportAddress(..)
     , addDiscoveredAddress
     , deriveRndStateAddress
     , findUnusedPath
@@ -156,6 +158,25 @@ mkRndState key seed = RndState
     , gen = mkStdGen seed
     }
 
+newtype ErrImportAddress
+    = ErrAddrDoesNotBelong Address
+    deriving (Generic, Eq, Show)
+
+-- | Import an address into the state. This fails if the address does not belong
+-- to the wallet. Import an address that is already known is a no-op.
+importAddress
+    :: Address
+    -> RndState n
+    -> Either ErrImportAddress (RndState n)
+importAddress addr s = do
+    case addressToPath addr (hdPassphrase s) of
+        Nothing ->
+            Left (ErrAddrDoesNotBelong addr)
+        Just path | Map.member path (addresses s) ->
+            Right s
+        Just path ->
+            Right (addDiscoveredAddress addr Unused path s)
+
 -- | Updates a 'RndState' by adding an address and its derivation path to the
 -- set of discovered addresses. If the address was in the 'pendingAddresses' set
 -- (i.e. it was a newly generated change address), then it is removed from