Address PR comments

MetaMask · Jan 9, 2024 · 33a93ee · 33a93ee
1 parent 8e597af
commit 33a93ee
Show file tree

Hide file tree

Showing 2 changed files with 72 additions and 61 deletions.
diff --git a/packages/snaps-controllers/coverage.json b/packages/snaps-controllers/coverage.json
@@ -1,6 +1,6 @@
 {
-  "branches": 91.01,
+  "branches": 90.9,
   "functions": 96.41,
-  "lines": 97.73,
-  "statements": 97.41
+  "lines": 97.72,
+  "statements": 97.4
 }
diff --git a/packages/snaps-controllers/src/snaps/SnapController.ts b/packages/snaps-controllers/src/snaps/SnapController.ts
@@ -244,9 +244,9 @@ export type PersistedSnapControllerState = SnapControllerState & {
 type RollbackSnapshot = {
   statePatches: Patch[];
   permissions: {
-    revoked: unknown;
-    granted: unknown[];
-    requestData: unknown;
+    revoked?: SubjectPermissions<ValidPermission<string, Caveat<string, any>>>;
+    granted?: RequestedPermissions;
+    requestData?: Record<string, unknown>;
   };
   newVersion: string;
 };
@@ -1025,16 +1025,16 @@ export class SnapController extends BaseController<
         ? virtualFiles.find((file) => file.path === iconPath)
         : undefined;
 
-      assert(sourceCode, 'Source code not provided for preinstalled snap');
+      assert(sourceCode, 'Source code not provided for preinstalled snap.');
 
       assert(
         manifest.source.files === undefined,
-        'Auxiliary files are not currently supported for preinstalled snaps',
+        'Auxiliary files are not currently supported for preinstalled snaps.',
       );
 
       assert(
         manifest.source.locales === undefined,
-        'Localization files are not currently supported for preinstalled snaps',
+        'Localization files are not currently supported for preinstalled snaps.',
       );
 
       const filesObject: FetchedSnapFiles = {
@@ -1048,7 +1048,7 @@ export class SnapController extends BaseController<
       // Add snap to the SnapController state
       this.#set({
         id: snapId,
-        origin: 'MetaMask',
+        origin: 'metamask',
         files: filesObject,
         removable,
         preinstalled: true,
@@ -1064,19 +1064,7 @@ export class SnapController extends BaseController<
       const { newPermissions, unusedPermissions } =
         this.#calculatePermissionsChange(snapId, processedPermissions);
 
-      const unusedPermissionsKeys = Object.keys(unusedPermissions);
-      if (isNonEmptyArray(unusedPermissionsKeys)) {
-        this.messagingSystem.call('PermissionController:revokePermissions', {
-          [snapId]: unusedPermissionsKeys,
-        });
-      }
-
-      if (isNonEmptyArray(Object.keys(newPermissions))) {
-        this.messagingSystem.call('PermissionController:grantPermissions', {
-          approvedPermissions: newPermissions,
-          subject: { origin: snapId },
-        });
-      }
+      this.#updatePermissions({ snapId, newPermissions, unusedPermissions });
 
       // Set status
       this.update((state) => {
@@ -2128,27 +2116,17 @@ export class SnapController extends BaseController<
         isUpdate: true,
       });
 
-      const unusedPermissionsKeys = Object.keys(unusedPermissions);
-      if (isNonEmptyArray(unusedPermissionsKeys)) {
-        this.messagingSystem.call('PermissionController:revokePermissions', {
-          [snapId]: unusedPermissionsKeys,
-        });
-      }
-
-      if (isNonEmptyArray(Object.keys(approvedNewPermissions))) {
-        this.messagingSystem.call('PermissionController:grantPermissions', {
-          approvedPermissions: approvedNewPermissions,
-          subject: { origin: snapId },
-          requestData,
-        });
-      }
+      this.#updatePermissions({
+        snapId,
+        unusedPermissions,
+        newPermissions: approvedNewPermissions,
+        requestData,
+      });
 
       const rollbackSnapshot = this.#getRollbackSnapshot(snapId);
       if (rollbackSnapshot !== undefined) {
         rollbackSnapshot.permissions.revoked = unusedPermissions;
-        rollbackSnapshot.permissions.granted = Object.keys(
-          approvedNewPermissions,
-        );
+        rollbackSnapshot.permissions.granted = approvedNewPermissions;
         rollbackSnapshot.permissions.requestData = requestData;
       }
 
