Refactor protocol caveats

Author: FrederikBolding

packages/snaps-controllers/src/multichain/MultichainRoutingController.ts

@@ -4,15 +4,10 @@ import type {
   ControllerStateChangeEvent,
 } from '@metamask/base-controller';
 import { BaseController } from '@metamask/base-controller';
-import type {
-  Caveat,
-  GetPermissions,
-  ValidPermission,
-} from '@metamask/permission-controller';
+import type { GetPermissions } from '@metamask/permission-controller';
 import { rpcErrors } from '@metamask/rpc-errors';
 import {
-  getProtocolCaveatChainIds,
-  getProtocolCaveatRpcMethods,
+  getProtocolCaveatChains,
   SnapEndowments,
 } from '@metamask/snaps-rpc-methods';
 import type {
@@ -103,9 +98,9 @@ export type MultichainRoutingControllerArgs = {
 
 export type MultichainRoutingControllerState = EmptyObject;
 
-type SnapWithPermission = {
+type ProtocolSnap = {
   snapId: SnapId;
-  permission: ValidPermission<string, Caveat<string, Json>>;
+  methods: string[];
 };
 
 const controllerName = 'MultichainRoutingController';
@@ -223,18 +218,18 @@ export class MultichainRoutingController extends BaseController<
     const allSnaps = this.messagingSystem.call('SnapController:getAll');
     const filteredSnaps = getRunnableSnaps(allSnaps);
 
-    return filteredSnaps.reduce<SnapWithPermission[]>((accumulator, snap) => {
+    return filteredSnaps.reduce<ProtocolSnap[]>((accumulator, snap) => {
       const permissions = this.messagingSystem.call(
         'PermissionController:getPermissions',
         snap.id,
       );
       if (permissions && hasProperty(permissions, SnapEndowments.Protocol)) {
         const permission = permissions[SnapEndowments.Protocol];
-        const chains = getProtocolCaveatChainIds(permission);
-        if (chains?.includes(chainId)) {
+        const chains = getProtocolCaveatChains(permission);
+        if (chains && hasProperty(chains, chainId)) {
           accumulator.push({
             snapId: snap.id,
-            permission,
+            methods: chains[chainId].methods,
           });
         }
       }
@@ -276,7 +271,7 @@ export class MultichainRoutingController extends BaseController<
     // but has a protocol Snap available, route it there.
     const protocolSnaps = this.#getProtocolSnaps(chainId);
     const protocolSnap = protocolSnaps.find((snap) =>
-      getProtocolCaveatRpcMethods(snap.permission)?.includes(method),
+      snap.methods.includes(method),
     );
     if (protocolSnap) {
       return this.messagingSystem.call('SnapController:handleRequest', {

packages/snaps-controllers/src/test-utils/multichain.ts

@@ -65,12 +65,8 @@ export const MOCK_SOLANA_SNAP_PERMISSIONS: Record<
   [SnapEndowments.Keyring]: {
     caveats: [
       {
-        type: SnapCaveatType.ChainIds,
-        value: [SOLANA_CAIP2],
-      },
-      {
-        type: SnapCaveatType.SnapRpcMethods,
-        value: ['getVersion'],
+        type: SnapCaveatType.KeyringOrigin,
+        value: { allowedOrigins: [] },
       },
     ],
     date: 1664187844588,
@@ -81,12 +77,8 @@ export const MOCK_SOLANA_SNAP_PERMISSIONS: Record<
   [SnapEndowments.Protocol]: {
     caveats: [
       {
-        type: SnapCaveatType.ChainIds,
-        value: [SOLANA_CAIP2],
-      },
-      {
-        type: SnapCaveatType.SnapRpcMethods,
-        value: ['getVersion'],
+        type: SnapCaveatType.ProtocolSnapChains,
+        value: { [SOLANA_CAIP2]: { methods: ['getVersion'] } },
       },
     ],
     date: 1664187844588,

packages/snaps-rpc-methods/src/endowments/index.ts

@@ -128,7 +128,4 @@ export { getChainIdsCaveat, getLookupMatchersCaveat } from './name-lookup';
 export { getKeyringCaveatOrigins } from './keyring';
 export { getMaxRequestTimeCaveat } from './caveats';
 export { getCronjobCaveatJobs } from './cronjob';
-export {
-  getProtocolCaveatChainIds,
-  getProtocolCaveatRpcMethods,
-} from './protocol';
+export { getProtocolCaveatChains } from './protocol';

packages/snaps-rpc-methods/src/endowments/protocol.ts

@@ -10,10 +10,8 @@ import type {
 } from '@metamask/permission-controller';
 import { PermissionType, SubjectType } from '@metamask/permission-controller';
 import { rpcErrors } from '@metamask/rpc-errors';
-import {
-  ProtocolRpcMethodsStruct,
-  SnapCaveatType,
-} from '@metamask/snaps-utils';
+import { ProtocolChainsStruct, SnapCaveatType } from '@metamask/snaps-utils';
+import type { Infer } from '@metamask/superstruct';
 import type { Json, NonEmptyArray } from '@metamask/utils';
 import {
   assertStruct,
@@ -51,15 +49,10 @@ const specificationBuilder: PermissionSpecificationBuilder<
   return {
     permissionType: PermissionType.Endowment,
     targetName: permissionName,
-    allowedCaveats: [
-      SnapCaveatType.ChainIds,
-      SnapCaveatType.SnapRpcMethods,
-      SnapCaveatType.MaxRequestTime,
-    ],
+    allowedCaveats: [SnapCaveatType.ChainIds, SnapCaveatType.MaxRequestTime],
     endowmentGetter: (_getterOptions?: EndowmentGetterParams) => null,
     validator: createGenericPermissionValidator([
-      { type: SnapCaveatType.ChainIds },
-      { type: SnapCaveatType.SnapRpcMethods },
+      { type: SnapCaveatType.ProtocolSnapChains },
       { type: SnapCaveatType.MaxRequestTime, optional: true },
     ]),
     subjectTypes: [SubjectType.Snap],
@@ -90,51 +83,29 @@ export function getProtocolCaveatMapper(
 
   if (value.chains) {
     caveats.push({
-      type: SnapCaveatType.ChainIds,
+      type: SnapCaveatType.ProtocolSnapChains,
       value: value.chains,
     });
   }
 
-  if (value.methods) {
-    caveats.push({
-      type: SnapCaveatType.SnapRpcMethods,
-      value: value.methods,
-    });
-  }
-
   return { caveats: caveats as NonEmptyArray<CaveatConstraint> };
 }
 
-/**
- * Getter function to get the {@link ChainIds} caveat value from a
- * permission.
- *
- * @param permission - The permission to get the caveat value from.
- * @returns The caveat value.
- */
-export function getProtocolCaveatChainIds(
-  permission?: PermissionConstraint,
-): string[] | null {
-  const caveat = permission?.caveats?.find(
-    (permCaveat) => permCaveat.type === SnapCaveatType.ChainIds,
-  ) as Caveat<string, string[]> | undefined;
-
-  return caveat ? caveat.value : null;
-}
+export type ProtocolChains = Infer<typeof ProtocolChainsStruct>;
 
 /**
- * Getter function to get the {@link SnapRpcMethods} caveat value from a
+ * Getter function to get the {@link ProtocolSnapChains} caveat value from a
  * permission.
  *
  * @param permission - The permission to get the caveat value from.
  * @returns The caveat value.
  */
-export function getProtocolCaveatRpcMethods(
+export function getProtocolCaveatChains(
   permission?: PermissionConstraint,
-): string[] | null {
+): ProtocolChains | null {
   const caveat = permission?.caveats?.find(
-    (permCaveat) => permCaveat.type === SnapCaveatType.SnapRpcMethods,
-  ) as Caveat<string, string[]> | undefined;
+    (permCaveat) => permCaveat.type === SnapCaveatType.ProtocolSnapChains,
+  ) as Caveat<string, ProtocolChains> | undefined;
 
   return caveat ? caveat.value : null;
 }
@@ -155,18 +126,18 @@ function validateCaveat(caveat: Caveat<string, any>): void {
   const { value } = caveat;
   assertStruct(
     value,
-    ProtocolRpcMethodsStruct,
-    'Invalid RPC methods specified',
+    ProtocolChainsStruct,
+    'Invalid chains specified',
     rpcErrors.invalidParams,
   );
 }
 
 export const protocolCaveatSpecifications: Record<
-  SnapCaveatType.SnapRpcMethods,
+  SnapCaveatType.ProtocolSnapChains,
   CaveatSpecificationConstraint
 > = {
-  [SnapCaveatType.SnapRpcMethods]: Object.freeze({
-    type: SnapCaveatType.SnapRpcMethods,
+  [SnapCaveatType.ProtocolSnapChains]: Object.freeze({
+    type: SnapCaveatType.ProtocolSnapChains,
     validator: (caveat: Caveat<string, any>) => validateCaveat(caveat),
   }),
 };

packages/snaps-utils/src/caveats.ts

@@ -57,5 +57,5 @@ export enum SnapCaveatType {
   /**
    * Caveat specifying a list of RPC methods serviced by an endowment.
    */
-  SnapRpcMethods = 'snapRpcMethods',
+  ProtocolSnapChains = 'protocolSnapChains',
 }

packages/snaps-utils/src/manifest/validation.ts

@@ -27,6 +27,7 @@ import {
   isValidSemVerRange,
   inMilliseconds,
   Duration,
+  CaipChainIdStruct,
 } from '@metamask/utils';
 
 import { isEqual } from '../array';
@@ -174,8 +175,10 @@ export const MaxRequestTimeStruct = size(
   MAXIMUM_REQUEST_TIMEOUT,
 );
 
-// TODO: Decide on the format for this
-export const ProtocolRpcMethodsStruct = array(string());
+export const ProtocolChainsStruct = record(
+  CaipChainIdStruct,
+  object({ methods: array(string()) }),
+);
 
 // Utility type to union with for all handler structs
 export const HandlerCaveatsStruct = object({
@@ -204,7 +207,7 @@ export const PermissionsStruct: Describe<InitialPermissions> = type({
   'endowment:protocol': optional(
     mergeStructs(
       HandlerCaveatsStruct,
-      object({ chains: ChainIdsStruct, methods: ProtocolRpcMethodsStruct }),
+      object({ chains: ProtocolChainsStruct }),
     ),
   ),
   'endowment:lifecycle-hooks': optional(HandlerCaveatsStruct),