diff --git a/packages/snaps-rpc-methods/src/permitted/getState.ts b/packages/snaps-rpc-methods/src/permitted/getState.ts index b0380f98ad..7c364bed01 100644 --- a/packages/snaps-rpc-methods/src/permitted/getState.ts +++ b/packages/snaps-rpc-methods/src/permitted/getState.ts @@ -19,6 +19,7 @@ import { hasProperty, isObject } from '@metamask/utils'; import { manageStateBuilder } from '../restricted/manageState'; import type { + GetUnlockPromise, JsonRpcRequestWithOrigin, SnapControllerGetSnapStateAction, } from '../types'; @@ -30,12 +31,7 @@ const hookNames: MethodHooksObject = { }; export type GetStateMethodHooks = { - /** - * Wait for the extension to be unlocked. - * - * @returns A promise that resolves once the extension is unlocked. - */ - getUnlockPromise: (shouldShowUnlockRequest: boolean) => Promise; + getUnlockPromise: GetUnlockPromise; }; export type GetStateMethodActions = @@ -134,7 +130,7 @@ async function getStateImplementation( const { key, encrypted = true } = validatedParams; if (encrypted) { - await getUnlockPromise(true); + await getUnlockPromise(true, true); } const state = await messenger.call( diff --git a/packages/snaps-rpc-methods/src/permitted/listEntropySources.test.ts b/packages/snaps-rpc-methods/src/permitted/listEntropySources.test.ts index 0094187c0f..2efd0b5d3e 100644 --- a/packages/snaps-rpc-methods/src/permitted/listEntropySources.test.ts +++ b/packages/snaps-rpc-methods/src/permitted/listEntropySources.test.ts @@ -104,7 +104,7 @@ describe('snap_listEntropySources', () => { method: 'snap_listEntropySources', }); - expect(getUnlockPromise).toHaveBeenCalledWith(true); + expect(getUnlockPromise).toHaveBeenCalledWith(true, true); expect(response).toStrictEqual({ jsonrpc: '2.0', id: 1, diff --git a/packages/snaps-rpc-methods/src/permitted/listEntropySources.ts b/packages/snaps-rpc-methods/src/permitted/listEntropySources.ts index dc05169064..c36db099b8 100644 --- a/packages/snaps-rpc-methods/src/permitted/listEntropySources.ts +++ b/packages/snaps-rpc-methods/src/permitted/listEntropySources.ts @@ -17,6 +17,7 @@ import { getBip32PublicKeyBuilder } from '../restricted/getBip32PublicKey'; import { getBip44EntropyBuilder } from '../restricted/getBip44Entropy'; import { getEntropyBuilder } from '../restricted/getEntropy'; import type { + GetUnlockPromise, JsonRpcRequestWithOrigin, KeyringControllerGetStateAction, } from '../types'; @@ -43,12 +44,7 @@ const hookNames: MethodHooksObject = { }; export type ListEntropySourcesMethodHooks = { - /** - * Wait for the extension to be unlocked. - * - * @returns A promise that resolves once the extension is unlocked. - */ - getUnlockPromise: (shouldShowUnlockRequest: boolean) => Promise; + getUnlockPromise: GetUnlockPromise; }; export type ListEntropySourcesMethodActions = @@ -139,7 +135,7 @@ async function listEntropySourcesImplementation( return end(providerErrors.unauthorized()); } - await getUnlockPromise(true); + await getUnlockPromise(true, true); const { keyrings } = messenger.call('KeyringController:getState'); diff --git a/packages/snaps-rpc-methods/src/permitted/setState.ts b/packages/snaps-rpc-methods/src/permitted/setState.ts index 96842dabe7..0a50aacdca 100644 --- a/packages/snaps-rpc-methods/src/permitted/setState.ts +++ b/packages/snaps-rpc-methods/src/permitted/setState.ts @@ -28,6 +28,7 @@ import { STORAGE_SIZE_LIMIT, } from '../restricted/manageState'; import type { + GetUnlockPromise, JsonRpcRequestWithOrigin, SnapControllerGetSnapAction, SnapControllerGetSnapStateAction, @@ -41,12 +42,7 @@ const hookNames: MethodHooksObject = { }; export type SetStateMethodHooks = { - /** - * Wait for the extension to be unlocked. - * - * @returns A promise that resolves once the extension is unlocked. - */ - getUnlockPromise: (shouldShowUnlockRequest: boolean) => Promise; + getUnlockPromise: GetUnlockPromise; }; export type SetStateMethodActions = @@ -192,7 +188,7 @@ async function setStateImplementation( } if (encrypted) { - await getUnlockPromise(true); + await getUnlockPromise(true, true); } const snapId = origin as SnapId; diff --git a/packages/snaps-rpc-methods/src/restricted/getBip32Entropy.ts b/packages/snaps-rpc-methods/src/restricted/getBip32Entropy.ts index a55ca235c1..5399eefd66 100644 --- a/packages/snaps-rpc-methods/src/restricted/getBip32Entropy.ts +++ b/packages/snaps-rpc-methods/src/restricted/getBip32Entropy.ts @@ -16,7 +16,10 @@ import { SnapCaveatType } from '@metamask/snaps-utils'; import type { NonEmptyArray } from '@metamask/utils'; import { assert } from '@metamask/utils'; -import type { KeyringControllerWithKeyringAction } from '../types'; +import type { + GetUnlockPromise, + KeyringControllerWithKeyringAction, +} from '../types'; import type { MethodHooksObject } from '../utils'; import { getMnemonic, @@ -29,12 +32,7 @@ import { const targetName = 'snap_getBip32Entropy'; export type GetBip32EntropyMethodHooks = { - /** - * Waits for the extension to be unlocked. - * - * @returns A promise that resolves once the extension is unlocked. - */ - getUnlockPromise: (shouldShowUnlockRequest: boolean) => Promise; + getUnlockPromise: GetUnlockPromise; /** * Get the cryptographic functions to use for the client. This may return an diff --git a/packages/snaps-rpc-methods/src/restricted/getBip32PublicKey.ts b/packages/snaps-rpc-methods/src/restricted/getBip32PublicKey.ts index 530641859e..c3c2e853f2 100644 --- a/packages/snaps-rpc-methods/src/restricted/getBip32PublicKey.ts +++ b/packages/snaps-rpc-methods/src/restricted/getBip32PublicKey.ts @@ -22,7 +22,10 @@ import { boolean, object, optional, string } from '@metamask/superstruct'; import type { NonEmptyArray } from '@metamask/utils'; import { assertStruct } from '@metamask/utils'; -import type { KeyringControllerWithKeyringAction } from '../types'; +import type { + GetUnlockPromise, + KeyringControllerWithKeyringAction, +} from '../types'; import type { MethodHooksObject } from '../utils'; import { getMnemonic, @@ -35,12 +38,7 @@ import { const targetName = 'snap_getBip32PublicKey'; export type GetBip32PublicKeyMethodHooks = { - /** - * Waits for the extension to be unlocked. - * - * @returns A promise that resolves once the extension is unlocked. - */ - getUnlockPromise: (shouldShowUnlockRequest: boolean) => Promise; + getUnlockPromise: GetUnlockPromise; /** * Get the cryptographic functions to use for the client. This may return an diff --git a/packages/snaps-rpc-methods/src/restricted/getBip44Entropy.ts b/packages/snaps-rpc-methods/src/restricted/getBip44Entropy.ts index e668cf811f..3bf7d8be63 100644 --- a/packages/snaps-rpc-methods/src/restricted/getBip44Entropy.ts +++ b/packages/snaps-rpc-methods/src/restricted/getBip44Entropy.ts @@ -16,19 +16,17 @@ import type { import { SnapCaveatType } from '@metamask/snaps-utils'; import type { NonEmptyArray } from '@metamask/utils'; -import type { KeyringControllerWithKeyringAction } from '../types'; +import type { + GetUnlockPromise, + KeyringControllerWithKeyringAction, +} from '../types'; import type { MethodHooksObject } from '../utils'; import { getMnemonicSeed, getValueFromEntropySource } from '../utils'; const targetName = 'snap_getBip44Entropy'; export type GetBip44EntropyMethodHooks = { - /** - * Waits for the extension to be unlocked. - * - * @returns A promise that resolves once the extension is unlocked. - */ - getUnlockPromise: (shouldShowUnlockRequest: boolean) => Promise; + getUnlockPromise: GetUnlockPromise; /** * Get the cryptographic functions to use for the client. This may return an diff --git a/packages/snaps-rpc-methods/src/restricted/getEntropy.ts b/packages/snaps-rpc-methods/src/restricted/getEntropy.ts index 7a874aad45..3d5dbb9362 100644 --- a/packages/snaps-rpc-methods/src/restricted/getEntropy.ts +++ b/packages/snaps-rpc-methods/src/restricted/getEntropy.ts @@ -14,7 +14,10 @@ import { literal, object, optional, string } from '@metamask/superstruct'; import type { NonEmptyArray } from '@metamask/utils'; import { assertStruct } from '@metamask/utils'; -import type { KeyringControllerWithKeyringAction } from '../types'; +import type { + GetUnlockPromise, + KeyringControllerWithKeyringAction, +} from '../types'; import type { MethodHooksObject } from '../utils'; import { deriveEntropyFromSeed, @@ -119,12 +122,7 @@ export const getEntropyBuilder = Object.freeze({ } as const); export type GetEntropyHooks = { - /** - * Waits for the extension to be unlocked. - * - * @returns A promise that resolves once the extension is unlocked. - */ - getUnlockPromise: (shouldShowUnlockRequest: boolean) => Promise; + getUnlockPromise: GetUnlockPromise; /** * Get the cryptographic functions to use for the client. This may return an @@ -169,7 +167,7 @@ function getEntropyImplementation({ rpcErrors.invalidParams, ); - await getUnlockPromise(true); + await getUnlockPromise(true, true); const seed = await getValueFromEntropySource( getMnemonicSeed.bind(null, messenger), diff --git a/packages/snaps-rpc-methods/src/restricted/manageAccounts.ts b/packages/snaps-rpc-methods/src/restricted/manageAccounts.ts index 52eadc9340..e4f6103fcb 100644 --- a/packages/snaps-rpc-methods/src/restricted/manageAccounts.ts +++ b/packages/snaps-rpc-methods/src/restricted/manageAccounts.ts @@ -20,6 +20,8 @@ import { import type { Json, NonEmptyArray } from '@metamask/utils'; import { JsonStruct } from '@metamask/utils'; +import type { GetUnlockPromise } from '../types'; + const SnapMessageStruct = union([ object({ method: string(), @@ -45,12 +47,7 @@ export type ManageAccountsMethodHooks = { ) => Promise; }>; - /** - * Wait for the client to be unlocked. - * - * @returns A promise that resolves once the client is unlocked. - */ - getUnlockPromise: (shouldShowUnlockRequest: boolean) => Promise; + getUnlockPromise: GetUnlockPromise; }; type ManageAccountsSpecificationBuilderOptions = { @@ -115,7 +112,7 @@ export function manageAccountsImplementation({ assert(params, SnapMessageStruct); - await getUnlockPromise(true); + await getUnlockPromise(true, true); const keyring = await getSnapKeyring(origin); return await keyring.handleKeyringSnapMessage(origin, params); diff --git a/packages/snaps-rpc-methods/src/restricted/manageState.ts b/packages/snaps-rpc-methods/src/restricted/manageState.ts index 775c699fb1..6b044e05e1 100644 --- a/packages/snaps-rpc-methods/src/restricted/manageState.ts +++ b/packages/snaps-rpc-methods/src/restricted/manageState.ts @@ -17,6 +17,7 @@ import type { NonEmptyArray } from '@metamask/utils'; import { isObject, isValidJson } from '@metamask/utils'; import type { + GetUnlockPromise, SnapControllerClearSnapStateAction, SnapControllerGetSnapAction, SnapControllerGetSnapStateAction, @@ -31,12 +32,7 @@ export const STATE_ENCRYPTION_SALT = 'snap_manageState encryption'; const methodName = 'snap_manageState'; export type ManageStateMethodHooks = { - /** - * Waits for the extension to be unlocked. - * - * @returns A promise that resolves once the extension is unlocked. - */ - getUnlockPromise: (shouldShowUnlockRequest: boolean) => Promise; + getUnlockPromise: GetUnlockPromise; }; export type ManageStateMessengerActions = diff --git a/packages/snaps-rpc-methods/src/types.ts b/packages/snaps-rpc-methods/src/types.ts index ebeeabdd7d..6b51cbdf55 100644 --- a/packages/snaps-rpc-methods/src/types.ts +++ b/packages/snaps-rpc-methods/src/types.ts @@ -234,3 +234,8 @@ export type WebSocketServiceGetAllAction = { type: `WebSocketService:getAll`; handler: (snapId: string) => GetWebSocketsResult; }; + +export type GetUnlockPromise = ( + shouldShowUnlockRequest: boolean, + skipCount?: boolean, +) => Promise; diff --git a/packages/snaps-simulation/src/simulation.ts b/packages/snaps-simulation/src/simulation.ts index f4c9465582..a6c75a3fe4 100644 --- a/packages/snaps-simulation/src/simulation.ts +++ b/packages/snaps-simulation/src/simulation.ts @@ -179,7 +179,10 @@ export type PermittedMiddlewareHooks = { * @param shouldShowUnlockRequest - Whether to show the unlock request. * @returns A promise that resolves once the extension is unlocked. */ - getUnlockPromise: (shouldShowUnlockRequest: boolean) => Promise; + getUnlockPromise: ( + shouldShowUnlockRequest: boolean, + skipCount?: boolean, + ) => Promise; /** * A hook that returns whether the client is active or not.