feat: add new RPC starkNet_switchAccount

packages/starknet-snap/src/index.tsx

@@ -39,6 +39,7 @@ import type {
   AddAccountParams,
   GetCurrentAccountParams,
   ListAccountsParams,
+  SwitchAccountParams,
 } from './rpcs';
 import {
   displayPrivateKey,
@@ -58,6 +59,7 @@ import {
   addAccount,
   getCurrentAccount,
   listAccounts,
+  switchAccount,
 } from './rpcs';
 import { signDeployAccountTransaction } from './signDeployAccountTransaction';
 import type {
@@ -296,6 +298,11 @@ export const onRpcRequest: OnRpcRequestHandler = async ({
           requestParams as unknown as GetCurrentAccountParams,
         );
 
+      case RpcMethod.SwitchAccount:
+        return await switchAccount.execute(
+          requestParams as unknown as SwitchAccountParams,
+        );
+
       default:
         throw new MethodNotFoundError() as unknown as Error;
     }

packages/starknet-snap/src/rpcs/index.ts

@@ -15,3 +15,4 @@ export * from './list-transactions';
 export * from './add-account';
 export * from './get-current-account';
 export * from './list-accounts';
+export * from './switch-account';

packages/starknet-snap/src/rpcs/switch-account.test.ts

@@ -0,0 +1,52 @@
+import type { constants } from 'starknet';
+
+import { STARKNET_SEPOLIA_TESTNET_NETWORK } from '../utils/constants';
+import { InvalidRequestParamsError } from '../utils/exceptions';
+import { AccountService } from '../wallet/account';
+import { setupAccountController } from './__tests__/helper';
+import { switchAccount } from './switch-account';
+import type { SwitchAccountParams } from './switch-account';
+
+jest.mock('../utils/snap');
+jest.mock('../utils/logger');
+
+describe('SwitchAccountRpc', () => {
+  const network = STARKNET_SEPOLIA_TESTNET_NETWORK;
+
+  const setupSwitchAccountTest = async () => {
+    const { account } = await setupAccountController({});
+
+    const switchAccountSpy = jest.spyOn(
+      AccountService.prototype,
+      'switchAccount',
+    );
+    switchAccountSpy.mockReturnThis();
+
+    const request = {
+      chainId: network.chainId as unknown as constants.StarknetChainId,
+      address: account.address,
+    };
+
+    return {
+      switchAccountSpy,
+      account,
+      request,
+    };
+  };
+
+  it('switch the current account and returns the account', async () => {
+    const { account, request, switchAccountSpy } =
+      await setupSwitchAccountTest();
+
+    const result = await switchAccount.execute(request);
+
+    expect(result).toStrictEqual(await account.serialize());
+    expect(switchAccountSpy).toHaveBeenCalledWith(network.chainId, account);
+  });
+
+  it('throws `InvalidRequestParamsError` when request parameter is not correct', async () => {
+    await expect(
+      switchAccount.execute({} as unknown as SwitchAccountParams),
+    ).rejects.toThrow(InvalidRequestParamsError);
+  });
+});

packages/starknet-snap/src/rpcs/switch-account.ts

@@ -0,0 +1,58 @@
+import { assign, object, type Infer } from 'superstruct';
+
+import { BaseRequestStruct, AccountStruct, AddressStruct } from '../utils';
+import { createAccountService } from '../utils/factory';
+import { AccountRpcController } from './abstract/account-rpc-controller';
+
+export const SwitchAccountRequestStruct = assign(
+  BaseRequestStruct,
+  object({
+    address: AddressStruct,
+  }),
+);
+
+export const SwitchAccountResponseStruct = AccountStruct;
+
+export type SwitchAccountParams = Infer<typeof SwitchAccountRequestStruct>;
+
+export type SwitchAccountResponse = Infer<typeof SwitchAccountResponseStruct>;
+
+/**
+ * The RPC handler to switch a account by network.
+ */
+export class SwitchAccountRpc extends AccountRpcController<
+  SwitchAccountParams,
+  SwitchAccountResponse
+> {
+  protected requestStruct = SwitchAccountRequestStruct;
+
+  protected responseStruct = SwitchAccountResponseStruct;
+
+  /**
+   * Execute the switch account request handler.
+   *
+   * @param params - The parameters of the request.
+   * @param params.chainId - The chain id of the network.
+   * @param params.address - The address of the account to switch to.
+   * @returns A promise that resolves to the switched account.
+   */
+  protected async handleRequest(
+    params: SwitchAccountParams,
+  ): Promise<SwitchAccountResponse> {
+    const accountService = createAccountService(this.network);
+
+    await accountService.switchAccount(params.chainId, this.account);
+
+    return (await this.account.serialize()) as unknown as SwitchAccountResponse;
+  }
+
+  // Switching an account does not require any verification.
+  // Hence, we overrided the `verifyAccount` method to mute the error,
+  // in case the account to switch for requires deploy/upgrade.
+  // eslint-disable-next-line @typescript-eslint/no-empty-function
+  protected override async verifyAccount(): Promise<void> {
+    // Do not throw any error.
+  }
+}
+
+export const switchAccount = new SwitchAccountRpc();

packages/starknet-snap/src/state/account-state-manager.test.ts

@@ -8,6 +8,7 @@ import {
   mockState,
 } from './__tests__/helper';
 import { AccountStateManager } from './account-state-manager';
+import { StateManagerError } from './state-manager';
 
 describe('AccountStateManager', () => {
   const testnetChainId = constants.StarknetChainId.SN_SEPOLIA;
@@ -293,4 +294,73 @@ describe('AccountStateManager', () => {
       expect(result).toBeNull();
     });
   });
+
+  describe('switchAccount', () => {
+    const setupSwitchAccountTest = async () => {
+      const [testnetCurrentAccount] = await generateTestnetAccounts();
+      const state = await mockStateWithMainnetAccounts([testnetCurrentAccount]);
+
+      return {
+        testnetCurrentAccount,
+        state,
+      };
+    };
+
+    it('switchs the current account', async () => {
+      const { testnetCurrentAccount, state } = await setupSwitchAccountTest();
+      // simulate the account to switch for that contains updated data.
+      const updatedAccountToSwitch = {
+        ...testnetCurrentAccount,
+        upgradeRequired: true,
+      };
+
+      const stateManager = new AccountStateManager();
+      await stateManager.switchAccount({
+        chainId: testnetChainId,
+        accountToSwitch: updatedAccountToSwitch,
+      });
+
+      const updatedAccountFromState = state.accContracts.find(
+        (acc) =>
+          acc.chainId === updatedAccountToSwitch.chainId &&
+          acc.address === updatedAccountToSwitch.address,
+      );
+
+      expect(state.currentAccount).toHaveProperty(testnetChainId);
+      expect(state.currentAccount[testnetChainId]).toStrictEqual(
+        updatedAccountToSwitch,
+      );
+      expect(updatedAccountFromState).toStrictEqual(updatedAccountToSwitch);
+    });
+
+    it('throws `Account does not exist` error if the account to switch for is not exist', async () => {
+      const { testnetCurrentAccount } = await setupSwitchAccountTest();
+      const accountNotExist = {
+        ...testnetCurrentAccount,
+        address: '0x123456789',
+      };
+
+      const stateManager = new AccountStateManager();
+      await expect(
+        stateManager.switchAccount({
+          chainId: testnetChainId,
+          accountToSwitch: accountNotExist,
+        }),
+      ).rejects.toThrow(new StateManagerError('Account does not exist'));
+    });
+
+    it('throws `Account to switch is not in the same chain` error if the account to switch for does not has the same chain Id as the given chain Id', async () => {
+      const { testnetCurrentAccount } = await setupSwitchAccountTest();
+
+      const stateManager = new AccountStateManager();
+      await expect(
+        stateManager.switchAccount({
+          chainId: mainnetChainId,
+          accountToSwitch: testnetCurrentAccount,
+        }),
+      ).rejects.toThrow(
+        new StateManagerError('Account to switch is not in the same chain'),
+      );
+    });
+  });
 });