@@ -2593,13 +2571,11 @@ export class SnapController extends BaseController<
       const { permissions: approvedPermissions, ...requestData } =
         (await pendingApproval.promise) as PermissionsRequest;
 
-      if (isNonEmptyArray(Object.keys(approvedPermissions))) {
-        this.messagingSystem.call('PermissionController:grantPermissions', {
-          approvedPermissions,
-          subject: { origin: snapId },
-          requestData,
-        });
-      }
+      this.#updatePermissions({
+        snapId,
+        newPermissions: approvedPermissions,
+        requestData,
+      });
     } finally {
       const runtime = this.#getRuntimeExpect(snapId);
       runtime.installPromise = null;
@@ -2922,7 +2898,7 @@ export class SnapController extends BaseController<
 
     this.#rollbackSnapshots.set(snapId, {
       statePatches: [],
-      permissions: { revoked: null, granted: [], requestData: null },
+      permissions: {},
       newVersion: '',
     });
 
@@ -2972,19 +2948,12 @@ export class SnapController extends BaseController<
       });
     }
 
-    if (permissions.revoked && Object.keys(permissions.revoked).length) {
-      this.messagingSystem.call('PermissionController:grantPermissions', {
-        approvedPermissions: permissions.revoked as RequestedPermissions,
-        subject: { origin: snapId },
-        requestData: permissions.requestData as Record<string, unknown>,
-      });
-    }
-
-    if (permissions.granted?.length) {
-      this.messagingSystem.call('PermissionController:revokePermissions', {
-        [snapId]: permissions.granted as NonEmptyArray<string>,
-      });
-    }
+    this.#updatePermissions({
+      snapId,
+      unusedPermissions: permissions.granted,
+      newPermissions: permissions.revoked,
+      requestData: permissions.requestData,
+    });
 
     const truncatedSnap = this.getTruncatedExpect(snapId);
 
@@ -3082,6 +3051,48 @@ export class SnapController extends BaseController<
     return { newPermissions, unusedPermissions, approvedPermissions };
   }
 
+  /**
+   * Updates the permissions for a snap following an install, update or rollback.
+   *
+   * Grants newly requested permissions and revokes unused/revoked permissions.
+   *
+   * @param args - An options bag.
+   * @param args.snapId - The snap ID.
+   * @param args.newPermissions - New permissions to be granted.
+   * @param args.unusedPermissions - Unused permissions to be revoked.
+   * @param args.requestData - Optional request data from an approval.
+   */
+  #updatePermissions({
+    snapId,
+    unusedPermissions = {},
+    newPermissions = {},
+    requestData,
+  }: {
+    snapId: SnapId;
+    newPermissions?:
+      | RequestedPermissions
+      | Record<string, Pick<PermissionConstraint, 'caveats'>>;
+    unusedPermissions?:
+      | RequestedPermissions
+      | SubjectPermissions<ValidPermission<string, Caveat<string, any>>>;
+    requestData?: Record<string, unknown>;
+  }) {
+    const unusedPermissionsKeys = Object.keys(unusedPermissions);
+    if (isNonEmptyArray(unusedPermissionsKeys)) {
+      this.messagingSystem.call('PermissionController:revokePermissions', {
+        [snapId]: unusedPermissionsKeys,
+      });
+    }
+
+    if (isNonEmptyArray(Object.keys(newPermissions))) {
+      this.messagingSystem.call('PermissionController:grantPermissions', {
+        approvedPermissions: newPermissions,
+        subject: { origin: snapId },
+        requestData,
+      });
+    }
+  }
+
   /**
    * Checks if a snap will pass version validation checks
    * with the new version range that is requested. The first