Skip to content

Commit

Permalink
BREAKING: Refactor encryption to enable caching (#2316)
Browse files Browse the repository at this point in the history 
This PR is a refactor of the Snap state encryption logic, it moves the
logic from the manageState RPC method implementation to the
SnapController. It re-implements the encryption/decryption logic using a
different set of methods, as the previously used implementation had
encapsulated the encryption key. With this refactor it is now possible
to cache the encryption key used by the Snap in the runtime, making
continuous requests to update the state much faster (numbers TBD).

With this refactor it is now the responsibility of the SnapController to
manage encrypting/decrypting the Snap state and thus the method hooks
`getSnapState` and `updateSnapState` have had breaking changes, they are
now asynchronous and expected to fully handle encryption/decryption.
Because of that, this change required changes to both snaps-jest and
snaps-simulator for compatibility.

This will not work as-is on mobile, since the encryptor on mobile does
not have feature parity with `browser-passworder`. This will need to be
addressed before we can use this change on mobile.

Closes #2315

---------

Co-authored-by: Maarten Zuidhoorn <maarten@zuidhoorn.com>
  • Loading branch information
@FrederikBolding @Mrtenz
FrederikBolding and Mrtenz authored Apr 10, 2024
1 parent 6f9ca36 commit f6ac2b0
Show file tree
Hide file tree
Showing 32 changed files with 657 additions and 545 deletions.
8 changes: 4 additions & 4 deletions packages/snaps-controllers/coverage.json
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
{
"branches": 90.82,
"functions": 96.57,
"lines": 97.8,
"statements": 97.46
"branches": 90.9,
"functions": 96.58,
"lines": 97.85,
"statements": 97.52
}
1 change: 1 addition & 0 deletions packages/snaps-controllers/package.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,7 @@
"@esbuild-plugins/node-modules-polyfill": "^0.2.2",
"@lavamoat/allow-scripts": "^3.0.3",
"@metamask/auto-changelog": "^3.4.4",
"@metamask/browser-passworder": "^5.0.0",
"@metamask/eslint-config": "^12.1.0",
"@metamask/eslint-config-jest": "^12.1.0",
"@metamask/eslint-config-nodejs": "^12.1.0",
Expand Down
251 changes: 208 additions & 43 deletions packages/snaps-controllers/src/snaps/SnapController.test.ts
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import { getPersistentState } from '@metamask/base-controller';
import { encrypt } from '@metamask/browser-passworder';
import {
createAsyncMiddleware,
JsonRpcEngine,
Expand Down Expand Up @@ -57,6 +58,7 @@ import {
stringToBytes,
} from '@metamask/utils';
import { File } from 'buffer';
import { webcrypto } from 'crypto';
import fetchMock from 'jest-fetch-mock';
import { pipeline } from 'readable-stream';
import type { Duplex } from 'readable-stream';
Expand All @@ -65,6 +67,7 @@ import { setupMultiplex } from '../services';
import type { NodeThreadExecutionService } from '../services/node';
import {
approvalControllerMock,
DEFAULT_ENCRYPTION_KEY_DERIVATION_OPTIONS,
ExecutionEnvironmentStub,
getControllerMessenger,
getNodeEESMessenger,
Expand Down Expand Up @@ -99,16 +102,34 @@ import {
SNAP_APPROVAL_UPDATE,
} from './SnapController';

Object.defineProperty(globalThis, 'crypto', {
value: {
// eslint-disable-next-line @typescript-eslint/no-require-imports, @typescript-eslint/no-var-requires
...require('node:crypto').webcrypto,
getRandomValues: jest.fn().mockReturnValue(new Uint32Array(32)),
},
});
if (!('CryptoKey' in globalThis)) {
// We can remove this once we drop Node 18
Object.defineProperty(globalThis, 'CryptoKey', {
value: webcrypto.CryptoKey,
});
}

globalThis.crypto ??= webcrypto as typeof globalThis.crypto;
globalThis.crypto.getRandomValues = <Type extends ArrayBufferView | null>(
array: Type,
) => {
if (array === null) {
return null as Type;
}

return new Uint8Array(array.buffer).fill(0) as unknown as Type;
};

fetchMock.enableMocks();

// Encryption key for `MOCK_SNAP_ID`.
const ENCRYPTION_KEY =
'0xd2f0a8e994b871ba4451ac383bf323cdaad8d554736355f2223e155692fbc446';

// Encryption key for `MOCK_LOCAL_SNAP_ID`.
const OTHER_ENCRYPTION_KEY =
'0x7cd340349a41e0f7af62a9d97c76e96b12485e0206791d6b5638dd59736af8f5';

describe('SnapController', () => {
beforeEach(() => {
// eslint-disable-next-line @typescript-eslint/require-await
Expand All @@ -125,28 +146,6 @@ describe('SnapController', () => {
await service.terminateAllSnaps();
});

it('creates a worker and snap controller, adds a snap, and update its state', async () => {
const [snapController, service] = getSnapControllerWithEES(
getSnapControllerWithEESOptions({
state: {
snaps: getPersistedSnapsState(),
},
}),
);

const snap = snapController.getExpect(MOCK_SNAP_ID);
const state = 'foo';

await snapController.startSnap(snap.id);
snapController.updateSnapState(snap.id, state, true);
const snapState = snapController.getSnapState(snap.id, true);
expect(snapState).toStrictEqual(state);

expect(snapController.state.snapStates[MOCK_SNAP_ID]).toStrictEqual(state);
snapController.destroy();
await service.terminateAllSnaps();
});

it('adds a snap and uses its JSON-RPC api with a NodeThreadExecutionService', async () => {
const [snapController, service] = getSnapControllerWithEES(
getSnapControllerWithEESOptions({
Expand Down Expand Up @@ -7427,7 +7426,15 @@ describe('SnapController', () => {
it(`gets the snap's state`, async () => {
const messenger = getSnapControllerMessenger();

const state = 'foo';
const state = { myVariable: 1 };

const mockEncryptedState = await encrypt(
ENCRYPTION_KEY,
state,
undefined,
undefined,
DEFAULT_ENCRYPTION_KEY_DERIVATION_OPTIONS,
);

const snapController = getSnapController(
getSnapControllerOptions({
Expand All @@ -7437,14 +7444,14 @@ describe('SnapController', () => {
[MOCK_SNAP_ID]: getPersistedSnapObject(),
},
snapStates: {
[MOCK_SNAP_ID]: state,
[MOCK_SNAP_ID]: mockEncryptedState,
},
},
}),
);

const getSnapStateSpy = jest.spyOn(snapController, 'getSnapState');
const result = messenger.call(
const result = await messenger.call(
'SnapController:getSnapState',
MOCK_SNAP_ID,
true,
Expand All @@ -7456,10 +7463,161 @@ describe('SnapController', () => {
snapController.destroy();
});

it('migrates user storage to latest key derivation options', async () => {
const messenger = getSnapControllerMessenger();

const state = { myVariable: 1 };

const initialEncryptedState = await encrypt(
ENCRYPTION_KEY,
state,
undefined,
undefined,
{
...DEFAULT_ENCRYPTION_KEY_DERIVATION_OPTIONS,
params: { iterations: 10_000 },
},
);

const snapController = getSnapController(
getSnapControllerOptions({
messenger,
state: {
snaps: {
[MOCK_SNAP_ID]: getPersistedSnapObject(),
},
snapStates: {
[MOCK_SNAP_ID]: initialEncryptedState,
},
},
}),
);

const newState = { myVariable: 2 };

await messenger.call(
'SnapController:updateSnapState',
MOCK_SNAP_ID,
newState,
true,
);

const upgradedEncryptedState = await encrypt(
ENCRYPTION_KEY,
newState,
undefined,
undefined,
DEFAULT_ENCRYPTION_KEY_DERIVATION_OPTIONS,
);

const result = await messenger.call(
'SnapController:getSnapState',
MOCK_SNAP_ID,
true,
);

expect(result).toStrictEqual(newState);
expect(snapController.state.snapStates[MOCK_SNAP_ID]).toStrictEqual(
upgradedEncryptedState,
);

snapController.destroy();
});

it('different snaps use different encryption keys', async () => {
const messenger = getSnapControllerMessenger();

const state = { foo: 'bar' };

const snapController = getSnapController(
getSnapControllerOptions({
messenger,
state: {
snaps: {
[MOCK_SNAP_ID]: getPersistedSnapObject(),
[MOCK_LOCAL_SNAP_ID]: getPersistedSnapObject({
id: MOCK_LOCAL_SNAP_ID,
}),
},
},
}),
);

await messenger.call(
'SnapController:updateSnapState',
MOCK_SNAP_ID,
state,
true,
);

await messenger.call(
'SnapController:updateSnapState',
MOCK_LOCAL_SNAP_ID,
state,
true,
);

const encryptedState1 = await encrypt(
ENCRYPTION_KEY,
state,
undefined,
undefined,
DEFAULT_ENCRYPTION_KEY_DERIVATION_OPTIONS,
);

const encryptedState2 = await encrypt(
OTHER_ENCRYPTION_KEY,
state,
undefined,
undefined,
DEFAULT_ENCRYPTION_KEY_DERIVATION_OPTIONS,
);

expect(snapController.state.snapStates[MOCK_SNAP_ID]).toStrictEqual(
encryptedState1,
);
expect(snapController.state.snapStates[MOCK_LOCAL_SNAP_ID]).toStrictEqual(
encryptedState2,
);
expect(snapController.state.snapStates[MOCK_SNAP_ID]).not.toStrictEqual(
snapController.state.snapStates[MOCK_LOCAL_SNAP_ID],
);

snapController.destroy();
});

it('throws an error if the state is corrupt', async () => {
const messenger = getSnapControllerMessenger();

const snapController = getSnapController(
getSnapControllerOptions({
messenger,
state: {
snaps: {
[MOCK_SNAP_ID]: getPersistedSnapObject(),
},
snapStates: {
[MOCK_SNAP_ID]: 'foo',
},
},
}),
);

await expect(
messenger.call('SnapController:getSnapState', MOCK_SNAP_ID, true),
).rejects.toThrow(
rpcErrors.internal({
message: 'Failed to decrypt snap state, the state must be corrupted.',
}),
);

snapController.destroy();
});

it(`gets the snap's unencrypted state`, async () => {
const messenger = getSnapControllerMessenger();

const state = 'foo';
const state = { foo: 'bar' };

const snapController = getSnapController(
getSnapControllerOptions({
Expand All @@ -7469,14 +7627,14 @@ describe('SnapController', () => {
[MOCK_SNAP_ID]: getPersistedSnapObject(),
},
unencryptedSnapStates: {
[MOCK_SNAP_ID]: state,
[MOCK_SNAP_ID]: JSON.stringify(state),
},
},
}),
);

const getSnapStateSpy = jest.spyOn(snapController, 'getSnapState');
const result = messenger.call(
const result = await messenger.call(
'SnapController:getSnapState',
MOCK_SNAP_ID,
false,
Expand Down Expand Up @@ -7535,8 +7693,15 @@ describe('SnapController', () => {
);

const updateSnapStateSpy = jest.spyOn(snapController, 'updateSnapState');
const state = 'bar';
messenger.call(
const state = { foo: 'bar' };
const mockEncryptedState = await encrypt(
ENCRYPTION_KEY,
state,
undefined,
undefined,
DEFAULT_ENCRYPTION_KEY_DERIVATION_OPTIONS,
);
await messenger.call(
'SnapController:updateSnapState',
MOCK_SNAP_ID,
state,
Expand All @@ -7545,7 +7710,7 @@ describe('SnapController', () => {

expect(updateSnapStateSpy).toHaveBeenCalledTimes(1);
expect(snapController.state.snapStates[MOCK_SNAP_ID]).toStrictEqual(
state,
mockEncryptedState,
);

snapController.destroy();
Expand All @@ -7564,8 +7729,8 @@ describe('SnapController', () => {
);

const updateSnapStateSpy = jest.spyOn(snapController, 'updateSnapState');
const state = 'bar';
messenger.call(
const state = { foo: 'bar' };
await messenger.call(
'SnapController:updateSnapState',
MOCK_SNAP_ID,
state,
Expand All @@ -7575,7 +7740,7 @@ describe('SnapController', () => {
expect(updateSnapStateSpy).toHaveBeenCalledTimes(1);
expect(
snapController.state.unencryptedSnapStates[MOCK_SNAP_ID],
).toStrictEqual(state);
).toStrictEqual(JSON.stringify(state));

snapController.destroy();
});
Expand All @@ -7600,7 +7765,7 @@ describe('SnapController', () => {
);

messenger.call('SnapController:clearSnapState', MOCK_SNAP_ID, true);
const clearedState = messenger.call(
const clearedState = await messenger.call(
'SnapController:getSnapState',
MOCK_SNAP_ID,
true,
Expand Down Expand Up @@ -7628,7 +7793,7 @@ describe('SnapController', () => {
);

messenger.call('SnapController:clearSnapState', MOCK_SNAP_ID, false);
const clearedState = messenger.call(
const clearedState = await messenger.call(
'SnapController:getSnapState',
MOCK_SNAP_ID,
false,
Expand Down
Loading

0 comments on commit f6ac2b0

Please sign in to comment.