packages/starknet-snap/src/state/account-state-manager.ts

@@ -259,4 +259,59 @@ export class AccountStateManager extends StateManager<AccContract> {
 
     return data.currentAccount?.[chainId] ?? null;
   }
+
+  /**
+   * Switches the current account for the specified chain ID.
+   * And updates the `AccContract` data for the current account from state.
+   *
+   * @param params - The parameters for switching the current account.
+   * @param params.chainId - The chain ID.
+   * @param params.accountToSwitch - The `AccContract` object to switch to.
+   * @throws {StateManagerError} If the account to switch does not exist.
+   */
+  async switchAccount({
+    chainId,
+    accountToSwitch,
+  }: {
+    chainId: string;
+    accountToSwitch: AccContract;
+  }): Promise<void> {
+    try {
+      await this.update(async (state: SnapState) => {
+        const { chainId: accountChainId, address } = accountToSwitch;
+
+        // We should not relied on the `chainId` from the `accountToSwitch` object
+        // therefore it is required to verify if the `accountToSwitch` object
+        // whether it has the same chain ID or not
+        if (chainId !== accountChainId) {
+          throw new Error(`Account to switch is not in the same chain`);
+        }
+
+        // a safeguard to ensure the `accountToSwitch` is exist in the state
+        const accountInState = await this.getAccount(
+          {
+            address,
+            chainId: accountChainId,
+          },
+          state,
+        );
+
+        if (!accountInState) {
+          throw new Error(`Account does not exist`);
+        }
+
+        if (!state.currentAccount) {
+          state.currentAccount = {};
+        }
+
+        state.currentAccount[chainId] = accountToSwitch;
+
+        // due to the account may contains legacy data,
+        // this is a hack to ensure the account is updated
+        this.updateEntity(accountInState, accountToSwitch);
+      });
+    } catch (error) {
+      throw new StateManagerError(error.message);
+    }
+  }
 }

