Merge pull request #1469 from input-output-hk/fix/dynamic-change-address-resolver-should-always-pick-lower-derivation-index

AngelCastilloB · web-flow · commit 7539b97dd37a · 2024-09-12T15:16:55.000+08:00
fix(wallet): dynamicChangeResolver gives preference to lower derivation indices
diff --git a/packages/wallet/src/services/AddressTracker.ts b/packages/wallet/src/services/AddressTracker.ts
@@ -17,6 +17,7 @@ import {
 } from 'rxjs';
 import { WalletStores } from '../persistence';
 import { groupedAddressesEquals } from './util';
+import { sortAddresses } from './util/sortAddresses';
 
 export type AddressTrackerDependencies = {
   store: WalletStores['addresses'];
@@ -86,7 +87,7 @@ export const createAddressTracker = ({ addressDiscovery$, store, logger }: Addre
         take(1)
       );
     },
-    addresses$,
+    addresses$: addresses$.pipe(map(sortAddresses)),
     shutdown: newAddresses$.complete.bind(newAddresses$)
   };
 };
diff --git a/packages/wallet/src/services/ChangeAddress/DynamicChangeAddressResolver.ts b/packages/wallet/src/services/ChangeAddress/DynamicChangeAddressResolver.ts
@@ -7,6 +7,7 @@ import { GroupedAddress } from '@cardano-sdk/key-management';
 import { InvalidStateError } from '@cardano-sdk/util';
 import { Logger } from 'ts-log';
 import { Observable, firstValueFrom } from 'rxjs';
+import { sortAddresses } from '../util/sortAddresses';
 import isEqual from 'lodash/isEqual.js';
 import uniq from 'lodash/uniq.js';
 
