refactor: revamp rpc starkNet_extractPrivateKey (#319)

* refactor: rename RPC method starkNet_extractPrivateKey to starkNet_displayPrivateKey

* fix: use snap dialog helper

* chore: revamp sign transaction

* chore: add comment

* chore: revamp display private key

* fix: apply suggestions from code review

Co-authored-by: Stanley Yuen <[email protected]>

* chore: update sign detail

* chore: update test

* chore: remove unnecessary mock

* chore: simplify test

* refactor: revamp `starkNet_displayPrivateKey`

* chore: update lint and unit test

* chore: update display private key rpc unit test

---------

Co-authored-by: Stanley Yuen <[email protected]>

packages/starknet-snap/index.html

@@ -612,13 +612,13 @@ <h1>Hello, Snaps!</h1>
       await callSnap('starkNet_createAccount', { addressIndex, deploy });
     }
 
-    // here we call the snap's "starkNet_extractPrivateKey" method
+    // here we call the snap's "starkNet_displayPrivateKey" method
     async function retrievePrivateKeyFromAddress(e) {
       e.preventDefault(); // to prevent default form behavior
 
       const userAddress = document.getElementById('userAddress').value;
 
-      await callSnap('starkNet_extractPrivateKey', { userAddress });
+      await callSnap('starkNet_displayPrivateKey', { userAddress });
     }
 
     // here we call the snap's "starkNet_extractPublicKey" method

packages/starknet-snap/openrpc/starknet_snap_api_openrpc.json

