keep track of pending change indexes for transaction signing

src/Cardano/Wallet.hs

@@ -62,7 +62,12 @@ import Cardano.Wallet.Primitive.AddressDerivation
     , publicKey
     )
 import Cardano.Wallet.Primitive.AddressDiscovery
-    ( AddressPoolGap, SeqState (..), generateChangeOutput, mkAddressPool )
+    ( AddressPoolGap
+    , SeqState (..)
+    , emptyPendingIxs
+    , generateChangeOutput
+    , mkAddressPool
+    )
 import Cardano.Wallet.Primitive.Model
     ( Wallet, applyBlocks, availableUTxO, currentTip, getState, initWallet )
 import Cardano.Wallet.Primitive.Signing
@@ -222,6 +227,7 @@ mkWalletLayer db network = WalletLayer
         let checkpoint = initWallet $ SeqState
                 { externalPool = extPool
                 , internalPool = intPool
+                , pendingChangeIxs = emptyPendingIxs
                 }
         now <- liftIO getCurrentTime
         let metadata = WalletMetadata

src/Cardano/Wallet/Primitive/AddressDiscovery.hs

@@ -16,11 +16,11 @@
 -- compatibility with random address scheme from the legacy Cardano wallets.
 
 module Cardano.Wallet.Primitive.AddressDiscovery
-    ( AddressScheme (..)
+    (
     -- * Sequential Derivation
 
     -- ** Address Pool Gap
-    , AddressPoolGap
+      AddressPoolGap
     , MkAddressPoolGapError (..)
     , defaultAddressPoolGap
     , getAddressPoolGap
@@ -35,9 +35,14 @@ module Cardano.Wallet.Primitive.AddressDiscovery
     , mkAddressPool
     , lookupAddress
 
+    -- * Pending Change Indexes
+    , PendingIxs
+    , emptyPendingIxs
+
     -- ** State
     , SeqState (..)
-  ) where
+    , AddressScheme (..)
+    ) where
 
 import Prelude
 
@@ -63,12 +68,8 @@ import Control.DeepSeq
     ( NFData, deepseq )
 import Data.Bifunctor
     ( first )
-import Data.Foldable
-    ( maximum )
 import Data.Function
     ( (&) )
-import Data.List
-    ( sortOn )
 import Data.Map.Strict
     ( Map )
 import Data.Maybe
@@ -82,6 +83,7 @@ import GHC.Generics
 import Text.Read
     ( readMaybe )
 
+import qualified Data.List as L
 import qualified Data.Map.Strict as Map
 import qualified Data.Text as T
 
@@ -95,7 +97,9 @@ import qualified Data.Text as T
 
 -------------------------------------------------------------------------------}
 
--- ** Address Pool Gap
+{-------------------------------------------------------------------------------
+                              Address Pool Gap
+-------------------------------------------------------------------------------}
 
 -- | Maximum number of consecutive undiscovered addresses allowed
 newtype AddressPoolGap = AddressPoolGap
@@ -149,8 +153,9 @@ defaultAddressPoolGap :: AddressPoolGap
 defaultAddressPoolGap =
     AddressPoolGap 20
 
-
--- ** Address Pool
+{-------------------------------------------------------------------------------
+                                 Address Pool
+-------------------------------------------------------------------------------}
 
 -- | An 'AddressPool' which keeps track of sequential addresses within a given
 -- Account and change chain. See 'mkAddressPool' to create a new or existing
@@ -181,7 +186,7 @@ instance NFData AddressPool
 --
 -- > mkAddressPool key g cc (addresses pool) == pool
 addresses :: AddressPool -> [Address]
-addresses = map fst . sortOn snd . Map.toList . indexedAddresses
+addresses = map fst . L.sortOn snd . Map.toList . indexedAddresses
 
 -- | Create a new Address pool from a list of addresses. Note that, the list is
 -- expected to be ordered in sequence (first indexes, first in the list).
@@ -254,31 +259,97 @@ nextAddresses !key (AddressPoolGap !g) !cc !fromIx =
         (>= fromIx)
     newAddress = keyToAddress . deriveAddressPublicKey key cc
 
