feat: validate snap response using superstruct (#15)

* fix: fix `keyring_listAccounts` struct and assertions

* fix: check if the request is a valid JSON-RPC request

* test: add unit tests to dispatcher

* feat: add `strictMask` function

* feat: validate response from snap using superstruct

src/keyring-client.test.ts

@@ -1,6 +1,5 @@
 import {
   KeyringAccount,
-  KeyringJsonRpcRequest,
   KeyringRequest,
   SubmitRequestResponse,
   KeyringClient,
@@ -176,11 +175,15 @@ describe('KeyringClient', () => {
   describe('getRequest', () => {
     it('should send a request to get a request and return the response', async () => {
       const id = '71621d8d-62a4-4bf4-97cc-fb8f243679b0';
-      const expectedResponse: KeyringJsonRpcRequest = {
-        jsonrpc: '2.0',
-        id,
-        method: 'personal_sign',
-        params: ['0xe9a74aacd7df8112911ca93260fc5a046f8a64ae', '0x0'],
+      const expectedResponse: KeyringRequest = {
+        account: '46b5ccd3-4786-427c-89d2-cef626dffe9b',
+        scope: 'eip155:1',
+        request: {
+          jsonrpc: '2.0',
+          id,
+          method: 'personal_sign',
+          params: ['0xe9a74aacd7df8112911ca93260fc5a046f8a64ae', '0x0'],
+        },
       };
 
       mockSender.send.mockResolvedValue(expectedResponse);

src/keyring-client.ts

@@ -7,12 +7,26 @@ import {
   KeyringAccount,
   KeyringRequest,
   SubmitRequestResponse,
+  SubmitRequestResponseStruct,
 } from './keyring-api';
-import { InternalRequest, InternalRequestStruct } from './keyring-internal-api';
-import { OmitUnion } from './utils';
+import {
+  ApproveRequestResponseStruct,
+  CreateAccountResponseStruct,
+  DeleteAccountResponseStruct,
+  FilterAccountChainsResponseStruct,
+  GetAccountResponseStruct,
+  GetRequestResponseStruct,
+  InternalRequest,
+  InternalResponse,
+  ListAccountsResponseStruct,
+  ListRequestsResponseStruct,
+  RejectRequestResponseStruct,
+  UpdateAccountResponseStruct,
+} from './keyring-internal-api';
+import { OmitUnion, strictMask } from './utils';
 
 export type Sender = {
-  send<Response extends Json>(request: InternalRequest): Promise<Response>;
+  send(request: InternalRequest): Promise<InternalResponse>;
 };
 
 export class KeyringClient implements Keyring {
@@ -33,93 +47,124 @@ export class KeyringClient implements Keyring {
    * @param partial - Partial internal request (method and params).
    * @returns A promise that resolves to the response to the request.
    */
-  async #send<Response extends Json>(
+  async #send(
     partial: OmitUnion<InternalRequest, 'jsonrpc' | 'id'>,
-  ): Promise<Response> {
-    const request = {
+  ): Promise<InternalResponse> {
+    return await this.#sender.send({
       jsonrpc: '2.0',
       id: uuid(),
       ...partial,
-    };
-    assert(request, InternalRequestStruct);
-    return await this.#sender.send<Response>(request);
+    });
   }
 
   async listAccounts(): Promise<KeyringAccount[]> {
-    return await this.#send<KeyringAccount[]>({
-      method: 'keyring_listAccounts',
-    });
+    return strictMask(
+      await this.#send({
+        method: 'keyring_listAccounts',
+      }),
+      ListAccountsResponseStruct,
+    );
   }
 
   async getAccount(id: string): Promise<KeyringAccount> {
-    return await this.#send<KeyringAccount>({
-      method: 'keyring_getAccount',
-      params: { id },
-    });
+    return strictMask(
+      await this.#send({
+        method: 'keyring_getAccount',
+        params: { id },
+      }),
+      GetAccountResponseStruct,
+    );
   }
 
   async createAccount(
     name: string,
     options: Record<string, Json> | null = null,
   ): Promise<KeyringAccount> {
-    return await this.#send<KeyringAccount>({
-      method: 'keyring_createAccount',
-      params: { name, options },
-    });
+    return strictMask(
+      await this.#send({
+        method: 'keyring_createAccount',
+        params: { name, options },
+      }),
+      CreateAccountResponseStruct,
+    );
   }
 
   async filterAccountChains(id: string, chains: string[]): Promise<string[]> {
-    return await this.#send<string[]>({
-      method: 'keyring_filterAccountChains',
-      params: { id, chains },
-    });
+    return strictMask(
+      await this.#send({
+        method: 'keyring_filterAccountChains',
+        params: { id, chains },
+      }),
+      FilterAccountChainsResponseStruct,
+    );
   }
 
   async updateAccount(account: KeyringAccount): Promise<void> {
-    await this.#send<null>({
-      method: 'keyring_updateAccount',
-      params: { account },
-    });
+    assert(
+      await this.#send({
+        method: 'keyring_updateAccount',
+        params: { account },
+      }),
+      UpdateAccountResponseStruct,
+    );
   }
 
   async deleteAccount(id: string): Promise<void> {
-    await this.#send<null>({
-      method: 'keyring_deleteAccount',
-      params: { id },
-    });
+    assert(
+      await this.#send({
+        method: 'keyring_deleteAccount',
+        params: { id },
+      }),
+      DeleteAccountResponseStruct,
+    );
   }
 
   async listRequests(): Promise<KeyringRequest[]> {
-    return await this.#send<KeyringRequest[]>({
-      method: 'keyring_listRequests',
-    });
+    return strictMask(
+      await this.#send({
+        method: 'keyring_listRequests',
+      }),
+      ListRequestsResponseStruct,
+    );
   }
 
   async getRequest(id: string): Promise<KeyringRequest> {
-    return await this.#send<KeyringRequest>({
-      method: 'keyring_getRequest',
-      params: { id },
-    });
+    return strictMask(
+      await this.#send({
+        method: 'keyring_getRequest',
+        params: { id },
+      }),
+      GetRequestResponseStruct,
+    );
   }
 
   async submitRequest(request: KeyringRequest): Promise<SubmitRequestResponse> {
-    return await this.#send<SubmitRequestResponse>({
-      method: 'keyring_submitRequest',
-      params: request,
-    });
+    return strictMask(
+      await this.#send({
+        method: 'keyring_submitRequest',
+        params: request,
+      }),
+      SubmitRequestResponseStruct,
+    );
   }
 
   async approveRequest(id: string): Promise<void> {
-    await this.#send<null>({
-      method: 'keyring_approveRequest',
-      params: { id },
-    });
+    assert(
+      await this.#send({
+        method: 'keyring_approveRequest',
+        params: { id },
+      }),
+      ApproveRequestResponseStruct,
+    );
   }
 
   async rejectRequest(id: string): Promise<void> {
-    await this.#send<null>({
-      method: 'keyring_rejectRequest',
-      params: { id },
-    });
+    assert(
+      await this.#send({
+        method: 'keyring_rejectRequest',
+        params: { id },
+      }),
+      RejectRequestResponseStruct,
+    );
   }
 }