packages/starknet-snap/src/utils/permission.test.ts

@@ -16,6 +16,7 @@ describe('validateOrigin', () => {
     RpcMethod.ReadContract,
     RpcMethod.GetStoredErc20Tokens,
     RpcMethod.AddAccount,
+    RpcMethod.SwitchAccount,
   ];
 
   it.each(walletUIDappPermissions)(

packages/starknet-snap/src/utils/permission.ts

@@ -17,6 +17,7 @@ export enum RpcMethod {
   GetCurrentAccount = 'starkNet_getCurrentAccount',
   ListAccounts = 'starkNet_listAccounts',
 
+  SwitchAccount = 'starkNet_swtichAccount',
   AddAccount = 'starkNet_addAccount',
   CreateAccount = 'starkNet_createAccount',
   DisplayPrivateKey = 'starkNet_displayPrivateKey',
@@ -68,6 +69,7 @@ const walletUIDappPermissions = publicPermissions.concat([
   RpcMethod.ReadContract,
   RpcMethod.GetStoredErc20Tokens,
   RpcMethod.AddAccount,
+  RpcMethod.SwitchAccount,
 ]);
 
 const publicPermissionsSet = new Set(publicPermissions);

packages/starknet-snap/src/wallet/account/service.test.ts

@@ -292,4 +292,24 @@ describe('AccountService', () => {
       expect(result.hdIndex).toStrictEqual(defaultIndex);
     });
   });
+
+  describe('switchAccount', () => {
+    it('switches the account for the network', async () => {
+      const switchAccountSpy = jest.spyOn(
+        AccountStateManager.prototype,
+        'switchAccount',
+      );
+      switchAccountSpy.mockResolvedValue();
+      mockAccountContractReader({});
+      const { accountObj } = await createAccountObject(network, 0);
+
+      const service = createAccountService(network);
+      await service.switchAccount(accountObj.chainId, accountObj);
+
+      expect(switchAccountSpy).toHaveBeenCalledWith({
+        chainId: accountObj.chainId,
+        accountToSwitch: await accountObj.serialize(),
+      });
+    });
+  });
 });

packages/starknet-snap/src/wallet/account/service.ts

@@ -133,4 +133,21 @@ export class AccountService {
       activeAccount ? activeAccount.addressIndex : 0,
     );
   }
+
+  /**
+   * Switches the account for the network.
+   * The account to switch must be in the same chain.
+   *
+   * @param chainId - The chain ID.
+   * @param accountToSwitch - The account to switch to.
+   */
+  async switchAccount(
+    chainId: string,
+    accountToSwitch: Account,
+  ): Promise<void> {
+    await this.accountStateMgr.switchAccount({
+      chainId,
+      accountToSwitch: await accountToSwitch.serialize(),
+    });
+  }
 }