+{-------------------------------------------------------------------------------
+                            Pending Change Indexes
+-------------------------------------------------------------------------------}
+
+-- | An ordered set of pending indexes. This keep track of indexes used
+newtype PendingIxs = PendingIxs [Index 'Soft 'AddressK]
+    deriving stock (Generic, Show)
+instance NFData PendingIxs
+
+-- | An empty pending set of change indexes.
+--
+-- NOTE: We do not define a 'Monoid' instance here because there's no rational
+-- of combining two pending sets.
+emptyPendingIxs :: PendingIxs
+emptyPendingIxs = PendingIxs mempty
+
+-- | Update the set of pending indexes by discarding every indexes _below_ the
+-- given index.
+--
+-- Why is that?
+--
+-- Because we really do care about the higher index that was last used in order
+-- to know from where we can generate new indexes.
+updatePendingIxs
+    :: Index 'Soft 'AddressK
+    -> PendingIxs
+    -> PendingIxs
+updatePendingIxs ix (PendingIxs ixs) =
+    PendingIxs $ L.filter (> ix) ixs
+
+-- | Get the next change index; If every available indexes have already been
+-- taken, we'll rotate the pending set and re-use already provided indexes.
+--
+-- This should not be a problem for users in practice, and remain okay for
+-- exchanges who care less about privacy / not-reusing addresses than
+-- regular users.
+nextChangeIndex
+    :: AddressPool
+    -> PendingIxs
+    -> (Index 'Soft 'AddressK, PendingIxs)
+nextChangeIndex pool (PendingIxs ixs) =
+    let
+        poolLen = length (addresses pool)
+        (firstUnused, lastUnused) =
+            ( toEnum $ poolLen - fromEnum (gap pool)
+            , toEnum $ poolLen - 1
+            )
+        (ix, ixs') = case ixs of
+            [] ->
+                (firstUnused, PendingIxs [firstUnused])
+            h:_ | length ixs < fromEnum (gap pool) ->
+                (succ h, PendingIxs (succ h:ixs))
+            h:q ->
+                (h, PendingIxs (q++[h]))
+    in
+        invariant "index is within first unused and last unused" (ix, ixs')
+            (\(i,_) -> i >= firstUnused && i <= lastUnused)
+
+{-------------------------------------------------------------------------------
+                                 State
+-------------------------------------------------------------------------------}
+
 data SeqState = SeqState
     { internalPool :: !AddressPool
+        -- ^ Addresses living on the 'InternalChain'
     , externalPool :: !AddressPool
+        -- ^ Addresses living on the 'ExternalChain'
+    , pendingChangeIxs :: !PendingIxs
+        -- ^ Indexes from the internal pool that have been used in pending
+        -- transactions. The list is maintained sorted in descending order
+        -- (cf: 'PendingIxs')
     }
     deriving stock (Generic, Show)
-
 instance NFData SeqState
 
 -- NOTE
 -- We have to scan both the internal and external chain. Note that, the
 -- account discovery algorithm is only specified for the external chain so
 -- in theory, there's nothing forcing a wallet to generate change
 -- addresses on the internal chain anywhere in the available range.
---
--- In practice, we may assume that user can't create change addresses and
--- that they are just created in sequence by the wallet software. Hence an
--- address pool with a gap of 1 should be sufficient for the internal chain.
 instance IsOurs SeqState where
-     isOurs addr (SeqState !s1 !s2) =
+     isOurs addr (SeqState !s1 !s2 !ixs) =
          let
-             (res1, !s1') = lookupAddress addr s1
-             (res2, !s2') = lookupAddress addr s2
-             ours = isJust (res1 <|> res2)
+             (internal, !s1') = lookupAddress addr s1
+             (external, !s2') = lookupAddress addr s2
+             !ixs' = case internal of
+                Nothing -> ixs
+                Just ix -> updatePendingIxs ix ixs
+             ours = isJust (internal <|> external)
          in
-             (ours `deepseq` ours, SeqState s1' s2')
+             (ixs' `deepseq` ours `deepseq` ours, SeqState s1' s2' ixs')
 
 -- TODO: We might want to move this abstraction / there is more work to be
 -- done here.
@@ -292,51 +363,55 @@ class AddressScheme s where
         :: Address
         -> (Key 'RootK XPrv, Passphrase "encryption")
         -> s
-        -> (Maybe (Key 'AddressK XPrv), s)
-
-    generateChangeOutput :: Coin -> s -> (TxOut, s)
+        -> Maybe (Key 'AddressK XPrv)
+        -- ^ Derive the private key corresponding to an address. Careful, this
+        -- operation can be costly. Note that the state is discarded from this
+        -- function as we do not intend to discover any addresses from this
+        -- operation; This is merely a lookup from known addresses.
+
+    generateChangeOutput
+        :: Coin
+        -> s
+        -> (TxOut, s)
+        -- ^ Generate a change output 'TxOut' from a given 'Coin'. This picks
+        -- up the first non-used known change address and use it. We keep track
+        -- of pending indexes in the state.
 
 instance AddressScheme SeqState where
-    keyFrom addr (rootPrv, pwd) (SeqState !s1 !s2) =
+    keyFrom addr (rootPrv, pwd) (SeqState !s1 !s2 _) =
         let
-            (xPrv1, !s1') = lookupAndDeriveXPrv s1 InternalChain
-            (xPrv2, !s2') = lookupAndDeriveXPrv s2 ExternalChain
-
-            xPrv = (xPrv1 <|> xPrv2)
-
+            xPrv1 = lookupAndDeriveXPrv s1 InternalChain
+            xPrv2 = lookupAndDeriveXPrv s2 ExternalChain
+            xPrv = xPrv1 <|> xPrv2
         in
-            (xPrv `deepseq` xPrv, SeqState s1' s2')
-
+            xPrv
       where
-        -- We are assuming there is only one account
-        account = minBound
-        accountPrv = deriveAccountPrivateKey pwd rootPrv account
-
         lookupAndDeriveXPrv
             :: AddressPool
             -> ChangeChain
-            -> (Maybe (Key 'AddressK XPrv), AddressPool)
+            -> Maybe (Key 'AddressK XPrv)
         lookupAndDeriveXPrv pool chain =
             let
-                (addrIx, pool') = lookupAddress addr pool
+                -- We are assuming there is only one account
+                accountPrv = deriveAccountPrivateKey pwd rootPrv minBound
+                (addrIx, _) = lookupAddress addr pool
             in
-                (deriveAddressPrivateKey pwd accountPrv chain <$> addrIx, pool')
-
-    generateChangeOutput c (SeqState intPool extPool) =
+                deriveAddressPrivateKey pwd accountPrv chain <$> addrIx
+
+    -- | We pick indexes in sequence from the first known available index (i.e.
+    -- @length addrs - gap@) but we do not generate _new change addresses_. As a
+    -- result, we can't generate more than @gap@ _pending_ change addresses and
+    -- therefore, rotate the change addresses when we need extra change outputs.
+    --
+    -- See also: 'nextChangeIndex'
+    generateChangeOutput c (SeqState intPool extPool pending) =
         let
-            indexes = map snd $ Map.toList $ indexedAddresses intPool
-            newIndex = succ $ maximum indexes
-            xPub = deriveAddressPublicKey
-                (accountPubKey intPool)
-                InternalChain
-                newIndex
-
-            next = Map.singleton addr newIndex
-            intPool' =
-                intPool { indexedAddresses = indexedAddresses intPool <> next }
-            addr = keyToAddress xPub
-
-        in (TxOut addr c, SeqState intPool' extPool)
-
-
-
+            (ix, pending') = nextChangeIndex intPool pending
+            accountXPub = accountPubKey intPool
+            addressXPub = deriveAddressPublicKey accountXPub InternalChain ix
+            txout = TxOut
+                { address = keyToAddress addressXPub
+                , coin = c
+                }
+        in
+            (txout, SeqState intPool extPool pending')

src/Cardano/Wallet/Primitive/Signing.hs

@@ -40,14 +40,12 @@ import Control.Monad
     ( forM )
 import Crypto.Hash
     ( Blake2b_256, hash )
-import Data.ByteArray
-    ( convert )
 import Data.ByteString
     ( ByteString )
 
 import qualified Cardano.Crypto.Wallet as CC
 import qualified Codec.CBOR.Encoding as CBOR
-
+import qualified Data.ByteArray as BA
 
 newtype SignTxError = KeyNotFoundForAddress Address
 
@@ -65,37 +63,25 @@ mkStdTx
     -> [TxOut]
     -- ^ Selected outputs (including change)
     -> Either SignTxError (Tx, [TxWitness])
-mkStdTx s (rootPrv, pwd) ownedIns outs = do
+mkStdTx s creds@(_, pwd) ownedIns outs = do
     let ins = (fmap fst ownedIns)
         tx = Tx ins outs
         txSigData = hashTx tx
-
-    txWitnesses <- forM ownedIns (\(_in, TxOut addr _c) ->
-        mkWitness txSigData <$> keyFrom' addr)
-
+    txWitnesses <- forM ownedIns $ \(_in, TxOut addr _c) -> mkWitness txSigData
+        <$> withEither (KeyNotFoundForAddress addr) (keyFrom addr creds s)
     return (tx, txWitnesses)
-
   where
-    keyFrom' addr =
-        -- We are ignoring the new state/pool. We won't discover any new
-        -- addresses when submitting transactions.
-        case (fst $ keyFrom addr (rootPrv, pwd) s) of
-            Just key -> Right key
-            Nothing -> Left $ KeyNotFoundForAddress addr
-
+    withEither :: e -> Maybe a -> Either e a
+    withEither e = maybe (Left e) Right
     hashTx :: Tx -> Hash "tx"
     hashTx txSigData = Hash
-        $ convert
-        $ (hash @ByteString @Blake2b_256)
+        $ BA.convert
+        $ (hash @_ @Blake2b_256)
         $ toByteString
         $ encodeTx txSigData
-
-
     mkWitness :: Hash "tx" -> Key 'AddressK XPrv -> TxWitness
-    mkWitness tx xPrv =
-        PublicKeyWitness $
-            encodeXPub (publicKey xPrv) <>
-            getHash (sign (SignTx tx) (xPrv, pwd))
+    mkWitness tx xPrv = PublicKeyWitness $
+        encodeXPub (publicKey xPrv) <> getHash (sign (SignTx tx) (xPrv, pwd))