src/keyring-snap-controller-client.test.ts

@@ -61,6 +61,7 @@ describe('KeyringSnapControllerClient', () => {
         controller: controller as unknown as SnapController,
       });
 
+      controller.handleRequest.mockResolvedValue(accountsList);
       await client.listAccounts();
       expect(controller.handleRequest).toHaveBeenCalledWith({
         ...request,

src/keyring-snap-controller-client.ts

@@ -1,10 +1,13 @@
 import type { SnapController } from '@metamask/snaps-controllers';
 import type { HandlerType, ValidatedSnapId } from '@metamask/snaps-utils';
-import type { Json } from '@metamask/utils';
-import { assert } from 'superstruct';
 
 import { KeyringClient, Sender } from './keyring-client';
-import { InternalRequest, InternalRequestStruct } from './keyring-internal-api';
+import {
+  InternalRequest,
+  InternalResponse,
+  InternalResponseStruct,
+} from './keyring-internal-api';
+import { strictMask } from './utils';
 
 /**
  * Implementation of the `Sender` interface that can be used to send requests
@@ -45,16 +48,16 @@ class SnapControllerSender implements Sender {
    * @param request - JSON-RPC request to send to the snap.
    * @returns A promise that resolves to the response of the request.
    */
-  async send<Response extends Json>(
-    request: InternalRequest,
-  ): Promise<Response> {
-    assert(request, InternalRequestStruct);
-    return (await this.#controller.handleRequest({
-      snapId: this.#snapId as ValidatedSnapId,
-      origin: this.#origin,
-      handler: this.#handler,
-      request,
-    })) as Response;
+  async send(request: InternalRequest): Promise<InternalResponse> {
+    return strictMask(
+      await this.#controller.handleRequest({
+        snapId: this.#snapId as ValidatedSnapId,
+        origin: this.#origin,
+        handler: this.#handler,
+        request,
+      }),
+      InternalResponseStruct,
+    );
   }
 }
 

src/keyring-snap-rpc-client.ts

@@ -1,8 +1,12 @@
 import type { MetaMaskInpageProvider } from '@metamask/providers';
-import type { Json } from '@metamask/utils';
 
 import { KeyringClient, Sender } from './keyring-client';
-import type { InternalRequest } from './keyring-internal-api';
+import {
+  InternalResponseStruct,
+  type InternalRequest,
+  InternalResponse,
+} from './keyring-internal-api';
+import { strictMask } from './utils';
 
 /**
  * Implementation of the `Sender` interface that can be used to send requests
@@ -30,17 +34,17 @@ export class SnapRpcSender implements Sender {
    * @param request - The JSON-RPC request to send to the snap.
    * @returns A promise that resolves to the response of the request.
    */
-  async send<Response extends Json>(
-    request: InternalRequest,
-  ): Promise<Response> {
-    const response = await this.#provider.request({
-      method: 'wallet_invokeSnap',
-      params: {
-        snapId: this.#origin,
-        request,
-      },
-    });
-    return response as Response;
+  async send(request: InternalRequest): Promise<InternalResponse> {
+    return strictMask(
+      await this.#provider.request({
+        method: 'wallet_invokeSnap',
+        params: {
+          snapId: this.#origin,
+          request,
+        },
+      }),
+      InternalResponseStruct,
+    );
   }
 }
 

src/utils.ts

@@ -1,4 +1,4 @@
-import { pattern, string } from 'superstruct';
+import { Struct, assert, pattern, string } from 'superstruct';
 
 /**
  * UUIDv4 struct.
@@ -19,3 +19,23 @@ export const UuidStruct = pattern(
 export type OmitUnion<T, K extends keyof any> = T extends any
   ? Omit<T, K>
   : never;
+
+/**
+ * Assert that a value is valid according to a struct.
+ *
+ * It is similar to superstruct's mask function, but it does not ignore extra
+ * properties.
+ *
+ * @param value - Value to check.
+ * @param struct - Struct to validate the value against.
+ * @param message - Error message to throw if the value is not valid.
+ * @returns The value if it is valid.
+ */
+export function strictMask<T, S>(
+  value: unknown,
+  struct: Struct<T, S>,
+  message?: string,
+): T {
+  assert(value, struct, message);
+  return value;
+}