@@ -315,7 +316,8 @@ export class DynamicChangeAddressResolver implements ChangeAddressResolver {
   async resolve(selection: Selection): Promise<Array<Cardano.TxOut>> {
     const delegationDistribution = [...(await firstValueFrom(this.#delegationDistribution)).values()];
     let portfolio = await this.#getDelegationPortfolio();
-    const addresses = await firstValueFrom(this.#addresses$);
+    const addresses = sortAddresses(await firstValueFrom(this.#addresses$));
+
     let updatedChange = [...selection.change];
 
     if (addresses.length === 0) throw new InvalidStateError('The wallet has no known addresses.');
diff --git a/packages/wallet/src/services/util/sortAddresses.ts b/packages/wallet/src/services/util/sortAddresses.ts
@@ -0,0 +1,29 @@
+import { GroupedAddress } from '@cardano-sdk/key-management';
+
+/**
+ * Sorts an array of addresses by their primary index and, if available, by the
+ * index of their stakeKeyDerivationPath.
+ *
+ * @param addresses - The array of addresses to sort.
+ * @returns A new sorted array of addresses.
+ */
+export const sortAddresses = (addresses: GroupedAddress[]): GroupedAddress[] =>
+  [...addresses].sort((a, b) => {
+    if (a.index !== b.index) {
+      return a.index - b.index;
+    }
+
+    if (a.stakeKeyDerivationPath && b.stakeKeyDerivationPath) {
+      return a.stakeKeyDerivationPath.index - b.stakeKeyDerivationPath.index;
+    }
+
+    if (a.stakeKeyDerivationPath && !b.stakeKeyDerivationPath) {
+      return -1;
+    }
+
+    if (!a.stakeKeyDerivationPath && b.stakeKeyDerivationPath) {
+      return 1;
+    }
+
+    return 0;
+  });
diff --git a/packages/wallet/test/services/AddressTracker.test.ts b/packages/wallet/test/services/AddressTracker.test.ts
@@ -1,14 +1,44 @@
 import { AddressTracker, createAddressTracker } from '../../src';
+import { AddressType, GroupedAddress, KeyRole } from '@cardano-sdk/key-management';
 import { Cardano } from '@cardano-sdk/core';
 import { EMPTY, firstValueFrom, of } from 'rxjs';
-import { GroupedAddress } from '@cardano-sdk/key-management';
 import { WalletStores } from '../../src/persistence';
+import { address_0_0, address_0_5, address_5_0, rewardAccount_0, rewardAccount_5 } from './ChangeAddress/testData';
 import { createTestScheduler, logger } from '@cardano-sdk/util-dev';
+import { sortAddresses } from '../../src/services/util/sortAddresses';
 
 describe('AddressTracker', () => {
   let store: jest.Mocked<WalletStores['addresses']>;
   let addressTracker: AddressTracker;
-  const discoveredAddresses = [{ address: 'addr1' as Cardano.PaymentAddress } as GroupedAddress];
+  const discoveredAddresses = [
+    {
+      accountIndex: 0,
+      address: address_5_0,
+      index: 5,
+      networkId: Cardano.NetworkId.Testnet,
+      rewardAccount: rewardAccount_0,
+      stakeKeyDerivationPath: { index: 0, role: KeyRole.Stake },
+      type: AddressType.External
+    },
+    {
+      accountIndex: 0,
+      address: address_0_5,
+      index: 0,
+      networkId: Cardano.NetworkId.Testnet,
+      rewardAccount: rewardAccount_5,
+      stakeKeyDerivationPath: { index: 5, role: KeyRole.Stake },
+      type: AddressType.External
+    },
+    {
+      accountIndex: 0,
+      address: address_0_0,
+      index: 0,
+      networkId: Cardano.NetworkId.Testnet,
+      rewardAccount: rewardAccount_0,
+      stakeKeyDerivationPath: { index: 0, role: KeyRole.Stake },
+      type: AddressType.External
+    }
+  ];
 
   beforeEach(() => {
     store = {
@@ -19,6 +49,8 @@ describe('AddressTracker', () => {
     };
   });
 
+  const sortedDiscoveredAddresses = sortAddresses(discoveredAddresses);
+
   afterEach(() => addressTracker.shutdown());
 
   describe('load', () => {
@@ -31,7 +63,7 @@ describe('AddressTracker', () => {
             logger,
             store
           });
-          expectObservable(addressTracker.addresses$).toBe('a', { a: discoveredAddresses });
+          expectObservable(addressTracker.addresses$).toBe('a', { a: sortedDiscoveredAddresses });
           expectSubscriptions(addressDiscovery$.subscriptions).toBe('^');
           flush();
           expect(store.get).toBeCalledTimes(1);
@@ -70,12 +102,12 @@ describe('AddressTracker', () => {
         logger,
         store
       });
-      await expect(firstValueFrom(addressTracker.addresses$)).resolves.toEqual(discoveredAddresses);
+      await expect(firstValueFrom(addressTracker.addresses$)).resolves.toEqual(sortedDiscoveredAddresses);
       const newAddress = { address: 'addr2' as Cardano.PaymentAddress } as GroupedAddress;
       const combinedAddresses = [...discoveredAddresses, newAddress];
 
       await expect(firstValueFrom(addressTracker.addAddresses([newAddress]))).resolves.toEqual(combinedAddresses);
-      await expect(firstValueFrom(addressTracker.addresses$)).resolves.toEqual(combinedAddresses);
+      await expect(firstValueFrom(addressTracker.addresses$)).resolves.toEqual(sortAddresses(combinedAddresses));
       expect(store.set).toBeCalledTimes(2);
       expect(store.set).toBeCalledWith(combinedAddresses);
     });
diff --git a/packages/wallet/test/services/ChangeAddress/DynamicChangeAddressResolver.test.ts b/packages/wallet/test/services/ChangeAddress/DynamicChangeAddressResolver.test.ts
@@ -25,7 +25,8 @@ import {
   rewardAccount_0,
   rewardAccount_1,
   rewardAccount_2,
-  rewardAccount_3
+  rewardAccount_3,
+  unorderedKnownAddresses$
 } from './testData';
 import { logger } from '@cardano-sdk/util-dev';
 
@@ -154,7 +155,7 @@ describe('DynamicChangeAddressResolver', () => {
 
   it('adds all change outputs at payment_stake address 0 if the wallet is currently not delegating to any pool', async () => {
     const changeAddressResolver = new DynamicChangeAddressResolver(
-      knownAddresses$,
+      unorderedKnownAddresses$,
       createMockDelegateTracker(new Map<Cardano.PoolId, DelegatedStake>([])).distribution$,
       getNullDelegationPortfolio,
       logger
@@ -191,7 +192,7 @@ describe('DynamicChangeAddressResolver', () => {
 
   it('distributes change equally between the currently delegated addresses if no portfolio is given, ', async () => {
     const changeAddressResolver = new DynamicChangeAddressResolver(
-      knownAddresses$,
+      unorderedKnownAddresses$,
       createMockDelegateTracker(
         new Map<Cardano.PoolId, DelegatedStake>([
           [
diff --git a/packages/wallet/test/services/ChangeAddress/testData.ts b/packages/wallet/test/services/ChangeAddress/testData.ts
@@ -212,6 +212,108 @@ export const knownAddresses$ = new BehaviorSubject<GroupedAddress[]>([
   }
 ]);
 
+export const unorderedKnownAddresses$ = new BehaviorSubject<GroupedAddress[]>([
+  {
+    accountIndex: 0,
+    address: address_5_0,
+    index: 5,
+    networkId: Cardano.NetworkId.Testnet,
+    rewardAccount: rewardAccount_0,
+    stakeKeyDerivationPath: { index: 0, role: KeyRole.Stake },
+    type: AddressType.External
+  },
+  {
+    accountIndex: 0,
+    address: address_1_0,
+    index: 1,
+    networkId: Cardano.NetworkId.Testnet,
+    rewardAccount: rewardAccount_0,
+    stakeKeyDerivationPath: { index: 0, role: KeyRole.Stake },
+    type: AddressType.External
+  },
+  {
+    accountIndex: 0,
+    address: address_0_3,
+    index: 0,
+    networkId: Cardano.NetworkId.Testnet,
+    rewardAccount: rewardAccount_3,
+    stakeKeyDerivationPath: { index: 3, role: KeyRole.Stake },
+    type: AddressType.External
+  },
+  {
+    accountIndex: 0,
+    address: address_0_1,
+    index: 0,
+    networkId: Cardano.NetworkId.Testnet,
+    rewardAccount: rewardAccount_1,
+    stakeKeyDerivationPath: { index: 1, role: KeyRole.Stake },
+    type: AddressType.External
+  },
+  {
+    accountIndex: 0,
+    address: address_0_2,
+    index: 0,
+    networkId: Cardano.NetworkId.Testnet,
+    rewardAccount: rewardAccount_2,
+    stakeKeyDerivationPath: { index: 2, role: KeyRole.Stake },
+    type: AddressType.External
+  },
+  {
+    accountIndex: 0,
+    address: address_2_0,
+    index: 2,
+    networkId: Cardano.NetworkId.Testnet,
+    rewardAccount: rewardAccount_0,
+    stakeKeyDerivationPath: { index: 0, role: KeyRole.Stake },
+    type: AddressType.External
+  },
+  {
+    accountIndex: 0,
+    address: address_0_4,
+    index: 0,
+    networkId: Cardano.NetworkId.Testnet,
+    rewardAccount: rewardAccount_4,
+    stakeKeyDerivationPath: { index: 4, role: KeyRole.Stake },
+    type: AddressType.External
+  },
+  {
+    accountIndex: 0,
+    address: address_3_0,
+    index: 3,
+    networkId: Cardano.NetworkId.Testnet,
+    rewardAccount: rewardAccount_0,
+    stakeKeyDerivationPath: { index: 0, role: KeyRole.Stake },
+    type: AddressType.External
+  },
+  {
+    accountIndex: 0,
+    address: address_0_5,
+    index: 0,
+    networkId: Cardano.NetworkId.Testnet,
+    rewardAccount: rewardAccount_5,
+    stakeKeyDerivationPath: { index: 5, role: KeyRole.Stake },
+    type: AddressType.External
+  },
+  {
+    accountIndex: 0,
+    address: address_0_0,
+    index: 0,
+    networkId: Cardano.NetworkId.Testnet,
+    rewardAccount: rewardAccount_0,
+    stakeKeyDerivationPath: { index: 0, role: KeyRole.Stake },
+    type: AddressType.External
+  },
+  {
+    accountIndex: 0,
+    address: address_4_0,
+    index: 4,
+    networkId: Cardano.NetworkId.Testnet,
+    rewardAccount: rewardAccount_0,
+    stakeKeyDerivationPath: { index: 0, role: KeyRole.Stake },
+    type: AddressType.External
+  }
+]);
+
 export const emptyKnownAddresses$ = new BehaviorSubject<GroupedAddress[]>([]);
 
 export const createMockDelegateTracker = (delegatedStake: Map<Cardano.PoolId, DelegatedStake>) => ({
