fix: move ephemeral state out of persisted wallet store (#1899)

Author: tyleroooo

src/bonsai/rest/compliance.ts

@@ -13,6 +13,7 @@ import {
   setLocalAddressScreenV2Raw,
   setSourceAddressScreenV2Raw,
 } from '@/state/raw';
+import { getHdKeyNonce } from '@/state/walletSelectors';
 
 import { signCompliancePayload } from '@/lib/compliance';
 import { mapIfPresent } from '@/lib/do';
@@ -62,9 +63,10 @@ export function setUpIndexerSourceAddressScreenV2Query(store: RootStore) {
 }
 
 const selectChainIdAndLocalAddress = createAppSelector(
-  [getSelectedDydxChainId, getUserWalletAddress, getSelectedNetwork],
-  (chainId, address, network) => ({
+  [getSelectedDydxChainId, getUserWalletAddress, getSelectedNetwork, getHdKeyNonce],
+  (chainId, address, network, hdKeyNonce) => ({
     chainId,
+    hdKeyNonce,
     address,
     network,
   })
@@ -81,12 +83,14 @@ const COMPLIANCE_PAYLOAD_MESSAGE = 'Verify account ownership';
 async function updateCompliance({
   chainId,
   address,
+  hdKeyNonce,
   network,
   status,
   action,
 }: {
   chainId: DydxChainId;
   address: string;
+  hdKeyNonce: number;
   network: DydxNetwork;
   status: ComplianceStatus;
   action: ComplianceAction;
@@ -104,7 +108,7 @@ async function updateCompliance({
     chainId,
   };
 
-  const signingResponse = await signCompliancePayload(address, payload);
+  const signingResponse = await signCompliancePayload(address, hdKeyNonce, payload);
   if (!signingResponse) {
     return { status: ComplianceStatus.UNKNOWN };
   }
@@ -163,8 +167,8 @@ export function setUpIndexerLocalAddressScreenV2Query(store: RootStore) {
       address,
       network,
     ],
-    getQueryFn: (indexerClient, { chainId, address, network }) => {
-      if (address == null) {
+    getQueryFn: (indexerClient, { chainId, address, network, hdKeyNonce }) => {
+      if (address == null || hdKeyNonce == null) {
         return null;
       }
       return async (): Promise<ComplianceResponse> => {
@@ -179,6 +183,7 @@ export function setUpIndexerLocalAddressScreenV2Query(store: RootStore) {
 
         const updateResult = updateCompliance({
           address,
+          hdKeyNonce,
           chainId,
           network,
           status: firstScreenResult.status,
@@ -211,15 +216,17 @@ export const triggerCompliance = (action: ComplianceAction) => {
       const currentLocalScreenStatus = selectRawLocalAddressScreenV2(state).data?.status;
       const chainId = getSelectedDydxChainId(state);
       const address = getUserWalletAddress(state);
+      const hdKeyNonce = getHdKeyNonce(state);
       const network = getSelectedNetwork(state);
 
-      if (!address || !currentLocalScreenStatus) {
+      if (!address || !currentLocalScreenStatus || !hdKeyNonce) {
         throw new Error('TriggerCompliance: No account connected or screen status not loaded');
       }
 
       const result = await updateCompliance({
         chainId,
         address,
+        hdKeyNonce,
         network,
         status: currentLocalScreenStatus,
         action,

src/lib/affiliates.ts

@@ -8,6 +8,7 @@ import { store } from '@/state/_store';
 import { getUserWalletAddress } from '@/state/accountInfoSelectors';
 import { appQueryClient } from '@/state/appQueryClient';
 import { getSelectedDydxChainId, getSelectedNetwork } from '@/state/appSelectors';
+import { getHdKeyNonce } from '@/state/walletSelectors';
 
 import { track } from './analytics/analytics';
 import { signCompliancePayload } from './compliance';
@@ -34,9 +35,10 @@ export const updateReferralCode = async (newCode: string) => {
     const chainId = getSelectedDydxChainId(state);
     const address = getUserWalletAddress(state);
     const network = getSelectedNetwork(state);
+    const hdKeyNonce = getHdKeyNonce(state);
 
-    if (!address) {
-      throw new Error('No account connected');
+    if (!address || !hdKeyNonce) {
+      throw new Error('No account connected or hdKey not set');
     }
 
     const networkConfig = ENVIRONMENT_CONFIG_MAP[network];
@@ -46,7 +48,7 @@ export const updateReferralCode = async (newCode: string) => {
       throw new Error('No indexer URL available');
     }
 
-    const signingResponse = await signCompliancePayload(address, {
+    const signingResponse = await signCompliancePayload(address, hdKeyNonce, {
       message: newCode,
       action: ReferralAction.UPDATE_CODE,
       status: '',

src/lib/compliance.ts

@@ -87,6 +87,7 @@ const signComplianceSignatureKeplr = async (
 
 export const signCompliancePayload = async (
   address: string,
+  nonce: number,
   params: {
     message: string;
     action: string;
@@ -95,7 +96,7 @@ export const signCompliancePayload = async (
   }
 ): Promise<string> => {
   try {
-    const hdkey = hdKeyManager.getHdkey(address);
+    const hdkey = hdKeyManager.getHdkey(address, nonce);
     if (hdkey?.privateKey && hdkey.publicKey) {
       const { signedMessage, timestamp } = await signComplianceSignature(
         params.message,

src/lib/hdKeyManager.ts

@@ -4,7 +4,7 @@ import { LocalWallet } from '@dydxprotocol/v4-client-js';
 import { Hdkey } from '@/constants/account';
 
 import type { RootStore } from '@/state/_store';
-import { setLocalWalletNonce } from '@/state/wallet';
+import { setHdKeyNonce, setLocalWalletNonce } from '@/state/wallet';
 
 import { log } from './telemetry';
 
@@ -13,22 +13,38 @@ class HDKeyManager {
 
   private hdkey: Hdkey | undefined;
 
+  private hdKeyNonce: number | undefined;
+
+  private store: RootStore | undefined;
+
+  setStore(store: RootStore) {
+    this.store = store;
+  }
+
   setHdkey(address: string | undefined, hdkey: Hdkey) {
+    this.hdKeyNonce = this.hdKeyNonce != null ? this.hdKeyNonce + 1 : 0;
     this.address = address;
     this.hdkey = hdkey;
+
+    if (!this.store) {
+      log('HDKeyManager: store has not been set');
+      return;
+    }
+    this.store.dispatch(setHdKeyNonce(this.hdKeyNonce));
   }
 
-  getHdkey(localWalletAddress: string): Hdkey | undefined {
-    if (localWalletAddress !== this.address) {
+  getHdkey(localWalletAddress: string, hdKeyNonce: number): Hdkey | undefined {
+    if (localWalletAddress !== this.address || hdKeyNonce !== this.hdKeyNonce) {
       return undefined;
     }
-
     return this.hdkey;
   }
 
   clearHdkey() {
     this.hdkey = undefined;
     this.address = undefined;
+    this.hdKeyNonce = undefined;
+    this.store?.dispatch(setHdKeyNonce(undefined));
   }
 }
 
@@ -54,9 +70,10 @@ class LocalWalletManager {
 
     if (!this.store) {
       log('LocalWalletManager: store has not been set');
+      return;
     }
 
-    this.store?.dispatch(setLocalWalletNonce(this.localWalletNonce));
+    this.store.dispatch(setLocalWalletNonce(this.localWalletNonce));
   }
 
   getLocalWallet(localWalletNonce: number): LocalWallet | undefined {

src/state/_store.ts

@@ -6,7 +6,7 @@ import { persistReducer, persistStore } from 'redux-persist';
 import autoMergeLevel2 from 'redux-persist/lib/stateReconciler/autoMergeLevel2';
 import storage from 'redux-persist/lib/storage';
 
-import { localWalletManager } from '@/lib/hdKeyManager';
+import { hdKeyManager, localWalletManager } from '@/lib/hdKeyManager';
 import { transformOntologyObject } from '@/lib/transformOntology';
 
 import { accountSlice } from './account';
@@ -33,7 +33,7 @@ import { tradingViewSlice } from './tradingView';
 import { transfersSlice } from './transfers';
 import { triggersFormSlice } from './triggersForm';
 import { vaultsSlice } from './vaults';
-import { walletSlice } from './wallet';
+import { walletEphemeralSlice, walletSlice } from './wallet';
 
 const reducers = {
   account: accountSlice.reducer,
@@ -55,6 +55,7 @@ const reducers = {
   transfers: transfersSlice.reducer,
   vaults: vaultsSlice.reducer,
   wallet: walletSlice.reducer,
+  walletEphemeral: walletEphemeralSlice.reducer,
   raw: rawSlice.reducer,
 } as const;
 
@@ -124,6 +125,7 @@ export const persistor = persistStore(store);
 
 // Set store so localWalletManager classes can getState and dispatch
 localWalletManager.setStore(store);
+hdKeyManager.setStore(store);
 
 export type RootStore = typeof store;
 export type RootState = ReturnType<typeof store.getState>;

src/state/wallet.ts

@@ -19,7 +19,6 @@ export interface WalletState {
     address?: string;
     subaccountNumber?: number;
   };
-  localWalletNonce?: number;
   turnkeyEmailOnboardingData?: TurnkeyEmailOnboardingData;
   turnkeyPrimaryWallet?: TurnkeyWallet;
 }
@@ -35,7 +34,6 @@ const initialState: WalletState = {
     address: undefined,
     subaccountNumber: 0,
   },
-  localWalletNonce: undefined,
   turnkeyEmailOnboardingData: undefined,
   turnkeyPrimaryWallet: undefined,
 };
@@ -44,9 +42,6 @@ export const walletSlice = createSlice({
   name: 'wallet',
   initialState,
   reducers: {
-    setLocalWalletNonce: (state, action: PayloadAction<number | undefined>) => {
-      state.localWalletNonce = action.payload;
-    },
     setSourceAddress: (
       state,
       action: PayloadAction<{ address: string; chain: WalletNetworkType }>
@@ -120,8 +115,30 @@ export const walletSlice = createSlice({
   },
 });
 
+export interface WalletEphemeralState {
+  localWalletNonce?: number;
+  hdKeyNonce?: number;
+}
+
+const initialEphemeralState: WalletEphemeralState = {
+  localWalletNonce: undefined,
+  hdKeyNonce: undefined,
+};
+
+export const walletEphemeralSlice = createSlice({
+  name: 'walletEphemeral',
+  initialState: initialEphemeralState,
+  reducers: {
+    setLocalWalletNonce: (state, action: PayloadAction<number | undefined>) => {
+      state.localWalletNonce = action.payload;
+    },
+    setHdKeyNonce: (state, action: PayloadAction<number | undefined>) => {
+      state.hdKeyNonce = action.payload;
+    },
+  },
+});
+
 export const {
-  setLocalWalletNonce,
   setSourceAddress,
   setRequiresAddressUpload,
   setWalletInfo,
@@ -134,3 +151,5 @@ export const {
   setTurnkeyPrimaryWallet,
   clearTurnkeyPrimaryWallet,
 } = walletSlice.actions;
+
+export const { setLocalWalletNonce, setHdKeyNonce } = walletEphemeralSlice.actions;

src/state/walletSelectors.ts

@@ -31,7 +31,8 @@ export const selectTurnkeyWalletInfo = createAppSelector([selectWalletInfo], (wa
   return undefined;
 });
 
-export const getLocalWalletNonce = (state: RootState) => state.wallet.localWalletNonce;
+export const getLocalWalletNonce = (state: RootState) => state.walletEphemeral.localWalletNonce;
+export const getHdKeyNonce = (state: RootState) => state.walletEphemeral.hdKeyNonce;
 
 export const getTurnkeyEmailOnboardingData = (state: RootState) =>
   state.wallet.turnkeyEmailOnboardingData;