chore: dont poll requests if the client is locked

packages/snap/snap.manifest.json

@@ -7,7 +7,7 @@
     "url": "git+https://github.com/MetaMask/snap-institutional-wallet.git"
   },
   "source": {
-    "shasum": "gmrs64oSJvzFljwbXMqZIW7HGlmPaECfY1sqo+ItKgk=",
+    "shasum": "61twpj95F7b92y4usr2OoZeGh/KO9tpG5GgWfQjy+G0=",
     "location": {
       "npm": {
         "filePath": "dist/bundle.js",

packages/snap/src/index.test.ts

@@ -6,6 +6,7 @@ import { getKeyring, getRequestManager } from './context';
 import { renderOnboarding } from './features/onboarding/render';
 import { CustodianType } from './lib/types/CustodianType';
 import { InternalMethod, originPermissions } from './permissions';
+import { getClientStatus } from './snap-state-manager/snap-util';
 
 // Mock the context module
 jest.mock('./context', () => ({
@@ -42,6 +43,14 @@ const mockRenderOnboarding = renderOnboarding as jest.MockedFunction<
   typeof renderOnboarding
 >;
 
+// Add this mock at the top of the test file with other mocks
+jest.mock('./snap-state-manager/snap-util');
+
+// Add type assertion for the mock
+const mockGetClientStatus = getClientStatus as jest.MockedFunction<
+  typeof getClientStatus
+>;
+
 describe('index', () => {
   const mockKeyring = {
     listAccounts: jest.fn().mockResolvedValue([]),
@@ -75,6 +84,9 @@ describe('index', () => {
         name: 'Test Account',
       },
     ]);
+
+    // Set default mock return value
+    mockGetClientStatus.mockResolvedValue({ locked: false });
   });
 
   describe('onRpcRequest', () => {
@@ -256,6 +268,35 @@ describe('index', () => {
       jest.clearAllMocks();
     });
 
+    it('should not poll when client is locked', async () => {
+      // Override the mock for this specific test
+      mockGetClientStatus.mockResolvedValueOnce({ locked: true });
+
+      await onCronjob({
+        request: {
+          method: 'execute',
+          id: 1,
+          jsonrpc: '2.0',
+        },
+      });
+
+      expect(mockRequestManager.poll).not.toHaveBeenCalled();
+      expect(setTimeoutSpy).not.toHaveBeenCalled();
+    });
+
+    it('should poll when client is unlocked', async () => {
+      await onCronjob({
+        request: {
+          method: 'execute',
+          id: 1,
+          jsonrpc: '2.0',
+        },
+      });
+
+      expect(mockRequestManager.poll).toHaveBeenCalled();
+      expect(setTimeoutSpy).toHaveBeenCalledTimes(5); // 5 delays between 6 polls
+    });
+
     it('should handle execute method', async () => {
       await onCronjob({
         request: {

packages/snap/src/index.ts

@@ -43,7 +43,7 @@ import type { CustodialKeyringAccount } from './lib/types/CustodialKeyringAccoun
 import { CustodianApiMap, CustodianType } from './lib/types/CustodianType';
 import logger from './logger';
 import { InternalMethod, originPermissions } from './permissions';
-// @audit - this file needs unittests
+import { getClientStatus } from './snap-state-manager/snap-util';
 
 /**
  * Verify if the caller can call the requested method.
@@ -249,6 +249,18 @@ const pollRequests = async (): Promise<void> => {
 };
 
 export const onCronjob: OnCronjobHandler = async ({ request }) => {
+  // Check if the client is locked
+  // If it is, return so as not to try and access the encrypted storage which would cause the user to have to enter their password
+
+  // TODO: In the future, we can have some unencrypted storage that is used to store whether the snap was ever used or not
+  // or even a list of all the accounts and pending requests and then we can optionally unlock to continue polliog
+  // But for now we just keep the client locked with the assumption that the user will unlock it when they want to complete the request
+
+  const clientStatus = await getClientStatus();
+  if (clientStatus.locked) {
+    return;
+  }
+
   switch (
     request.method // @audit-info this is execute every minute * * * * *  acc. to manifest
   ) {

packages/snap/src/snap-state-manager/snap-util.test.ts

@@ -0,0 +1,108 @@
+import { jest } from '@jest/globals';
+
+import { getStateData, setStateData, getClientStatus } from './snap-util';
+
+// Mock the global snap object
+const mockSnap = {
+  request: jest.fn<(...args: any[]) => Promise<any>>(),
+};
+
+(global as any).snap = mockSnap;
+
+describe('snap-util', () => {
+  beforeEach(() => {
+    // Clear mock calls between tests
+    jest.clearAllMocks();
+  });
+
+  describe('getStateData', () => {
+    it('should retrieve snap state successfully', async () => {
+      const mockState = { foo: 'bar' } as any;
+      mockSnap.request.mockResolvedValueOnce(mockState);
+
+      const result = await getStateData(false);
+
+      expect(result).toStrictEqual(mockState);
+      expect(mockSnap.request).toHaveBeenCalledWith({
+        method: 'snap_manageState',
+        params: { operation: 'get', encrypted: false },
+      });
+    });
+
+    it('should return null when no state exists', async () => {
+      mockSnap.request.mockResolvedValueOnce(null);
+
+      const result = await getStateData(false);
+
+      expect(result).toBeNull();
+      expect(mockSnap.request).toHaveBeenCalledWith({
+        method: 'snap_manageState',
+        params: { operation: 'get', encrypted: false },
+      });
+    });
+
+    it('should throw error when snap request fails', async () => {
+      const error = new Error('Failed to get state');
+      mockSnap.request.mockRejectedValueOnce(error);
+
+      await expect(getStateData(false)).rejects.toThrow('Failed to get state');
+    });
+  });
+
+  describe('setStateData', () => {
+    it('should set snap state successfully', async () => {
+      const newState = { foo: 'bar' };
+      mockSnap.request.mockResolvedValueOnce(undefined);
+
+      await setStateData({ data: newState, encrypted: false });
+
+      expect(mockSnap.request).toHaveBeenCalledWith({
+        method: 'snap_manageState',
+        params: {
+          operation: 'update',
+          newState,
+          encrypted: false,
+        },
+      });
+    });
+
+    it('should throw error when snap request fails', async () => {
+      const newState = { foo: 'bar' };
+      const error = new Error('Failed to set state');
+      mockSnap.request.mockRejectedValueOnce(error);
+
+      await expect(
+        setStateData({ data: newState, encrypted: false }),
+      ).rejects.toThrow('Failed to set state');
+    });
+
+    it('should handle setting null state', async () => {
+      mockSnap.request.mockResolvedValueOnce(undefined);
+
+      await setStateData({ data: null, encrypted: false });
+
+      expect(mockSnap.request).toHaveBeenCalledWith({
+        method: 'snap_manageState',
+        params: {
+          operation: 'update',
+          newState: null,
+          encrypted: false,
+        },
+      });
+    });
+  });
+
+  describe('getClientStatus', () => {
+    it('should retrieve client status successfully', async () => {
+      const mockStatus = { locked: true };
+      mockSnap.request.mockResolvedValueOnce(mockStatus);
+
+      const result = await getClientStatus();
+
+      expect(result).toStrictEqual(mockStatus);
+      expect(mockSnap.request).toHaveBeenCalledWith({
+        method: 'snap_getClientStatus',
+      });
+    });
+  });
+});

packages/snap/src/snap-state-manager/snap-util.ts

@@ -1,4 +1,4 @@
-import type { Json } from '@metamask/snaps-sdk';
+import type { GetClientStatusResult, Json } from '@metamask/snaps-sdk';
 
 /**
  * Retrieves the current state data.
@@ -39,3 +39,15 @@ export async function setStateData<State>({
     },
   });
 }
+
+/**
+ * Retrieves the client status (locked/unlocked) in this case from MM.
+ * Used to check if the client is locked or not.
+ *
+ * @returns An object containing the status.
+ */
+export async function getClientStatus(): Promise<GetClientStatusResult> {
+  return await snap.request({
+    method: 'snap_getClientStatus',
+  });
+}