Skip to content

Commit f776230

Browse files
MrtenzMrtenz
committed
Add support for state methods to snaps-simulation
1 parent 1e16931 commit f776230

File tree

11 files changed

+503
-126
lines changed

11 files changed

+503
-126
lines changed

‎packages/snaps-simulation/src/controllers.test.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,10 @@ import {
55
} from '@metamask/permission-controller';
66

77
import { getControllers } from './controllers';
8-
import type { MiddlewareHooks } from './simulation';
8+
import type { RestrictedMiddlewareHooks } from './simulation';
99
import { getMockOptions } from './test-utils';
1010

11-
const MOCK_HOOKS: MiddlewareHooks = {
11+
const MOCK_HOOKS: RestrictedMiddlewareHooks = {
1212
getIsLocked: jest.fn(),
1313
getMnemonic: jest.fn(),
1414
getSnapFile: jest.fn(),

‎packages/snaps-simulation/src/controllers.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ import { getSafeJson } from '@metamask/utils';
2929
import { getPermissionSpecifications } from './methods';
3030
import { UNRESTRICTED_METHODS } from './methods/constants';
3131
import type { SimulationOptions } from './options';
32-
import type { MiddlewareHooks } from './simulation';
32+
import type { RestrictedMiddlewareHooks } from './simulation';
3333
import type { RunSagaFunction } from './store';
3434

3535
export type RootControllerAllowedActions =
@@ -49,7 +49,7 @@ export type RootControllerMessenger = ControllerMessenger<
4949

5050
export type GetControllersOptions = {
5151
controllerMessenger: ControllerMessenger<any, any>;
52-
hooks: MiddlewareHooks;
52+
hooks: RestrictedMiddlewareHooks;
5353
runSaga: RunSagaFunction;
5454
options: SimulationOptions;
5555
};

‎packages/snaps-simulation/src/methods/hooks/index.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
export * from './get-preferences';
2+
export * from './interface';
23
export * from './notifications';
4+
export * from './permitted';
35
export * from './request-user-approval';
46
export * from './state';
5-
export * from './interface';

‎packages/snaps-simulation/src/methods/hooks/permitted/index.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
export * from './state';

‎packages/snaps-simulation/src/methods/hooks/permitted/state.test.ts‎

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
1+
import { createStore, getState, setState } from '../../../store';
2+
import { getMockOptions } from '../../../test-utils';
3+
import {
4+
getPermittedClearSnapStateMethodImplementation,
5+
getPermittedGetSnapStateMethodImplementation,
6+
getPermittedUpdateSnapStateMethodImplementation,
7+
} from './state';
8+
9+
describe('getPermittedGetSnapStateMethodImplementation', () => {
10+
it('returns the implementation of the `getSnapState` hook', async () => {
11+
const { store, runSaga } = createStore(getMockOptions());
12+
const fn = getPermittedGetSnapStateMethodImplementation(runSaga);
13+
14+
expect(await fn(true)).toBeNull();
15+
16+
store.dispatch(
17+
setState({
18+
state: JSON.stringify({
19+
foo: 'bar',
20+
}),
21+
encrypted: true,
22+
}),
23+
);
24+
25+
expect(await fn(true)).toStrictEqual({
26+
foo: 'bar',
27+
});
28+
});
29+
30+
it('returns the implementation of the `getSnapState` hook for unencrypted state', async () => {
31+
const { store, runSaga } = createStore(getMockOptions());
32+
const fn = getPermittedGetSnapStateMethodImplementation(runSaga);
33+
34+
expect(await fn(false)).toBeNull();
35+
36+
store.dispatch(
37+
setState({
38+
state: JSON.stringify({
39+
foo: 'bar',
40+
}),
41+
encrypted: false,
42+
}),
43+
);
44+
45+
expect(await fn(false)).toStrictEqual({
46+
foo: 'bar',
47+
});
48+
});
49+
});
50+
51+
describe('getPermittedUpdateSnapStateMethodImplementation', () => {
52+
it('returns the implementation of the `updateSnapState` hook', async () => {
53+
const { store, runSaga } = createStore(getMockOptions());
54+
const fn = getPermittedUpdateSnapStateMethodImplementation(runSaga);
55+
56+
expect(getState(true)(store.getState())).toBeNull();
57+
58+
await fn({ foo: 'bar' }, true);
59+
60+
expect(getState(true)(store.getState())).toStrictEqual(
61+
JSON.stringify({
62+
foo: 'bar',
63+
}),
64+
);
65+
});
66+
67+
it('returns the implementation of the `updateSnapState` hook for unencrypted state', async () => {
68+
const { store, runSaga } = createStore(getMockOptions());
69+
const fn = getPermittedUpdateSnapStateMethodImplementation(runSaga);
70+
71+
expect(getState(false)(store.getState())).toBeNull();
72+
73+
await fn({ foo: 'bar' }, false);
74+
75+
expect(getState(false)(store.getState())).toStrictEqual(
76+
JSON.stringify({
77+
foo: 'bar',
78+
}),
79+
);
80+
});
81+
});
82+
83+
describe('getPermittedClearSnapStateMethodImplementation', () => {
84+
it('returns the implementation of the `clearSnapState` hook', async () => {
85+
const { store, runSaga } = createStore(getMockOptions());
86+
const fn = getPermittedClearSnapStateMethodImplementation(runSaga);
87+
88+
store.dispatch(
89+
setState({
90+
state: JSON.stringify({
91+
foo: 'bar',
92+
}),
93+
encrypted: true,
94+
}),
95+
);
96+
97+
await fn(true);
98+
99+
expect(getState(true)(store.getState())).toBeNull();
100+
});
101+
102+
it('returns the implementation of the `clearSnapState` hook for unencrypted state', async () => {
103+
const { store, runSaga } = createStore(getMockOptions());
104+
const fn = getPermittedClearSnapStateMethodImplementation(runSaga);
105+
106+
store.dispatch(
107+
setState({
108+
state: JSON.stringify({
109+
foo: 'bar',
110+
}),
111+
encrypted: false,
112+
}),
113+
);
114+
115+
await fn(false);
116+
117+
expect(getState(false)(store.getState())).toBeNull();
118+
});
119+
});

‎packages/snaps-simulation/src/methods/hooks/permitted/state.ts‎

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
import { parseJson } from '@metamask/snaps-utils';
2+
import type { Json } from '@metamask/utils';
3+
import type { SagaIterator } from 'redux-saga';
4+
import { put, select } from 'redux-saga/effects';
5+
6+
import type { RunSagaFunction } from '../../../store';
7+
import { clearState, getState, setState } from '../../../store';
8+
9+
/**
10+
* Get the Snap state from the store.
11+
*
12+
* @param encrypted - Whether to get the encrypted or unencrypted state.
13+
* Defaults to encrypted state.
14+
* @returns The state of the Snap.
15+
* @yields Selects the state from the store.
16+
*/
17+
function* getSnapStateImplementation(encrypted: boolean): SagaIterator<string> {
18+
const state = yield select(getState(encrypted));
19+
// TODO: Use actual decryption implementation
20+
return parseJson(state);
21+
}
22+
23+
/**
24+
* Get the implementation of the `getSnapState` hook.
25+
*
26+
* @param runSaga - The function to run a saga outside the usual Redux flow.
27+
* @returns The implementation of the `getSnapState` hook.
28+
*/
29+
export function getPermittedGetSnapStateMethodImplementation(
30+
runSaga: RunSagaFunction,
31+
) {
32+
return async (...args: Parameters<typeof getSnapStateImplementation>) => {
33+
return await runSaga(getSnapStateImplementation, ...args).toPromise();
34+
};
35+
}
36+
37+
/**
38+
* Update the Snap state in the store.
39+
*
40+
* @param newState - The new state.
41+
* @param encrypted - Whether to update the encrypted or unencrypted state.
42+
* Defaults to encrypted state.
43+
* @yields Puts the new state in the store.
44+
*/
45+
function* updateSnapStateImplementation(
46+
newState: Record<string, Json>,
47+
encrypted: boolean,
48+
): SagaIterator<void> {
49+
// TODO: Use actual encryption implementation
50+
yield put(setState({ state: JSON.stringify(newState), encrypted }));
51+
}
52+
53+
/**
54+
* Get the implementation of the `updateSnapState` hook.
55+
*
56+
* @param runSaga - The function to run a saga outside the usual Redux flow.
57+
* @returns The implementation of the `updateSnapState` hook.
58+
*/
59+
export function getPermittedUpdateSnapStateMethodImplementation(
60+
runSaga: RunSagaFunction,
61+
) {
62+
return async (...args: Parameters<typeof updateSnapStateImplementation>) => {
63+
return await runSaga(updateSnapStateImplementation, ...args).toPromise();
64+
};
65+
}
66+
67+
/**
68+
* Clear the Snap state in the store.
69+
*
70+
* @param encrypted - Whether to clear the encrypted or unencrypted state.
71+
* Defaults to encrypted state.
72+
* @yields Puts the new state in the store.
73+
*/
74+
function* clearSnapStateImplementation(encrypted: boolean): SagaIterator<void> {
75+
yield put(clearState({ encrypted }));
76+
}
77+
78+
/**
79+
* Get the implementation of the `clearSnapState` hook.
80+
*
81+
* @param runSaga - The function to run a saga outside the usual Redux flow.
82+
* @returns The implementation of the `clearSnapState` hook.
83+
*/
84+
export function getPermittedClearSnapStateMethodImplementation(
85+
runSaga: RunSagaFunction,
86+
) {
87+
return async (...args: Parameters<typeof clearSnapStateImplementation>) => {
88+
runSaga(clearSnapStateImplementation, ...args).result();
89+
};
90+
}

‎packages/snaps-simulation/src/methods/specifications.test.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ import {
55
} from '@metamask/snaps-utils/test-utils';
66

77
import { getControllers, registerSnap } from '../controllers';
8-
import type { MiddlewareHooks } from '../simulation';
8+
import type { RestrictedMiddlewareHooks } from '../simulation';
99
import { getMockOptions } from '../test-utils/options';
1010
import {
1111
asyncResolve,
@@ -14,7 +14,7 @@ import {
1414
resolve,
1515
} from './specifications';
1616

17-
const MOCK_HOOKS: MiddlewareHooks = {
17+
const MOCK_HOOKS: RestrictedMiddlewareHooks = {
1818
getMnemonic: jest.fn(),
1919
getSnapFile: jest.fn(),
2020
createInterface: jest.fn(),

‎packages/snaps-simulation/src/middleware/engine.test.ts‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,11 +7,18 @@ describe('createJsonRpcEngine', () => {
77
const { store } = createStore(getMockOptions());
88
const engine = createJsonRpcEngine({
99
store,
10-
hooks: {
10+
restrictedHooks: {
1111
getMnemonic: jest.fn(),
12-
getSnapFile: jest.fn().mockResolvedValue('foo'),
1312
getIsLocked: jest.fn(),
13+
getClientCryptography: jest.fn(),
14+
},
15+
permittedHooks: {
16+
getSnapFile: jest.fn().mockResolvedValue('foo'),
17+
getSnapState: jest.fn(),
18+
updateSnapState: jest.fn(),
19+
clearSnapState: jest.fn(),
1420
getInterfaceState: jest.fn(),
21+
getInterfaceContext: jest.fn(),
1522
createInterface: jest.fn(),
1623
updateInterface: jest.fn(),
1724
resolveInterface: jest.fn(),

‎packages/snaps-simulation/src/middleware/engine.ts‎

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -6,14 +6,18 @@ import { createSnapsMethodMiddleware } from '@metamask/snaps-rpc-methods';
66
import type { Json } from '@metamask/utils';
77

88
import { DEFAULT_JSON_RPC_ENDPOINT } from '../constants';
9-
import type { MiddlewareHooks } from '../simulation';
9+
import type {
10+
PermittedMiddlewareHooks,
11+
RestrictedMiddlewareHooks,
12+
} from '../simulation';
1013
import type { Store } from '../store';
1114
import { createInternalMethodsMiddleware } from './internal-methods';
1215
import { createMockMiddleware } from './mock';
1316

1417
export type CreateJsonRpcEngineOptions = {
1518
store: Store;
16-
hooks: MiddlewareHooks;
19+
restrictedHooks: RestrictedMiddlewareHooks;
20+
permittedHooks: PermittedMiddlewareHooks;
1721
permissionMiddleware: JsonRpcMiddleware<RestrictedMethodParameters, Json>;
1822
endpoint?: string;
1923
};
@@ -26,21 +30,23 @@ export type CreateJsonRpcEngineOptions = {
2630
*
2731
* @param options - The options to use when creating the engine.
2832
* @param options.store - The Redux store to use.
29-
* @param options.hooks - Any hooks used by the middleware handlers.
33+
* @param options.restrictedHooks - Any hooks used by the middleware handlers.
34+
* @param options.permittedHooks - Any hooks used by the middleware handlers.
3035
* @param options.permissionMiddleware - The permission middleware to use.
3136
* @param options.endpoint - The JSON-RPC endpoint to use for Ethereum requests.
3237
* @returns A JSON-RPC engine.
3338
*/
3439
export function createJsonRpcEngine({
3540
store,
36-
hooks,
41+
restrictedHooks,
42+
permittedHooks,
3743
permissionMiddleware,
3844
endpoint = DEFAULT_JSON_RPC_ENDPOINT,
3945
}: CreateJsonRpcEngineOptions) {
4046
const engine = new JsonRpcEngine();
4147
engine.push(createMockMiddleware(store));
42-
engine.push(createInternalMethodsMiddleware(hooks));
43-
engine.push(createSnapsMethodMiddleware(true, hooks));
48+
engine.push(createInternalMethodsMiddleware(restrictedHooks));
49+
engine.push(createSnapsMethodMiddleware(true, permittedHooks));
4450
engine.push(permissionMiddleware);
4551
engine.push(
4652
createFetchMiddleware({

0 commit comments

Comments
 (0)