diff --git a/packages/snaps-controllers/src/snaps/SnapController.test.tsx b/packages/snaps-controllers/src/snaps/SnapController.test.tsx index e75b5cf2f0..1b55c086c9 100644 --- a/packages/snaps-controllers/src/snaps/SnapController.test.tsx +++ b/packages/snaps-controllers/src/snaps/SnapController.test.tsx @@ -8092,7 +8092,17 @@ describe('SnapController', () => { }) ).manifest.result; - const manifest2 = ( + const manifest2Old = ( + await getMockSnapFilesWithUpdatedChecksum({ + manifest: getSnapManifest({ + initialConnections: { + [MOCK_ORIGIN]: {}, + }, + }), + }) + ).manifest.result; + + const manifest2New = ( await getMockSnapFilesWithUpdatedChecksum({ manifest: getSnapManifest({ version: newVersion, @@ -8105,7 +8115,9 @@ describe('SnapController', () => { const detect = jest .fn() .mockImplementationOnce(() => new LoopbackLocation()) - .mockImplementationOnce(() => new LoopbackLocation()) + .mockImplementationOnce( + () => new LoopbackLocation({ manifest: manifest2Old }), + ) .mockImplementationOnce(() => new LoopbackLocation()) .mockImplementationOnce( () => @@ -8116,7 +8128,7 @@ describe('SnapController', () => { .mockImplementationOnce( () => new LoopbackLocation({ - manifest: manifest2, + manifest: manifest2New, files: [ new VirtualFile({ value: 'foo', @@ -8130,8 +8142,46 @@ describe('SnapController', () => { }), ); + const rootMessenger = getControllerMessenger(); + + let revokedConnection = false; + + rootMessenger.registerActionHandler( + 'PermissionController:revokePermissions', + () => { + revokedConnection = true; + return {}; + }, + ); + + rootMessenger.registerActionHandler( + 'PermissionController:getPermissions', + (origin) => { + if (origin === MOCK_ORIGIN && !revokedConnection) { + return { + [WALLET_SNAP_PERMISSION_KEY]: { + caveats: [ + { + type: SnapCaveatType.SnapIds, + value: { + [snapId2]: {}, + }, + }, + ], + date: 1664187844588, + id: 'izn0WGUO8cvq_jqvLQuQP', + invoker: MOCK_ORIGIN, + parentCapability: WALLET_SNAP_PERMISSION_KEY, + }, + }; + } + return MOCK_SNAP_PERMISSIONS; + }, + ); + const options = getSnapControllerWithEESOptions({ detectSnapLocation: detect, + rootMessenger, }); const { messenger } = options; @@ -8176,6 +8226,34 @@ describe('SnapController', () => { expect.anything(), ); + expect(messenger.call).toHaveBeenNthCalledWith( + 48, + 'PermissionController:revokePermissions', + { + [MOCK_ORIGIN]: [WALLET_SNAP_PERMISSION_KEY], + }, + ); + + expect(messenger.call).toHaveBeenNthCalledWith( + 59, + 'PermissionController:grantPermissions', + { + approvedPermissions: { + [WALLET_SNAP_PERMISSION_KEY]: { + caveats: [ + { + type: SnapCaveatType.SnapIds, + value: { + [snapId2]: {}, + }, + }, + ], + }, + }, + subject: { origin: MOCK_ORIGIN }, + }, + ); + controller.destroy(); await service.terminateAllSnaps(); }); diff --git a/packages/snaps-controllers/src/snaps/SnapController.ts b/packages/snaps-controllers/src/snaps/SnapController.ts index f490c70793..e9c32477d5 100644 --- a/packages/snaps-controllers/src/snaps/SnapController.ts +++ b/packages/snaps-controllers/src/snaps/SnapController.ts @@ -58,6 +58,7 @@ import type { OnAssetsMarketDataResponse, AssetMarketData, AssetMetadata, + EmptyObject, } from '@metamask/snaps-sdk'; import { AuxiliaryFileEncoding, @@ -333,6 +334,8 @@ type RollbackSnapshot = { granted?: RequestedPermissions; requestData?: Record; }; + previousInitialConnections?: Record | null; + newInitialConnections?: Record; newVersion: string; }; @@ -2994,19 +2997,22 @@ export class SnapController extends BaseController< requestData, }); - if (manifest.initialConnections) { - this.#handleInitialConnections( - snapId, - oldManifest.initialConnections ?? null, - manifest.initialConnections, - ); - } + const previousInitialConnections = oldManifest.initialConnections ?? null; + const newInitialConnections = manifest.initialConnections ?? {}; + this.#handleInitialConnections( + snapId, + previousInitialConnections, + newInitialConnections, + ); const rollbackSnapshot = this.#getRollbackSnapshot(snapId); if (rollbackSnapshot !== undefined) { rollbackSnapshot.permissions.revoked = unusedPermissions; rollbackSnapshot.permissions.granted = approvedNewPermissions; rollbackSnapshot.permissions.requestData = requestData; + rollbackSnapshot.previousInitialConnections = + previousInitialConnections; + rollbackSnapshot.newInitialConnections = newInitialConnections; } const sourceCode = sourceCodeFile.toString(); @@ -4139,7 +4145,12 @@ export class SnapController extends BaseController< this.#transition(snapId, SnapStatusEvents.Stop); } - const { statePatches, permissions } = rollbackSnapshot; + const { + statePatches, + permissions, + previousInitialConnections, + newInitialConnections, + } = rollbackSnapshot; if (statePatches?.length) { this.applyPatches(statePatches); @@ -4160,6 +4171,13 @@ export class SnapController extends BaseController< requestData: permissions.requestData, }); + // Calling this in reverse order to undo the changes + this.#handleInitialConnections( + snapId, + newInitialConnections ?? null, + previousInitialConnections ?? {}, + ); + const truncatedSnap = this.getTruncatedExpect(snapId); this.messagingSystem.publish(