diff --git a/packages/snap-networks-utils/CHANGELOG.md b/packages/snap-networks-utils/CHANGELOG.md index 57c5a84b..aed06335 100644 --- a/packages/snap-networks-utils/CHANGELOG.md +++ b/packages/snap-networks-utils/CHANGELOG.md @@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added +- Add shared snap state helpers `IStateManager`, `State`, and `InMemoryState` (Tron-style write mutex plus blob/path locking). ([#288](https://github.com/MetaMask/internal-snaps/pull/288)) - Add shared proof-of-ownership message parsing utilities, batch request/response structs, and batch request/response types. ([#268](https://github.com/MetaMask/internal-snaps/pull/268)) - Add a `UuidStruct` Superstruct for validating UUID v4 strings. ([#243](https://github.com/MetaMask/internal-snaps/pull/243)) - Add helpers `serialize`, `deserialize`, and `Serializable` for round-tripping `BigNumber`, `bigint`, `Uint8Array`, and `undefined` through snap state ([#197](https://github.com/MetaMask/internal-snaps/pull/197)) diff --git a/packages/snap-networks-utils/package.json b/packages/snap-networks-utils/package.json index 636a0a6a..7a70b479 100644 --- a/packages/snap-networks-utils/package.json +++ b/packages/snap-networks-utils/package.json @@ -76,6 +76,7 @@ "@metamask/snaps-sdk": "^12.0.1", "@metamask/superstruct": "^3.4.1", "@metamask/utils": "^11.11.0", + "async-mutex": "^0.5.0", "bignumber.js": "^9.3.1", "lodash": "^4.17.21" }, diff --git a/packages/snap-networks-utils/src/index.ts b/packages/snap-networks-utils/src/index.ts index 77ed424d..860ce9b9 100644 --- a/packages/snap-networks-utils/src/index.ts +++ b/packages/snap-networks-utils/src/index.ts @@ -62,6 +62,10 @@ export { } from './utils/originPermissions/createOriginPermissions'; export type { CreateOriginPermissionsParams } from './utils/originPermissions/createOriginPermissions'; export { validateOrigin } from './utils/originPermissions/validateOrigin'; +export { State } from './utils/state/State'; +export type { StateConfig } from './utils/state/State'; +export { InMemoryState } from './utils/state/InMemoryState'; +export type { IStateManager } from './utils/state/IStateManager'; export { createSnapErrorHandling, createTrackError, diff --git a/packages/snap-networks-utils/src/utils/state/IStateManager.ts b/packages/snap-networks-utils/src/utils/state/IStateManager.ts new file mode 100644 index 00000000..cdb9e922 --- /dev/null +++ b/packages/snap-networks-utils/src/utils/state/IStateManager.ts @@ -0,0 +1,76 @@ +import type { Serializable } from '../serialization/types'; + +export type IStateManager> = { + /** + * Gets the whole state object. + * + * ⚠️ WARNING: Use with caution because it transfers the whole state, which might contain a lot of data. + * If you need to retrieve only a specific part of the state, use IStateManager.getKey instead. + */ + get(): Promise; + /** + * Gets the value of the passed key in the state object. + * The key is the JSON path to the value to get. + * + * @returns The value of the key, or undefined if the key does not exist. + */ + getKey( + key: string, + ): Promise; + /** + * Sets the value of the passed key in the state object. + * The key is a JSON path to the value to set. + * + * @param key - The key to set, which is a JSON path to the location. + * @param value - The value to set. + */ + setKey(key: string, value: Serializable): Promise; + /** + * Atomically reads the current value at `key`, applies `updater`, and writes the result back. + * + * Implementations must ensure that no concurrent state write can interleave between the + * read and the write. This makes the method safe for updates where the next value depends + * on the current value, such as merging objects. + * + * Prefer this over a manual `getKey` + `setKey` sequence whenever the new value depends on + * the current one. + * + * @param key - The JSON-path key to update. + * @param updater - Receives the current value (or `undefined` when the key is absent) and + * returns the new value to store. + */ + setKeyWith( + key: string, + updater: (currentValue: TValue | undefined) => TValue, + ): Promise; + /** + * Updates the whole state object. + * + * Typically used for bulk `set`s or `delete`s, because: + * - Atomicity: Using a single `state.update` ensures that all changes are applied atomically. + * - Performance: One round trip instead of many `setKey` / `deleteKey` calls. + * - State Consistency: Read once, modify in memory, write the complete updated state back. + * + * ⚠️ WARNING: Use with caution because: + * - it will override the whole state. + * - it transfers the whole state back and forth to the data store. + * + * For single updates, use `setKey` or `deleteKey` instead. + * + * @param updaterFunction - The function that updates the state. + * @returns The updated state. + */ + update( + updaterFunction: (state: TStateValue) => TStateValue, + ): Promise; + /** + * Deletes the value of the passed key in the state object. + * The key is a JSON path to the value to delete. + */ + deleteKey(key: string): Promise; + /** + * Deletes multiple keys in the state object in a single operation. + * The keys are JSON paths to the values to delete. + */ + deleteKeys(keys: string[]): Promise; +}; diff --git a/packages/snap-networks-utils/src/utils/state/InMemoryState.test.ts b/packages/snap-networks-utils/src/utils/state/InMemoryState.test.ts new file mode 100644 index 00000000..987458b7 --- /dev/null +++ b/packages/snap-networks-utils/src/utils/state/InMemoryState.test.ts @@ -0,0 +1,40 @@ +import { InMemoryState } from './InMemoryState'; + +type MockStateValue = { + users: { name: string; age?: number }[]; +}; + +describe('InMemoryState', () => { + it('gets, sets, and deletes keys', async () => { + const state = new InMemoryState({ + users: [{ name: 'John', age: 30 }], + }); + + expect(await state.get()).toStrictEqual({ + users: [{ name: 'John', age: 30 }], + }); + expect(await state.getKey('users.0.name')).toBe('John'); + + await state.setKey('users.0.name', 'Jane'); + expect(await state.getKey('users.0.name')).toBe('Jane'); + + await state.setKeyWith('users.0.age', (age) => (age ?? 0) + 1); + expect(await state.getKey('users.0.age')).toBe(31); + + await state.deleteKey('users.0.age'); + expect(await state.get()).toStrictEqual({ users: [{ name: 'Jane' }] }); + + await state.deleteKeys(['users']); + expect(await state.get()).toStrictEqual({}); + }); + + it('replaces the whole state via update', async () => { + const state = new InMemoryState({ + users: [{ name: 'John' }], + }); + + expect( + await state.update(() => ({ users: [{ name: 'Bob', age: 50 }] })), + ).toStrictEqual({ users: [{ name: 'Bob', age: 50 }] }); + }); +}); diff --git a/packages/snap-networks-utils/src/utils/state/InMemoryState.ts b/packages/snap-networks-utils/src/utils/state/InMemoryState.ts new file mode 100644 index 00000000..bca1648f --- /dev/null +++ b/packages/snap-networks-utils/src/utils/state/InMemoryState.ts @@ -0,0 +1,57 @@ +import { get, set, unset } from 'lodash'; + +import type { Serializable } from '../serialization/types'; +import type { IStateManager } from './IStateManager'; + +/** + * A simple implementation of the `IStateManager` interface that relies on an in-memory + * state. Intended for tests. + */ +export class InMemoryState< + TStateValue extends Record, +> implements IStateManager { + #state: TStateValue; + + constructor(initialState: TStateValue) { + this.#state = initialState; + } + + async get(): Promise { + return this.#state; + } + + async getKey( + key: string, + ): Promise { + return get(this.#state, key) as TResponse | undefined; + } + + async setKey(key: string, value: Serializable): Promise { + set(this.#state, key, value); + } + + async setKeyWith( + key: string, + updater: (currentValue: TValue | undefined) => TValue, + ): Promise { + const oldValue = get(this.#state, key) as TValue | undefined; + set(this.#state, key, updater(oldValue)); + } + + async update( + callback: (state: TStateValue) => TStateValue, + ): Promise { + this.#state = callback(this.#state); + return this.#state; + } + + async deleteKey(key: string): Promise { + unset(this.#state, key); + } + + async deleteKeys(keys: string[]): Promise { + keys.forEach((key) => { + unset(this.#state, key); + }); + } +} diff --git a/packages/snap-networks-utils/src/utils/state/State.test.ts b/packages/snap-networks-utils/src/utils/state/State.test.ts new file mode 100644 index 00000000..d73c19fe --- /dev/null +++ b/packages/snap-networks-utils/src/utils/state/State.test.ts @@ -0,0 +1,571 @@ +/* eslint-disable jest/prefer-strict-equal */ +import { BigNumber } from 'bignumber.js'; + +import { State } from './State'; + +const snap = { + request: jest.fn(), +}; + +(globalThis as typeof globalThis & { snap: typeof snap }).snap = snap; + +const flushPromises = async (): Promise => { + await new Promise((resolve) => { + setTimeout(resolve, 0); + }); +}; + +type DelayedStateRequest = + | { + method: 'snap_getState'; + params: { key: string }; + } + | { + method: 'snap_setState'; + params: { key: string; value: Record }; + }; + +/** + * Mocks state reads so tests can hold pending `snap_getState` calls. + * + * @param storedValues - Mutable backing store returned from mocked state reads. + * @param getResolvers - Resolver queue for delayed mocked state reads. + */ +function mockDelayedStateRequests( + storedValues: Record>, + getResolvers: (() => void)[], +): void { + snap.request.mockImplementation(async (request: DelayedStateRequest) => { + if (request.method === 'snap_getState') { + await new Promise((resolve) => { + getResolvers.push(resolve); + }); + return storedValues[request.params.key] ?? null; + } + + if (request.method === 'snap_setState') { + storedValues[request.params.key] = request.params.value; + } + + return undefined; + }); +} + +type User = { + name: string; + age: BigNumber | bigint | number | undefined | null; +}; + +type MockStateValue = { + users: User[]; +}; + +const DEFAULT_STATE: MockStateValue = { + users: [ + { + name: 'John', + age: 30, + }, + { + name: 'Jane', + age: 25, + }, + ], +}; + +describe('State', () => { + let state: State; + + beforeEach(() => { + state = new State({ + encrypted: false, + defaultState: DEFAULT_STATE, + }); + + jest.clearAllMocks(); + }); + + afterEach(() => { + snap.request.mockReset(); + }); + + describe('get', () => { + it('gets the state', async () => { + const mockUnderlyingState = DEFAULT_STATE; + snap.request.mockResolvedValue(mockUnderlyingState); + + const stateValue = await state.get(); + + expect(snap.request).toHaveBeenCalledWith({ + method: 'snap_getState', + params: { encrypted: false }, + }); + expect(stateValue).toStrictEqual(mockUnderlyingState); + }); + + it('gets the default state if the snap state is empty', async () => { + snap.request.mockResolvedValue({}); + + expect(await state.get()).toStrictEqual(DEFAULT_STATE); + }); + + it('preserves defaults when persisted state values are undefined', async () => { + snap.request.mockResolvedValue({ users: undefined }); + + expect(await state.get()).toStrictEqual(DEFAULT_STATE); + }); + + it('allows concurrent path reads', async () => { + const getResolvers: (() => void)[] = []; + snap.request.mockImplementation(async () => { + await new Promise((resolve) => { + getResolvers.push(resolve); + }); + return {}; + }); + + const first = state.get(); + const second = state.get(); + + await flushPromises(); + + expect(getResolvers).toHaveLength(2); + + getResolvers.forEach((resolve) => resolve()); + + await Promise.all([first, second]); + }); + + describe('when getting serialized non-JSON values', () => { + it('deserializes undefined values', async () => { + snap.request.mockResolvedValue({ + users: [ + { + name: 'John', + age: { + __type: 'undefined', + }, + }, + ], + }); + + expect(await state.get()).toEqual({ + users: [ + { + name: 'John', + age: undefined, + }, + ], + }); + }); + + it('deserializes BigNumber values', async () => { + snap.request.mockResolvedValue({ + users: [ + { + name: 'John', + age: { + __type: 'BigNumber', + value: '30', + }, + }, + ], + }); + + expect(await state.get()).toStrictEqual({ + users: [ + { + name: 'John', + age: new BigNumber(30), + }, + ], + }); + }); + + it('deserializes bigint values', async () => { + snap.request.mockResolvedValue({ + users: [ + { + name: 'John', + age: { + __type: 'bigint', + value: '30', + }, + }, + ], + }); + + expect(await state.get()).toStrictEqual({ + users: [ + { + name: 'John', + age: BigInt(30), + }, + ], + }); + }); + }); + }); + + describe('getKey', () => { + it('calls the snap_getState method with the correct parameters', async () => { + snap.request.mockResolvedValue(DEFAULT_STATE); + + await state.getKey('users.1.name'); + + expect(snap.request).toHaveBeenCalledWith({ + method: 'snap_getState', + params: { key: 'users.1.name', encrypted: false }, + }); + }); + + it('returns undefined if the key does not exist', async () => { + snap.request.mockResolvedValue(null); + + expect(await state.getKey('users.1.name')).toBeUndefined(); + }); + }); + + describe('setKey', () => { + it('sets the value of a key', async () => { + await state.setKey('users.1.name', 'Bob'); + + expect(snap.request).toHaveBeenCalledWith({ + method: 'snap_setState', + params: { + key: 'users.1.name', + value: 'Bob', + encrypted: false, + }, + }); + }); + }); + + describe('setKeyWith', () => { + it('reads the current value, applies the updater, and writes the result', async () => { + snap.request.mockResolvedValueOnce({ alice: 10 }); + + await state.setKeyWith>('scores', (current) => ({ + ...current, + bob: 20, + })); + + expect(snap.request).toHaveBeenNthCalledWith(1, { + method: 'snap_getState', + params: { key: 'scores', encrypted: false }, + }); + expect(snap.request).toHaveBeenNthCalledWith(2, { + method: 'snap_setState', + params: { + key: 'scores', + value: { alice: 10, bob: 20 }, + encrypted: false, + }, + }); + }); + + it('passes undefined to the updater when the key does not exist', async () => { + snap.request.mockResolvedValueOnce(null); + + const updater = jest.fn().mockReturnValue({ bob: 20 }); + + await state.setKeyWith('scores', updater); + + expect(updater).toHaveBeenCalledWith(undefined); + }); + + it('serializes concurrent read-modify-write updates', async () => { + const storedValues: Record> = { + scores: { alice: 10 }, + }; + const getResolvers: (() => void)[] = []; + + mockDelayedStateRequests(storedValues, getResolvers); + + const firstUpdate = state.setKeyWith>( + 'scores', + (current) => ({ + ...current, + bob: 20, + }), + ); + const secondUpdate = state.setKeyWith>( + 'scores', + (current) => ({ + ...current, + carol: 30, + }), + ); + + await flushPromises(); + + expect(getResolvers).toHaveLength(1); + + getResolvers[0]?.(); + + await flushPromises(); + + expect(getResolvers).toHaveLength(2); + + getResolvers[1]?.(); + + await Promise.all([firstUpdate, secondUpdate]); + + expect(storedValues.scores).toStrictEqual({ + alice: 10, + bob: 20, + carol: 30, + }); + }); + }); + + describe('update', () => { + it('does not admit path operations while an update is waiting', async () => { + const getResolvers: (() => void)[] = []; + snap.request.mockImplementation(async (request) => { + if (request.method === 'snap_getState') { + await new Promise((resolve) => { + getResolvers.push(resolve); + }); + return {}; + } + + return null; + }); + + const firstRead = state.get(); + await flushPromises(); + expect(getResolvers).toHaveLength(1); + + const update = state.update((currentState) => currentState); + const secondRead = state.get(); + await flushPromises(); + + // The update owns admission while it waits for the first read to finish. + expect(getResolvers).toHaveLength(1); + + getResolvers[0]?.(); + await flushPromises(); + expect(getResolvers).toHaveLength(2); + + // The update's read has started, but the second read remains blocked. + getResolvers[1]?.(); + await flushPromises(); + expect(getResolvers).toHaveLength(3); + + getResolvers[2]?.(); + await Promise.all([firstRead, update, secondRead]); + }); + + it('updates the state', async () => { + await state.update((currentState) => ({ + users: [ + ...currentState.users, + { + name: 'Bob', + age: 50, + }, + ], + })); + + expect(snap.request).toHaveBeenCalledWith({ + method: 'snap_getState', + params: { encrypted: false }, + }); + + expect(snap.request).toHaveBeenCalledWith({ + method: 'snap_manageState', + params: { + operation: 'update', + encrypted: false, + newState: { + users: [ + ...DEFAULT_STATE.users, + { + name: 'Bob', + age: 50, + }, + ], + }, + }, + }); + }); + + describe('when updating serialized non-JSON values', () => { + it('serializes undefined values', async () => { + await state.update((currentState) => ({ + users: [ + ...currentState.users, + { + name: 'Bob', + age: undefined, + }, + ], + })); + + expect(snap.request).toHaveBeenNthCalledWith(2, { + method: 'snap_manageState', + params: { + operation: 'update', + encrypted: false, + newState: { + users: [ + ...DEFAULT_STATE.users, + { + name: 'Bob', + age: { + __type: 'undefined', + }, + }, + ], + }, + }, + }); + }); + + it('serializes BigNumber values', async () => { + await state.update((currentState) => ({ + users: [ + ...currentState.users, + { + name: 'Bob', + age: new BigNumber(50), + }, + ], + })); + + expect(snap.request).toHaveBeenNthCalledWith(2, { + method: 'snap_manageState', + params: { + operation: 'update', + encrypted: false, + newState: { + users: [ + ...DEFAULT_STATE.users, + { + name: 'Bob', + age: { + __type: 'BigNumber', + value: '50', + }, + }, + ], + }, + }, + }); + }); + + it('serializes bigint values', async () => { + await state.update((currentState) => ({ + users: [ + ...currentState.users, + { + name: 'Bob', + age: BigInt(50), + }, + ], + })); + + expect(snap.request).toHaveBeenNthCalledWith(2, { + method: 'snap_manageState', + params: { + operation: 'update', + encrypted: false, + newState: { + users: [ + ...DEFAULT_STATE.users, + { + name: 'Bob', + age: { + __type: 'bigint', + value: '50', + }, + }, + ], + }, + }, + }); + }); + + it('serializes null values', async () => { + await state.update((currentState) => ({ + users: [...currentState.users, { name: 'Bob', age: null }], + })); + + expect(snap.request).toHaveBeenNthCalledWith(2, { + method: 'snap_manageState', + params: { + operation: 'update', + encrypted: false, + newState: { + users: [...DEFAULT_STATE.users, { name: 'Bob', age: null }], + }, + }, + }); + }); + }); + }); + + describe('deleteKey', () => { + it('deletes a key', async () => { + await state.deleteKey('users'); + + expect(snap.request).toHaveBeenCalledWith({ + method: 'snap_manageState', + params: { + operation: 'update', + newState: {}, + encrypted: false, + }, + }); + }); + + it('does not mutate the shared default state', async () => { + // Empty persisted state means the defaults are what the updater receives. + snap.request.mockResolvedValue({}); + + await state.deleteKey('users[0].age'); + + expect(DEFAULT_STATE.users[0]).toStrictEqual({ name: 'John', age: 30 }); + }); + + it('deletes a nested key', async () => { + await state.deleteKey('users[0].age'); + + expect(snap.request).toHaveBeenCalledWith({ + method: 'snap_manageState', + params: { + operation: 'update', + newState: { + users: [ + { + name: 'John', + }, + { + name: 'Jane', + age: 25, + }, + ], + }, + encrypted: false, + }, + }); + }); + }); + + describe('deleteKeys', () => { + it('deletes multiple keys', async () => { + await state.deleteKeys(['users.0.age', 'users.1.name']); + + expect(snap.request).toHaveBeenCalledWith({ + method: 'snap_manageState', + params: { + operation: 'update', + newState: { users: [{ name: 'John' }, { age: 25 }] }, + encrypted: false, + }, + }); + }); + }); +}); +/* eslint-enable jest/prefer-strict-equal */ diff --git a/packages/snap-networks-utils/src/utils/state/State.ts b/packages/snap-networks-utils/src/utils/state/State.ts new file mode 100644 index 00000000..301a860f --- /dev/null +++ b/packages/snap-networks-utils/src/utils/state/State.ts @@ -0,0 +1,221 @@ +import type { Json, SnapsProvider } from '@metamask/snaps-sdk'; +import type { MutexInterface } from 'async-mutex'; +import { Mutex } from 'async-mutex'; +import { cloneDeep, unset } from 'lodash'; + +import { safeMerge } from '../safeMerge/safeMerge'; +import { deserialize, serialize } from '../serialization/serialization'; +import type { Serializable } from '../serialization/types'; +import type { IStateManager } from './IStateManager'; + +export type StateConfig> = { + encrypted: boolean; + defaultState: TValue; +}; + +/** + * Resolves the Snap RPC client from the global `snap` object provided by the Snap runtime. + * + * @returns The Snap `request` function. + */ +const getSnapRequest = (): SnapsProvider['request'] => + (globalThis as typeof globalThis & { snap: SnapsProvider }).snap.request; + +/** + * Because we use both snap_manageState and snap_setState, we must protect against them + * being used at the same time. We must also protect against multiple parallel requests + * to snap_manageState. + * + * Path writes (`setKey` / `setKeyWith`) are serialized with a dedicated write mutex so + * concurrent read-modify-write updates cannot interleave. Path reads may still run in + * parallel with each other. A blob (`snap_manageState`) operation waits for in-flight path + * operations to finish and blocks new ones from starting until it completes. + */ +class StateLock { + // Gate every operation must pass through to start. Path operations hold it only while + // registering; a blob operation holds it for its whole duration, so no path operation + // can start after a blob operation begins waiting for in-flight ones to finish. + readonly #operationAdmissionMutex = new Mutex(); + + // Held while at least one path operation is in flight. + readonly #regularStateUpdateMutex = new Mutex(); + + readonly #regularStateWriteMutex = new Mutex(); + + #pendingRegularStateUpdates = 0; + + #releaseRegularStateUpdateMutex: MutexInterface.Releaser | null = null; + + async #acquireRegularStateUpdateMutex(): Promise { + if (!this.#regularStateUpdateMutex.isLocked()) { + this.#releaseRegularStateUpdateMutex = + await this.#regularStateUpdateMutex.acquire(); + } + } + + async wrapRegularStateOperation( + callback: MutexInterface.Worker, + ): Promise { + await this.#operationAdmissionMutex.runExclusive(async () => { + await this.#acquireRegularStateUpdateMutex(); + this.#pendingRegularStateUpdates += 1; + }); + + try { + return await callback(); + } finally { + this.#pendingRegularStateUpdates -= 1; + + if ( + this.#pendingRegularStateUpdates === 0 && + this.#releaseRegularStateUpdateMutex + ) { + this.#releaseRegularStateUpdateMutex(); + this.#releaseRegularStateUpdateMutex = null; + } + } + } + + async wrapRegularStateWriteOperation( + callback: MutexInterface.Worker, + ): Promise { + return await this.#regularStateWriteMutex.runExclusive(async () => + this.wrapRegularStateOperation(callback), + ); + } + + async wrapManageStateOperation( + callback: MutexInterface.Worker, + ): Promise { + return await this.#operationAdmissionMutex.runExclusive(async () => { + await this.#regularStateUpdateMutex.waitForUnlock(); + + return await callback(); + }); + } +} + +/** + * Layer on top of `snap_manageState` / `snap_getState` / `snap_setState`: + * + * - Serializes values before storing them and deserializes after reading. + * - Merges `defaultState` on full-blob reads via `safeMerge`. + * - Serializes path writes and full-blob updates (lock strategy B). + */ +export class State< + TStateValue extends Record, +> implements IStateManager { + readonly #lock = new StateLock(); + + readonly #config: StateConfig; + + constructor(config: StateConfig) { + this.#config = config; + } + + /** + * Reads and deserializes the value at `key`, or the whole state blob when `key` is omitted. + * + * @param key - The JSON-path key to read. Omit to read the whole blob. + * @returns The deserialized value, or `undefined` when the key is absent. + */ + async #read( + key?: string, + ): Promise { + const value = await getSnapRequest()({ + method: 'snap_getState', + params: { + ...(key === undefined ? {} : { key }), + encrypted: this.#config.encrypted, + }, + }); + + return value === null || value === undefined + ? undefined + : (deserialize(value) as TValue); + } + + /** + * Serializes `value` and writes it to `key`. + * + * @param key - The JSON-path key to write. + * @param value - The value to store. + */ + async #write(key: string, value: Serializable): Promise { + await getSnapRequest()({ + method: 'snap_setState', + params: { + key, + value: serialize(value), + encrypted: this.#config.encrypted, + }, + }); + } + + async #unsafeGet(): Promise { + const state = (await this.#read()) ?? ({} as TStateValue); + + // Clone the defaults so updaters that mutate the returned state (e.g. `deleteKey` + // via lodash `unset`) never leak into the shared `defaultState` object. + return safeMerge(cloneDeep(this.#config.defaultState), state); + } + + async get(): Promise { + return this.#lock.wrapRegularStateOperation(async () => this.#unsafeGet()); + } + + async getKey( + key: string, + ): Promise { + return this.#lock.wrapRegularStateOperation(async () => + this.#read(key), + ); + } + + async setKey(key: string, value: Serializable): Promise { + await this.#lock.wrapRegularStateWriteOperation(async () => + this.#write(key, value), + ); + } + + async setKeyWith( + key: string, + updater: (currentValue: TValue | undefined) => TValue, + ): Promise { + await this.#lock.wrapRegularStateWriteOperation(async () => + this.#write(key, updater(await this.#read(key))), + ); + } + + async update( + updaterFunction: (state: TStateValue) => TStateValue, + ): Promise { + return await this.#lock.wrapManageStateOperation(async () => { + const newState = updaterFunction(await this.#unsafeGet()); + + await getSnapRequest()({ + method: 'snap_manageState', + params: { + operation: 'update', + newState: serialize(newState) as Record, + encrypted: this.#config.encrypted, + }, + }); + + return newState; + }); + } + + async deleteKey(key: string): Promise { + await this.deleteKeys([key]); + } + + async deleteKeys(keys: string[]): Promise { + await this.update((state) => { + keys.forEach((key) => { + unset(state, key); + }); + return state; + }); + } +} diff --git a/yarn.lock b/yarn.lock index 7cc65247..545e9e1f 100644 --- a/yarn.lock +++ b/yarn.lock @@ -3335,6 +3335,7 @@ __metadata: "@ts-bridge/cli": "npm:^0.6.4" "@types/jest": "npm:^30.0.0" "@types/lodash": "npm:^4.17.15" + async-mutex: "npm:^0.5.0" bignumber.js: "npm:^9.3.1" deepmerge: "npm:^4.2.2" jest: "npm:30.0.3"