diff --git a/eslint-suppressions.json b/eslint-suppressions.json index 63a2c50c5..acd48b2b7 100644 --- a/eslint-suppressions.json +++ b/eslint-suppressions.json @@ -1502,14 +1502,6 @@ "count": 7 } }, - "packages/tron-wallet-snap/src/services/state/State.test.ts": { - "@typescript-eslint/explicit-function-return-type": { - "count": 1 - }, - "@typescript-eslint/no-explicit-any": { - "count": 1 - } - }, "packages/tron-wallet-snap/src/services/transaction-scan/TransactionScanService.ts": { "no-restricted-syntax": { "count": 5 diff --git a/packages/snap-networks-utils/src/utils/cache/StateCache.test.ts b/packages/snap-networks-utils/src/utils/cache/StateCache.test.ts index 4c7884d51..2bb980f1a 100644 --- a/packages/snap-networks-utils/src/utils/cache/StateCache.test.ts +++ b/packages/snap-networks-utils/src/utils/cache/StateCache.test.ts @@ -1,50 +1,10 @@ /* eslint-disable jest/prefer-strict-equal */ -import { get, set, unset } from 'lodash'; - import { Logger, LogLevel } from '../logger/Logger'; -import type { Serializable } from '../serialization/types'; -import type { CacheStateManager } from './StateCache'; -import type { StateValue } from './StateCache'; +import { InMemoryState } from '../state/InMemoryState'; +import type { CacheStateManager, StateValue } from './StateCache'; import { StateCache } from './StateCache'; -/** - * A simple implementation of a state manager that relies on an in-memory state, - * used for testing purposes. - */ -class InMemoryState implements CacheStateManager { - #state: StateValue; - - constructor(initialState: StateValue) { - this.#state = initialState; - } - - async get(): Promise { - return this.#state; - } - - async getKey( - key: string, - ): Promise { - return get(this.#state, key) as TKey | undefined; - } - - async setKey(key: string, value: Serializable): Promise { - set(this.#state, key, value); // Use lodash to set the value using a json path - } - - async update( - callback: (state: StateValue) => StateValue, - ): Promise { - return (this.#state = callback(this.#state)); - } - - async deleteKey(key: string): Promise { - // Using lodash's unset to leverage the json path capabilities - unset(this.#state, key); - } -} - describe('StateCache', () => { let logger: Logger; diff --git a/packages/tron-wallet-snap/package.json b/packages/tron-wallet-snap/package.json index e97493133..c116bef10 100644 --- a/packages/tron-wallet-snap/package.json +++ b/packages/tron-wallet-snap/package.json @@ -61,7 +61,6 @@ "@metamask/superstruct": "^3.4.1", "@metamask/utils": "^11.11.0", "@types/lodash": "^4.17.15", - "async-mutex": "^0.5.0", "bignumber.js": "^9.3.1", "concurrently": "^10.0.3", "dotenv": "^17.2.1", diff --git a/packages/tron-wallet-snap/snap.manifest.json b/packages/tron-wallet-snap/snap.manifest.json index 8e5f3d2c8..d63bd8cb6 100644 --- a/packages/tron-wallet-snap/snap.manifest.json +++ b/packages/tron-wallet-snap/snap.manifest.json @@ -7,7 +7,7 @@ "url": "https://github.com/MetaMask/internal-snaps.git" }, "source": { - "shasum": "nbBLLJNTo42N6FSp8TuKDy1SlsTrBCYHXEvSk8dHIQA=", + "shasum": "Ur39uZHrtG9fCP8q7mzi14zsrtqB3F/4BKHq0dwZxiU=", "location": { "npm": { "filePath": "dist/bundle.js", diff --git a/packages/tron-wallet-snap/src/context.ts b/packages/tron-wallet-snap/src/context.ts index 6745eda27..a5026a490 100644 --- a/packages/tron-wallet-snap/src/context.ts +++ b/packages/tron-wallet-snap/src/context.ts @@ -2,10 +2,12 @@ import { AssetsProvider, InMemoryCache, RemoteFeatureFlagsProvider, + State, StateCache, } from '@metamask/snap-networks-utils'; import type { AssetsProviderMessenger, + IStateManager, RemoteFeatureFlagsProviderMessenger, } from '@metamask/snap-networks-utils'; import { getMessenger } from '@metamask/snaps-sdk'; @@ -33,8 +35,8 @@ import { ConfirmationHandler } from './services/confirmation/ConfirmationHandler import { FeeCalculatorService } from './services/send/FeeCalculatorService'; import { SendService } from './services/send/SendService'; import { StakingService } from './services/staking/StakingService'; -import type { UnencryptedStateValue } from './services/state/State'; -import { State } from './services/state/State'; +import { DEFAULT_UNENCRYPTED_STATE } from './services/state/stateTypes'; +import type { UnencryptedStateValue } from './services/state/stateTypes'; import { TransactionExpirationRefresherService } from './services/transaction-expiration-refresher/TransactionExpirationRefresherService'; import { TransactionScanService } from './services/transaction-scan/TransactionScanService'; import { TransactionsRepository } from './services/transactions/TransactionsRepository'; @@ -59,12 +61,7 @@ export const configProvider = new ConfigProvider(); const state = new State({ encrypted: false, - defaultState: { - keyringAccounts: {}, - assets: {}, - transactions: {}, - mapInterfaceNameToId: {}, - }, + defaultState: DEFAULT_UNENCRYPTED_STATE, }); const snapClient = new SnapClient(); @@ -261,7 +258,7 @@ export type SnapExecutionContext = { /** * Services */ - state: State; + state: IStateManager; priceApiClient: PriceApiClient; feeCalculatorService: FeeCalculatorService; assetsService: AssetsService; diff --git a/packages/tron-wallet-snap/src/handlers/cronjob/cronjob.test.tsx b/packages/tron-wallet-snap/src/handlers/cronjob/cronjob.test.tsx index c6c02294a..6d4b78e70 100644 --- a/packages/tron-wallet-snap/src/handlers/cronjob/cronjob.test.tsx +++ b/packages/tron-wallet-snap/src/handlers/cronjob/cronjob.test.tsx @@ -1,3 +1,5 @@ +import type { IStateManager } from '@metamask/snap-networks-utils'; + import type { PriceApiClient } from '../../clients/price-api/PriceApiClient'; import type { SnapClient } from '../../clients/snap/SnapClient'; import type { TronHttpClient } from '../../clients/tron-http/TronHttpClient'; @@ -8,7 +10,7 @@ import { TRACK_TX_MAX_ATTEMPTS, } from '../../constants'; import type { AccountsService } from '../../services/accounts/AccountsService'; -import type { State, UnencryptedStateValue } from '../../services/state/State'; +import type { UnencryptedStateValue } from '../../services/state/stateTypes'; import { TransactionExpirationRefresherService } from '../../services/transaction-expiration-refresher/TransactionExpirationRefresherService'; import type { JsonTransactionRawData } from '../../services/transaction-expiration-refresher/types'; import type { TransactionScanService } from '../../services/transaction-scan/TransactionScanService'; @@ -387,7 +389,7 @@ function buildCronHandler({ logger: mockLogger, accountsService: {} as AccountsService, snapClient: mockSnapClient as unknown as SnapClient, - state: mockState as unknown as State, + state: mockState as unknown as IStateManager, priceApiClient: {} as PriceApiClient, tronHttpClient: {} as TronHttpClient, transactionScanService: @@ -949,7 +951,7 @@ describe('CronHandler', () => { logger: mockLogger, accountsService: mockAccountsService as unknown as AccountsService, snapClient: mockSnapClient as unknown as SnapClient, - state: {} as unknown as State, + state: {} as unknown as IStateManager, priceApiClient: {} as PriceApiClient, tronHttpClient: mockTronHttpClient as unknown as TronHttpClient, transactionScanService: {} as unknown as TransactionScanService, diff --git a/packages/tron-wallet-snap/src/handlers/cronjob/cronjob.tsx b/packages/tron-wallet-snap/src/handlers/cronjob/cronjob.tsx index b3e75b3fe..4d00a9717 100644 --- a/packages/tron-wallet-snap/src/handlers/cronjob/cronjob.tsx +++ b/packages/tron-wallet-snap/src/handlers/cronjob/cronjob.tsx @@ -1,4 +1,4 @@ -import type { Logger } from '@metamask/snap-networks-utils'; +import type { IStateManager, Logger } from '@metamask/snap-networks-utils'; import type { JsonRpcRequest } from '@metamask/snaps-sdk'; import type { PriceApiClient } from '../../clients/price-api/PriceApiClient'; @@ -8,7 +8,7 @@ import type { Network } from '../../constants'; import { TRACK_TX_INTERVAL, TRACK_TX_MAX_ATTEMPTS } from '../../constants'; import type { TronKeyringAccount } from '../../entities/keyring-account'; import type { AccountsService } from '../../services/accounts/AccountsService'; -import type { State, UnencryptedStateValue } from '../../services/state/State'; +import type { UnencryptedStateValue } from '../../services/state/stateTypes'; import type { TransactionExpirationRefresherService } from '../../services/transaction-expiration-refresher/TransactionExpirationRefresherService'; import type { JsonTransactionRawData, @@ -52,7 +52,7 @@ export class CronHandler { readonly #snapClient: SnapClient; - readonly #state: State; + readonly #state: IStateManager; readonly #priceApiClient: PriceApiClient; @@ -75,7 +75,7 @@ export class CronHandler { logger: Logger; accountsService: AccountsService; snapClient: SnapClient; - state: State; + state: IStateManager; priceApiClient: PriceApiClient; tronHttpClient: TronHttpClient; transactionScanService: TransactionScanService; diff --git a/packages/tron-wallet-snap/src/services/accounts/AccountsRepository.test.ts b/packages/tron-wallet-snap/src/services/accounts/AccountsRepository.test.ts index 7224b8d79..c902fa951 100644 --- a/packages/tron-wallet-snap/src/services/accounts/AccountsRepository.test.ts +++ b/packages/tron-wallet-snap/src/services/accounts/AccountsRepository.test.ts @@ -1,8 +1,8 @@ import { TrxAccountType, TrxScope } from '@metamask/keyring-api'; +import { InMemoryState } from '@metamask/snap-networks-utils'; import type { TronKeyringAccount } from '../../entities/keyring-account'; -import { InMemoryState } from '../state/InMemoryState'; -import type { UnencryptedStateValue } from '../state/State'; +import type { UnencryptedStateValue } from '../state/stateTypes'; import { AccountsRepository } from './AccountsRepository'; /** diff --git a/packages/tron-wallet-snap/src/services/accounts/AccountsRepository.ts b/packages/tron-wallet-snap/src/services/accounts/AccountsRepository.ts index 31879c32e..0573a5f0d 100644 --- a/packages/tron-wallet-snap/src/services/accounts/AccountsRepository.ts +++ b/packages/tron-wallet-snap/src/services/accounts/AccountsRepository.ts @@ -1,8 +1,8 @@ import type { EntropySourceId } from '@metamask/keyring-api'; +import type { IStateManager } from '@metamask/snap-networks-utils'; import type { TronKeyringAccount } from '../../entities/keyring-account'; -import type { IStateManager } from '../state/IStateManager'; -import type { UnencryptedStateValue } from '../state/State'; +import type { UnencryptedStateValue } from '../state/stateTypes'; /** * Range of inclusive account indices to create. @@ -198,10 +198,10 @@ export class AccountsRepository { } async delete(id: string): Promise { - await Promise.all([ - this.#state.deleteKey(`${this.#storageKey}.${id}`), - this.#state.deleteKey(`assets.${id}`), - this.#state.deleteKey(`transactions.${id}`), + await this.#state.deleteKeys([ + `${this.#storageKey}.${id}`, + `assets.${id}`, + `transactions.${id}`, ]); } } diff --git a/packages/tron-wallet-snap/src/services/assets/AssetsRepository.test.ts b/packages/tron-wallet-snap/src/services/assets/AssetsRepository.test.ts index 25fb54dcc..a8b886262 100644 --- a/packages/tron-wallet-snap/src/services/assets/AssetsRepository.test.ts +++ b/packages/tron-wallet-snap/src/services/assets/AssetsRepository.test.ts @@ -1,11 +1,12 @@ +import type { IStateManager } from '@metamask/snap-networks-utils'; + import { KnownCaip19Id, Network } from '../../constants'; import type { AssetEntity, NativeAsset, TokenAsset, } from '../../entities/assets'; -import type { IStateManager } from '../state/IStateManager'; -import type { UnencryptedStateValue } from '../state/State'; +import type { UnencryptedStateValue } from '../state/stateTypes'; import { AssetsRepository } from './AssetsRepository'; import type { NativeCaipAssetType, TokenCaipAssetType } from './types'; @@ -109,6 +110,7 @@ describe('AssetsRepository', () => { return stateValue; }, deleteKey: async () => undefined, + deleteKeys: async () => undefined, }; return { diff --git a/packages/tron-wallet-snap/src/services/assets/AssetsRepository.ts b/packages/tron-wallet-snap/src/services/assets/AssetsRepository.ts index d4c480f2c..0c7027e78 100644 --- a/packages/tron-wallet-snap/src/services/assets/AssetsRepository.ts +++ b/packages/tron-wallet-snap/src/services/assets/AssetsRepository.ts @@ -1,8 +1,8 @@ +import type { IStateManager } from '@metamask/snap-networks-utils'; import { cloneDeep } from 'lodash'; import type { AssetEntity } from '../../entities/assets'; -import type { IStateManager } from '../state/IStateManager'; -import type { UnencryptedStateValue } from '../state/State'; +import type { UnencryptedStateValue } from '../state/stateTypes'; export class AssetsRepository { readonly #state: IStateManager; diff --git a/packages/tron-wallet-snap/src/services/assets/adapters/SnapAssetsAdapter.ts b/packages/tron-wallet-snap/src/services/assets/adapters/SnapAssetsAdapter.ts index c9fd16bc8..592ff4c60 100644 --- a/packages/tron-wallet-snap/src/services/assets/adapters/SnapAssetsAdapter.ts +++ b/packages/tron-wallet-snap/src/services/assets/adapters/SnapAssetsAdapter.ts @@ -5,7 +5,7 @@ import type { KeyringAccount, } from '@metamask/keyring-api'; import { emitSnapKeyringEvent } from '@metamask/keyring-snap-sdk'; -import type { Logger } from '@metamask/snap-networks-utils'; +import type { IStateManager, Logger } from '@metamask/snap-networks-utils'; import type { AssetMetadata, FungibleAssetMetadata } from '@metamask/snaps-sdk'; import type { CaipAssetType } from '@metamask/utils'; import { parseCaipAssetType } from '@metamask/utils'; @@ -39,7 +39,7 @@ import { } from '../../../constants'; import type { AssetEntity } from '../../../entities/assets'; import { toUiAmount } from '../../../utils/conversion'; -import type { State, UnencryptedStateValue } from '../../state/State'; +import type { UnencryptedStateValue } from '../../state/stateTypes'; import type { AssetsRepository } from '../AssetsRepository'; import type { InLockPeriodCaipAssetType, @@ -86,7 +86,7 @@ export class SnapAssetsAdapter { readonly #assetsRepository: AssetsRepository; - readonly #state: State; + readonly #state: IStateManager; readonly #trongridApiClient: TrongridApiClient; @@ -110,7 +110,7 @@ export class SnapAssetsAdapter { }: { logger: Logger; assetsRepository: AssetsRepository; - state: State; + state: IStateManager; trongridApiClient: TrongridApiClient; tronHttpClient: TronHttpClient; priceApiClient: PriceApiClient; diff --git a/packages/tron-wallet-snap/src/services/confirmation/ConfirmationHandler.ts b/packages/tron-wallet-snap/src/services/confirmation/ConfirmationHandler.ts index 83f3bbdb5..3be85365a 100644 --- a/packages/tron-wallet-snap/src/services/confirmation/ConfirmationHandler.ts +++ b/packages/tron-wallet-snap/src/services/confirmation/ConfirmationHandler.ts @@ -1,4 +1,4 @@ -import type { Logger } from '@metamask/snap-networks-utils'; +import type { IStateManager, Logger } from '@metamask/snap-networks-utils'; import { InternalError } from '@metamask/snaps-sdk'; import { assert } from '@metamask/superstruct'; import { BigNumber } from 'bignumber.js'; @@ -28,14 +28,14 @@ import { assertTransactionStructure } from '../../validation/transaction'; import type { AssetsService } from '../assets/AssetsService'; import type { FeeCalculatorService } from '../send/FeeCalculatorService'; import type { ComputeFeeResult } from '../send/types'; -import type { State, UnencryptedStateValue } from '../state/State'; +import type { UnencryptedStateValue } from '../state/stateTypes'; export class ConfirmationHandler { readonly #logger: Logger; readonly #snapClient: SnapClient; - readonly #state: State; + readonly #state: IStateManager; readonly #tronWebFactory: TronWebFactory; @@ -52,7 +52,7 @@ export class ConfirmationHandler { logger, }: { snapClient: SnapClient; - state: State; + state: IStateManager; tronWebFactory: TronWebFactory; assetsService: AssetsService; feeCalculatorService: FeeCalculatorService; diff --git a/packages/tron-wallet-snap/src/services/state/IStateManager.ts b/packages/tron-wallet-snap/src/services/state/IStateManager.ts deleted file mode 100644 index 9923ca4b0..000000000 --- a/packages/tron-wallet-snap/src/services/state/IStateManager.ts +++ /dev/null @@ -1,105 +0,0 @@ -import type { Serializable } from '@metamask/snap-networks-utils'; - -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. - * - * @example - * ```typescript - * // state is { users: [ { name: 'Alice', age: 20 }, { name: 'Bob', age: 25 } ], countries: ['Spain', 'France'] } - * - * const value = await stateManager.get(); - * // value is { users: [ { name: 'Alice', age: 20 }, { name: 'Bob', age: 25 } ], countries: ['Spain', 'France'] } - * ``` - */ - get(): Promise; - /** - * Gets the value of passed key in the state object. - * The key is the json path to the value to get. - * - * @example - * ```typescript - * // state is { users: [ { name: 'Alice', age: 20 }, { name: 'Bob', age: 25 } ], countries: ['Spain', 'France'] } - * - * const value = await stateManager.getKey('users.1.name'); - * // value is 'Bob' - * - * @returns The value of the key, or undefined if the key does not exist. - */ - getKey( - key: string, - ): Promise; - /** - * Sets the value of passed key in the state object. - * The key is a json path to the value to set. - * - * @example - * ```typescript - * const state = await stateManager.get(); - * // state is { users: [ { name: 'Alice', age: 20 }, { name: 'Bob', age: 25 } ] } - * - * await stateManager.set('users.1.name', 'John'); - * // state is now { users: [ { name: 'Alice', age: 20 }, { name: 'John', age: 25 } ] } - * ``` - * @param key - The key to set, which is a json path to the location. - * @param value - The value to set. - */ - // TODO: Replace `any` with type - // eslint-disable-next-line @typescript-eslint/no-explicit-any - setKey(key: string, value: any): 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. If any part of the operation fails, none of the changes will be applied. This prevents partial updates that could leave the underlying data store in an inconsistent state. - * - Performance: Making multiple individual `state.set` or `state.delete` calls would require multiple round trips to the state storage system, causing potential overheads. - * - State Consistency: Maintains better state consistency by reading the state once, making all modifications in memory and writing the complete updated state back. - * - * ⚠️ WARNING: Use with caution because: - * - it will override the whole state. - * - it transfers the whole state back and forth the data store, which might consume a lot of bandwidth. - * - * For single updates, use instead `setKey` or `deleteKey`. - * - * @param updaterFunction - The function that updates the state. - * @returns The updated state. - */ - update( - updaterFunction: (state: TStateValue) => TStateValue, - ): Promise; - /** - * Deletes the value of passed key in the state object. - * The key is a json path to the value to delete. - * - * @example - * ```typescript - * const state = await stateManager.get(); - * // state is { users: [ { name: 'Alice', age: 20 }, { name: 'Bob', age: 25 } ] } - * - * await stateManager.deleteKey('users.1'); - * // state is now { users: [ { name: 'Alice', age: 20 } ] } - * ``` - */ - deleteKey(key: string): Promise; -}; diff --git a/packages/tron-wallet-snap/src/services/state/InMemoryState.ts b/packages/tron-wallet-snap/src/services/state/InMemoryState.ts deleted file mode 100644 index 8fd4e1991..000000000 --- a/packages/tron-wallet-snap/src/services/state/InMemoryState.ts +++ /dev/null @@ -1,56 +0,0 @@ -import type { Serializable } from '@metamask/snap-networks-utils'; -import { get, set, unset } from 'lodash'; - -import type { IStateManager } from './IStateManager'; - -/** - * A simple implementation of the `IStateManager` interface that relies on an in memory state that can be used for testing purposes. - */ -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 { - const value = get(this.#state, key); - - return value as TResponse | undefined; - } - - async setKey(key: string, value: Serializable): Promise { - set(this.#state, key, value); // Use lodash to set the value using a json path - } - - async setKeyWith( - key: string, - updater: (currentValue: TValue | undefined) => TValue, - ): Promise { - const oldValue = get(this.#state, key) as TValue | undefined; - const newValue = updater(oldValue); - - set(this.#state, key, newValue); - } - - async update( - callback: (state: TStateValue) => TStateValue, - ): Promise { - this.#state = callback(this.#state); - - return this.#state; - } - - async deleteKey(key: string): Promise { - // Using lodash's unset to leverage the json path capabilities - unset(this.#state, key); - } -} diff --git a/packages/tron-wallet-snap/src/services/state/State.test.ts b/packages/tron-wallet-snap/src/services/state/State.test.ts deleted file mode 100644 index 76147538e..000000000 --- a/packages/tron-wallet-snap/src/services/state/State.test.ts +++ /dev/null @@ -1,504 +0,0 @@ -/* eslint-disable jest/prefer-strict-equal */ - -import { BigNumber } from 'bignumber.js'; - -import { State } from './State'; - -const snap = { - request: jest.fn(), -}; - -(globalThis as any).snap = snap; - -const flushPromises = async () => { - 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 () => { - const mockUnderlyingState = {}; - snap.request.mockResolvedValue(mockUnderlyingState); - - const stateValue = await state.get(); - - expect(stateValue).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); - }); - - describe('when getting serialized non-JSON values', () => { - it('deserializes undefined values', async () => { - const mockUnderlyingState = { - users: [ - { - name: 'John', - age: { - __type: 'undefined', - }, - }, - ], - }; - snap.request.mockResolvedValue(mockUnderlyingState); - - const stateValue = await state.get(); - - expect(stateValue).toEqual({ - users: [ - { - name: 'John', - age: undefined, - }, - ], - }); - }); - - it('deserializes BigNumber values', async () => { - const mockUnderlyingState = { - users: [ - { - name: 'John', - age: { - __type: 'BigNumber', - value: '30', - }, - }, - ], - }; - snap.request.mockResolvedValue(mockUnderlyingState); - - const stateValue = await state.get(); - - expect(stateValue).toStrictEqual({ - users: [ - { - name: 'John', - age: new BigNumber(30), - }, - ], - }); - }); - - it('deserializes bigint values', async () => { - const mockUnderlyingState = { - users: [ - { - name: 'John', - age: { - __type: 'bigint', - value: '30', - }, - }, - ], - }; - snap.request.mockResolvedValue(mockUnderlyingState); - - const stateValue = await state.get(); - - expect(stateValue).toStrictEqual({ - users: [ - { - name: 'John', - age: BigInt(30), - }, - ], - }); - }); - }); - }); - - describe('getKey', () => { - it('calls the snap_getState method with the correct parameters', async () => { - const mockUnderlyingState = DEFAULT_STATE; - snap.request.mockResolvedValue(mockUnderlyingState); - - 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); - - const value = await state.getKey('users.1.name'); - - expect(value).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 }); // getState (read) - - 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); // getState returns null → key absent - - 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('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('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, - }, - }); - }); - }); -}); diff --git a/packages/tron-wallet-snap/src/services/state/State.ts b/packages/tron-wallet-snap/src/services/state/State.ts deleted file mode 100644 index 15c941894..000000000 --- a/packages/tron-wallet-snap/src/services/state/State.ts +++ /dev/null @@ -1,263 +0,0 @@ -import type { Transaction } from '@metamask/keyring-api'; -import { - deserialize, - safeMerge, - serialize, -} from '@metamask/snap-networks-utils'; -import type { Serializable } from '@metamask/snap-networks-utils'; -import type { Json } from '@metamask/snaps-sdk'; -import type { MutexInterface } from 'async-mutex'; -import { Mutex } from 'async-mutex'; -import { unset } from 'lodash'; - -import type { AssetEntity } from '../../entities/assets'; -import type { TronKeyringAccount } from '../../entities/keyring-account'; -import type { IStateManager } from './IStateManager'; - -export type AccountId = string; - -export type UnencryptedStateValue = { - keyringAccounts: Record; - assets: Record; - transactions: Record; - mapInterfaceNameToId: Record; -}; - -export const DEFAULT_UNENCRYPTED_STATE: UnencryptedStateValue = { - keyringAccounts: {}, - assets: {}, - transactions: {}, - mapInterfaceNameToId: {}, -}; - -export type StateConfig> = { - encrypted: boolean; - defaultState: TValue; -}; - -/** - * 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. - * snap_setState, snap_getState etc does not have this limitation and can be accessed safely as long as - * an ongoing manageState operation is not occurring. - */ -class StateLock { - readonly #blobModificationMutex = new Mutex(); - - 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 { - // If we are currently doing a full blob update, wait it out. - // Signal that regular state operations are ongoing by acquring the mutex. - // Other regular state operations can skip this, as they are safe to do in parallel. - await Promise.all([ - this.#blobModificationMutex.waitForUnlock(), - this.#acquireRegularStateUpdateMutex(), - ]); - - try { - this.#pendingRegularStateUpdates += 1; - return await callback(); - } finally { - this.#pendingRegularStateUpdates -= 1; - - if ( - this.#pendingRegularStateUpdates === 0 && - this.#releaseRegularStateUpdateMutex - ) { - this.#releaseRegularStateUpdateMutex(); - } - } - } - - async wrapRegularStateWriteOperation( - callback: MutexInterface.Worker, - ): Promise { - return await this.#regularStateWriteMutex.runExclusive(async () => - this.wrapRegularStateOperation(callback), - ); - } - - async wrapManageStateOperation( - callback: MutexInterface.Worker, - ): Promise { - await this.#regularStateUpdateMutex.waitForUnlock(); - - return await this.#blobModificationMutex.runExclusive(callback); - } -} - -/** - * This class is a layer on top the the `snap_manageState` API that facilitates its usage: - * - * Basic usage: - * - Get and update the sate of the snap - * - * Serialization: - * - It serializes the data before storing it in the snap state because only JSON-assignable data can be stored. - * - It deserializes the data after retrieving it from the snap state. - * - So you don't need to worry about the data format when storing or retrieving data. - * - * Default values: - * - It merges the default state with the underlying snap state to ensure that we always have default values, - * letting us avoid a ton of null checks everywhere. - */ -export class State< - TStateValue extends Record, -> implements IStateManager { - readonly #lock = new StateLock(); - - readonly #config: StateConfig; - - constructor(config: StateConfig) { - this.#config = config; - } - - async #unsafeGet(): Promise { - const state = await snap.request({ - method: 'snap_getState', - params: { - encrypted: this.#config.encrypted, - }, - }); - - const stateDeserialized = deserialize(state ?? {}) as TStateValue; - - // Merge the default state with the underlying snap state - // to ensure that we always have default values. It lets us avoid a ton of null checks everywhere. - const stateWithDefaults = safeMerge( - this.#config.defaultState, - stateDeserialized, - ); - - return stateWithDefaults; - } - - async get(): Promise { - return this.#lock.wrapRegularStateOperation(async () => this.#unsafeGet()); - } - - async getKey( - key: string, - ): Promise { - return this.#lock.wrapRegularStateOperation(async () => { - const value = await snap.request({ - method: 'snap_getState', - params: { - key, - encrypted: this.#config.encrypted, - }, - }); - - if (value === null) { - return undefined; - } - - return deserialize(value) as TResponse; - }); - } - - async setKey(key: string, value: Serializable): Promise { - await this.#lock.wrapRegularStateWriteOperation(async () => { - const serializedValue = serialize(value); - - await snap.request({ - method: 'snap_setState', - params: { - key, - value: serializedValue, - encrypted: this.#config.encrypted, - }, - }); - }); - } - - async setKeyWith( - key: string, - updater: (currentValue: TValue | undefined) => TValue, - ): Promise { - await this.#lock.wrapRegularStateWriteOperation(async () => { - const rawValue = await snap.request({ - method: 'snap_getState', - params: { - key, - encrypted: this.#config.encrypted, - }, - }); - - const oldValue = - rawValue === null ? undefined : (deserialize(rawValue) as TValue); - - const newValue = updater(oldValue); - - const serializedValue = serialize(newValue); - - await snap.request({ - method: 'snap_setState', - params: { - key, - value: serializedValue, - encrypted: this.#config.encrypted, - }, - }); - }); - } - - async update( - updaterFunction: (state: TStateValue) => TStateValue, - ): Promise { - // Because this function modifies the entire state blob, - // we must protect against parallel requests. - return await this.#lock.wrapManageStateOperation(async () => { - const currentState = await this.#unsafeGet(); - - const newState = updaterFunction(currentState); - - // Generally we should try to use snap_getState and snap_setState over this - // as snap_manageState is slower and error-prone due to requiring manual mutex management. - await snap.request({ - method: 'snap_manageState', - params: { - operation: 'update', - // State values are always objects, so the serialized result is too. - newState: serialize(newState) as Record, - encrypted: this.#config.encrypted, - }, - }); - - return newState; - }); - } - - async deleteKey(key: string): Promise { - await this.update((state) => { - // Using lodash's unset to leverage the json path capabilities - unset(state, key); - return state; - }); - } - - async deleteKeys(keys: string[]): Promise { - await this.update((state) => { - keys.forEach((key) => { - unset(state, key); - }); - return state; - }); - } -} diff --git a/packages/tron-wallet-snap/src/services/state/stateTypes.ts b/packages/tron-wallet-snap/src/services/state/stateTypes.ts new file mode 100644 index 000000000..fb84d7244 --- /dev/null +++ b/packages/tron-wallet-snap/src/services/state/stateTypes.ts @@ -0,0 +1,18 @@ +import type { Transaction } from '@metamask/keyring-api'; + +import type { AssetEntity } from '../../entities/assets'; +import type { TronKeyringAccount } from '../../entities/keyring-account'; + +export type UnencryptedStateValue = { + keyringAccounts: Record; + assets: Record; + transactions: Record; + mapInterfaceNameToId: Record; +}; + +export const DEFAULT_UNENCRYPTED_STATE: UnencryptedStateValue = { + keyringAccounts: {}, + assets: {}, + transactions: {}, + mapInterfaceNameToId: {}, +}; diff --git a/packages/tron-wallet-snap/src/services/transactions/TransactionsRepository.test.ts b/packages/tron-wallet-snap/src/services/transactions/TransactionsRepository.test.ts index 33e0d3e09..ead83c006 100644 --- a/packages/tron-wallet-snap/src/services/transactions/TransactionsRepository.test.ts +++ b/packages/tron-wallet-snap/src/services/transactions/TransactionsRepository.test.ts @@ -1,13 +1,14 @@ import type { Transaction } from '@metamask/keyring-api'; import { TransactionStatus, TransactionType } from '@metamask/keyring-api'; +import type { IStateManager } from '@metamask/snap-networks-utils'; import { KnownCaip19Id, Network } from '../../constants'; -import type { State, UnencryptedStateValue } from '../state/State'; +import type { UnencryptedStateValue } from '../state/stateTypes'; import { TransactionsRepository } from './TransactionsRepository'; describe('TransactionsRepository', () => { let transactionsRepository: TransactionsRepository; - let mockState: jest.Mocked>; + let mockState: jest.Mocked>; const mockAccountId = 'test-account-id'; @@ -53,7 +54,7 @@ describe('TransactionsRepository', () => { setKey: jest.fn(), setKeyWith: jest.fn(), update: jest.fn(), - } as unknown as jest.Mocked>; + } as unknown as jest.Mocked>; transactionsRepository = new TransactionsRepository(mockState); }); diff --git a/packages/tron-wallet-snap/src/services/transactions/TransactionsRepository.ts b/packages/tron-wallet-snap/src/services/transactions/TransactionsRepository.ts index aec81f62c..214879fb6 100644 --- a/packages/tron-wallet-snap/src/services/transactions/TransactionsRepository.ts +++ b/packages/tron-wallet-snap/src/services/transactions/TransactionsRepository.ts @@ -1,15 +1,16 @@ import type { Transaction } from '@metamask/keyring-api'; import { TransactionStatus } from '@metamask/keyring-api'; +import type { IStateManager } from '@metamask/snap-networks-utils'; import { chain } from 'lodash'; -import type { State, UnencryptedStateValue } from '../state/State'; +import type { UnencryptedStateValue } from '../state/stateTypes'; export class TransactionsRepository { - readonly #state: State; + readonly #state: IStateManager; readonly #stateKey = 'transactions'; - constructor(state: State) { + constructor(state: IStateManager) { this.#state = state; } diff --git a/packages/tron-wallet-snap/src/ui/confirmation/views/ConfirmTransactionRequest/render.test.tsx b/packages/tron-wallet-snap/src/ui/confirmation/views/ConfirmTransactionRequest/render.test.tsx index 51f113904..462d78e11 100644 --- a/packages/tron-wallet-snap/src/ui/confirmation/views/ConfirmTransactionRequest/render.test.tsx +++ b/packages/tron-wallet-snap/src/ui/confirmation/views/ConfirmTransactionRequest/render.test.tsx @@ -1,3 +1,4 @@ +import type { IStateManager } from '@metamask/snap-networks-utils'; import { Types as TronwebTypes } from 'tronweb'; import { @@ -8,10 +9,7 @@ import type { SnapClient } from '../../../../clients/snap/SnapClient'; import { Network } from '../../../../constants'; import type { AssetEntity } from '../../../../entities/assets'; import { BackgroundEventMethod } from '../../../../handlers/cronjob/cronjob'; -import type { - State, - UnencryptedStateValue, -} from '../../../../services/state/State'; +import type { UnencryptedStateValue } from '../../../../services/state/stateTypes'; import type { TransactionScanService } from '../../../../services/transaction-scan/TransactionScanService'; import { SimulationStatus } from '../../../../services/transaction-scan/types'; import type { TransactionScanResult } from '../../../../services/transaction-scan/types'; @@ -231,7 +229,7 @@ async function withRender( ) => render( mockSnapClient as unknown as SnapClient, - mockState as unknown as State, + mockState as unknown as IStateManager, { ...defaultIncomingContext, ...contextOverrides }, ); diff --git a/packages/tron-wallet-snap/src/ui/confirmation/views/ConfirmTransactionRequest/render.tsx b/packages/tron-wallet-snap/src/ui/confirmation/views/ConfirmTransactionRequest/render.tsx index c2a53bd04..73571cedf 100644 --- a/packages/tron-wallet-snap/src/ui/confirmation/views/ConfirmTransactionRequest/render.tsx +++ b/packages/tron-wallet-snap/src/ui/confirmation/views/ConfirmTransactionRequest/render.tsx @@ -1,3 +1,4 @@ +import type { IStateManager } from '@metamask/snap-networks-utils'; import type { DialogResult, Json } from '@metamask/snaps-sdk'; import type { Types as TronwebTypes } from 'tronweb'; @@ -8,10 +9,7 @@ import type { AssetEntity } from '../../../../entities/assets'; import type { TronKeyringAccount } from '../../../../entities/keyring-account'; import { BackgroundEventMethod } from '../../../../handlers/cronjob/cronjob'; import type { ComputeFeeResult } from '../../../../services/send/types'; -import type { - State, - UnencryptedStateValue, -} from '../../../../services/state/State'; +import type { UnencryptedStateValue } from '../../../../services/state/stateTypes'; import { TRX_IMAGE_SVG } from '../../../../static/tron-logo'; import { FetchStatus } from '../../../../types/snap'; import { getIconUrlForKnownAsset } from '../../utils/getIconUrlForKnownAsset'; @@ -76,7 +74,7 @@ export const DEFAULT_CONFIRMATION_CONTEXT: ConfirmTransactionRequestContext = { */ export async function render( snapClient: SnapClient, - state: State, + state: IStateManager, incomingContext: { scope: Network; fromAddress: string; diff --git a/yarn.lock b/yarn.lock index 545e9e1fd..77cf460e0 100644 --- a/yarn.lock +++ b/yarn.lock @@ -3839,7 +3839,6 @@ __metadata: "@metamask/superstruct": "npm:^3.4.1" "@metamask/utils": "npm:^11.11.0" "@types/lodash": "npm:^4.17.15" - async-mutex: "npm:^0.5.0" bignumber.js: "npm:^9.3.1" concurrently: "npm:^10.0.3" dotenv: "npm:^17.2.1"