@@ -8,7 +8,7 @@
   "servers": [],
   "methods": [
     {
-      "name": "starkNet_extractPrivateKey",
+      "name": "starkNet_displayPrivateKey",
       "summary": "Extract private key from deployed Starknet account and display in MetaMask",
       "paramStructure": "by-name",
       "params": [

packages/starknet-snap/src/index.ts

@@ -23,7 +23,6 @@ import { estimateAccDeployFee } from './estimateAccountDeployFee';
 import { estimateFee } from './estimateFee';
 import { estimateFees } from './estimateFees';
 import { executeTxn } from './executeTxn';
-import { extractPrivateKey } from './extractPrivateKey';
 import { extractPublicKey } from './extractPublicKey';
 import { getCurrentNetwork } from './getCurrentNetwork';
 import { getErc20TokenBalance } from './getErc20TokenBalance';
@@ -36,8 +35,12 @@ import { getTransactions } from './getTransactions';
 import { getTransactionStatus } from './getTransactionStatus';
 import { getValue } from './getValue';
 import { recoverAccounts } from './recoverAccounts';
-import type { SignMessageParams, SignTransactionParams } from './rpcs';
-import { signMessage, signTransaction } from './rpcs';
+import type {
+  DisplayPrivateKeyParams,
+  SignMessageParams,
+  SignTransactionParams,
+} from './rpcs';
+import { displayPrivateKey, signMessage, signTransaction } from './rpcs';
 import { sendTransaction } from './sendTransaction';
 import { signDeclareTransaction } from './signDeclareTransaction';
 import { signDeployAccountTransaction } from './signDeployAccountTransaction';
@@ -161,10 +164,9 @@ export const onRpcRequest: OnRpcRequestHandler = async ({ request }) => {
       case 'starkNet_getStoredUserAccounts':
         return await getStoredUserAccounts(apiParams);
 
-      case 'starkNet_extractPrivateKey':
-        apiParams.keyDeriver = await getAddressKeyDeriver(snap);
-        return await extractPrivateKey(
-          apiParams as unknown as ApiParamsWithKeyDeriver,
+      case 'starkNet_displayPrivateKey':
+        return await displayPrivateKey.execute(
+          apiParams.requestParams as unknown as DisplayPrivateKeyParams,
         );
 
       case 'starkNet_extractPublicKey':

packages/starknet-snap/src/rpcs/__tests__/helper.ts

@@ -48,10 +48,21 @@ export function prepareMockAccount(account: StarknetAccount, state: SnapState) {
 /**
  *
  */
-export function prepareSignConfirmDialog() {
+export function prepareConfirmDialog() {
   const confirmDialogSpy = jest.spyOn(snapHelper, 'confirmDialog');
   confirmDialogSpy.mockResolvedValue(true);
   return {
     confirmDialogSpy,
   };
 }
+
+/**
+ *
+ */
+export function prepareAlertDialog() {
+  const alertDialogSpy = jest.spyOn(snapHelper, 'alertDialog');
+  alertDialogSpy.mockResolvedValue(true);
+  return {
+    alertDialogSpy,
+  };
+}

packages/starknet-snap/src/rpcs/displayPrivateKey.test.ts

@@ -0,0 +1,116 @@
+import {
+  InvalidParamsError,
+  UserRejectedRequestError,
+} from '@metamask/snaps-sdk';
+import { constants } from 'starknet';
+
+import type { SnapState } from '../types/snapState';
+import { STARKNET_SEPOLIA_TESTNET_NETWORK } from '../utils/constants';
+import {
+  mockAccount,
+  prepareAlertDialog,
+  prepareMockAccount,
+  prepareConfirmDialog,
+} from './__tests__/helper';
+import { displayPrivateKey } from './displayPrivateKey';
+import type { DisplayPrivateKeyParams } from './displayPrivateKey';
+
+describe('displayPrivateKey', () => {
+  const state: SnapState = {
+    accContracts: [],
+    erc20Tokens: [],
+    networks: [STARKNET_SEPOLIA_TESTNET_NETWORK],
+    transactions: [],
+  };
+
+  const createRequestParam = (
+    chainId: constants.StarknetChainId,
+    address: string,
+  ): DisplayPrivateKeyParams => {
+    const request: DisplayPrivateKeyParams = {
+      chainId,
+      address,
+    };
+    return request;
+  };
+
+  it('displays private key correctly', async () => {
+    const chainId = constants.StarknetChainId.SN_SEPOLIA;
+    const account = await mockAccount(chainId);
+    prepareMockAccount(account, state);
+    prepareConfirmDialog();
+    const { alertDialogSpy } = prepareAlertDialog();
+
+    const request = createRequestParam(chainId, account.address);
+
+    await displayPrivateKey.execute(request);
+
+    expect(alertDialogSpy).toHaveBeenCalledTimes(1);
+
+    const calls = alertDialogSpy.mock.calls[0][0];
+
+    expect(calls).toStrictEqual([
+      { type: 'text', value: 'Starknet Account Private Key' },
+      { type: 'copyable', value: account.privateKey },
+    ]);
+  });
+
+  it('renders confirmation dialog', async () => {
+    const chainId = constants.StarknetChainId.SN_SEPOLIA;
+    const account = await mockAccount(chainId);
+    prepareMockAccount(account, state);
+    const { confirmDialogSpy } = prepareConfirmDialog();
+    prepareAlertDialog();
+
+    const request = createRequestParam(chainId, account.address);
+
+    await displayPrivateKey.execute(request);
+
+    expect(confirmDialogSpy).toHaveBeenCalledTimes(1);
+
+    const calls = confirmDialogSpy.mock.calls[0][0];
+
+    expect(calls).toStrictEqual([
+      { type: 'text', value: 'Do you want to export your private key?' },
+    ]);
+  });
+
+  it('throws `UserRejectedRequestError` if user denies the operation', async () => {
+    const chainId = constants.StarknetChainId.SN_SEPOLIA;
+    const account = await mockAccount(chainId);
+    prepareMockAccount(account, state);
+    const { confirmDialogSpy } = prepareConfirmDialog();
+    prepareAlertDialog();
+
+    confirmDialogSpy.mockResolvedValue(false);
+
+    const request = createRequestParam(chainId, account.address);
+
+    await expect(displayPrivateKey.execute(request)).rejects.toThrow(
+      UserRejectedRequestError,
+    );
+  });
+
+  it.each([
+    {
+      case: 'user address is omitted',
+      request: {
+        chainId: constants.StarknetChainId.SN_SEPOLIA,
+      },
+    },
+    {
+      case: 'user address is invalid',
+      request: {
+        chainId: constants.StarknetChainId.SN_SEPOLIA,
+        address: 'invalid_address',
+      },
+    },
+  ])(
+    'throws `InvalidParamsError` when $case',
+    async ({ request }: { request: unknown }) => {
+      await expect(
+        displayPrivateKey.execute(request as DisplayPrivateKeyParams),
+      ).rejects.toThrow(InvalidParamsError);
+    },
+  );
+});

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

@@ -0,0 +1,77 @@
+import { copyable, text, UserRejectedRequestError } from '@metamask/snaps-sdk';
+import { type Infer, object, literal, assign } from 'superstruct';
+
+import {
+  AccountRpcController,
+  AddressStruct,
+  confirmDialog,
+  alertDialog,
+  BaseRequestStruct,
+} from '../utils';
+
+export const DisplayPrivateKeyRequestStruct = assign(
+  object({
+    address: AddressStruct,
+  }),
+  BaseRequestStruct,
+);
+
+export const DisplayPrivateKeyResponseStruct = literal(null);
+
+export type DisplayPrivateKeyParams = Infer<
+  typeof DisplayPrivateKeyRequestStruct
+>;
+
+export type DisplayPrivateKeyResponse = Infer<
+  typeof DisplayPrivateKeyResponseStruct
+>;
+
+/**
+ * The RPC handler to display a private key.
+ */
+export class DisplayPrivateKeyRpc extends AccountRpcController<
+  DisplayPrivateKeyParams,
+  DisplayPrivateKeyResponse
+> {
+  protected requestStruct = DisplayPrivateKeyRequestStruct;
+
+  protected responseStruct = DisplayPrivateKeyResponseStruct;
+
+  /**
+   * Execute the display private key request handler.
+   * The private key will be display via a confirmation dialog.
+   *
+   * @param params - The parameters of the request.
+   * @param params.address - The account address.
+   * @param params.chainId - The chain id of the network.
+   */
+  async execute(
+    params: DisplayPrivateKeyParams,
+  ): Promise<DisplayPrivateKeyResponse> {
+    return super.execute(params);
+  }
+
+  protected async handleRequest(
+    // eslint-disable-next-line @typescript-eslint/no-unused-vars
+    params: DisplayPrivateKeyParams,
+  ): Promise<DisplayPrivateKeyResponse> {
+    const confirmComponents = [text('Do you want to export your private key?')];
+
+    if (!(await confirmDialog(confirmComponents))) {
+      throw new UserRejectedRequestError() as unknown as Error;
+    }
+
+    const alertComponents = [
+      text('Starknet Account Private Key'),
+      copyable(this.account.privateKey),
+    ];
+
+    await alertDialog(alertComponents);
+
+    return null;
+  }
+}
+
+export const displayPrivateKey = new DisplayPrivateKeyRpc({
+  showInvalidAccountAlert: false,
+});

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

@@ -1,2 +1,3 @@
 export * from './signMessage';
+export * from './displayPrivateKey';
 export * from './signTransaction';