diff --git a/src/BaseControllerV2.test.ts b/src/BaseControllerV2.test.ts index 96dc07197c5..262e53a48ca 100644 --- a/src/BaseControllerV2.test.ts +++ b/src/BaseControllerV2.test.ts @@ -1,12 +1,19 @@ import type { Draft } from 'immer'; import * as sinon from 'sinon'; -import { BaseController } from './BaseControllerV2'; +import { BaseController, getAnonymizedState, getPersistentState } from './BaseControllerV2'; interface MockControllerState { count: number; } +const mockControllerSchema = { + count: { + persist: true, + anonymous: true, + }, +}; + class MockController extends BaseController { update(callback: (state: Draft) => void | MockControllerState) { super.update(callback); @@ -19,13 +26,19 @@ class MockController extends BaseController { describe('BaseController', () => { it('should set initial state', () => { - const controller = new MockController({ count: 0 }); + const controller = new MockController({ count: 0 }, mockControllerSchema); expect(controller.state).toEqual({ count: 0 }); }); + it('should set initial schema', () => { + const controller = new MockController({ count: 0 }, mockControllerSchema); + + expect(controller.schema).toEqual(mockControllerSchema); + }); + it('should not allow mutating state directly', () => { - const controller = new MockController({ count: 0 }); + const controller = new MockController({ count: 0 }, mockControllerSchema); expect(() => { controller.state = { count: 1 }; @@ -33,7 +46,7 @@ describe('BaseController', () => { }); it('should allow updating state by modifying draft', () => { - const controller = new MockController({ count: 0 }); + const controller = new MockController({ count: 0 }, mockControllerSchema); controller.update((draft) => { draft.count += 1; @@ -43,7 +56,7 @@ describe('BaseController', () => { }); it('should allow updating state by return a value', () => { - const controller = new MockController({ count: 0 }); + const controller = new MockController({ count: 0 }, mockControllerSchema); controller.update(() => { return { count: 1 }; @@ -53,7 +66,7 @@ describe('BaseController', () => { }); it('should throw an error if update callback modifies draft and returns value', () => { - const controller = new MockController({ count: 0 }); + const controller = new MockController({ count: 0 }, mockControllerSchema); expect(() => { controller.update((draft) => { @@ -64,7 +77,7 @@ describe('BaseController', () => { }); it('should inform subscribers of state changes', () => { - const controller = new MockController({ count: 0 }); + const controller = new MockController({ count: 0 }, mockControllerSchema); const listener1 = sinon.stub(); const listener2 = sinon.stub(); @@ -81,7 +94,7 @@ describe('BaseController', () => { }); it('should inform a subscriber of each state change once even after multiple subscriptions', () => { - const controller = new MockController({ count: 0 }); + const controller = new MockController({ count: 0 }, mockControllerSchema); const listener1 = sinon.stub(); controller.subscribe(listener1); @@ -95,7 +108,7 @@ describe('BaseController', () => { }); it('should no longer inform a subscriber about state changes after unsubscribing', () => { - const controller = new MockController({ count: 0 }); + const controller = new MockController({ count: 0 }, mockControllerSchema); const listener1 = sinon.stub(); controller.subscribe(listener1); @@ -108,7 +121,7 @@ describe('BaseController', () => { }); it('should no longer inform a subscriber about state changes after unsubscribing once, even if they subscribed many times', () => { - const controller = new MockController({ count: 0 }); + const controller = new MockController({ count: 0 }, mockControllerSchema); const listener1 = sinon.stub(); controller.subscribe(listener1); @@ -122,7 +135,7 @@ describe('BaseController', () => { }); it('should allow unsubscribing listeners who were never subscribed', () => { - const controller = new MockController({ count: 0 }); + const controller = new MockController({ count: 0 }, mockControllerSchema); const listener1 = sinon.stub(); expect(() => { @@ -131,7 +144,7 @@ describe('BaseController', () => { }); it('should no longer update subscribers after being destroyed', () => { - const controller = new MockController({ count: 0 }); + const controller = new MockController({ count: 0 }, mockControllerSchema); const listener1 = sinon.stub(); const listener2 = sinon.stub(); @@ -146,3 +159,105 @@ describe('BaseController', () => { expect(listener2.callCount).toEqual(0); }); }); + +describe('getAnonymizedState', () => { + it('should return empty state', () => { + expect(getAnonymizedState({}, {})).toEqual({}); + }); + + it('should return empty state when no properties are anonymized', () => { + const anonymizedState = getAnonymizedState({ count: 1 }, { count: { anonymous: false, persist: false } }); + expect(anonymizedState).toEqual({}); + }); + + it('should return state that is already anonymized', () => { + const anonymizedState = getAnonymizedState( + { + password: 'secret password', + privateKey: '123', + network: 'mainnet', + tokens: ['DAI', 'USDC'], + }, + { + password: { + anonymous: false, + persist: false, + }, + privateKey: { + anonymous: false, + persist: false, + }, + network: { + anonymous: true, + persist: false, + }, + tokens: { + anonymous: true, + persist: false, + }, + }, + ); + expect(anonymizedState).toEqual({ network: 'mainnet', tokens: ['DAI', 'USDC'] }); + }); + + it('should use anonymizing function to anonymize state', () => { + const anonymizeTransactionHash = (hash: string) => { + return hash.split('').reverse().join(''); + }; + + const anonymizedState = getAnonymizedState( + { + transactionHash: '0x1234', + }, + { + transactionHash: { + anonymous: anonymizeTransactionHash, + persist: false, + }, + }, + ); + + expect(anonymizedState).toEqual({ transactionHash: '4321x0' }); + }); +}); + +describe('getPersistentState', () => { + it('should return empty state', () => { + expect(getPersistentState({}, {})).toEqual({}); + }); + + it('should return empty state when no properties are persistent', () => { + const persistentState = getPersistentState({ count: 1 }, { count: { anonymous: false, persist: false } }); + expect(persistentState).toEqual({}); + }); + + it('should return persistent state', () => { + const persistentState = getPersistentState( + { + password: 'secret password', + privateKey: '123', + network: 'mainnet', + tokens: ['DAI', 'USDC'], + }, + { + password: { + anonymous: false, + persist: true, + }, + privateKey: { + anonymous: false, + persist: true, + }, + network: { + anonymous: false, + persist: false, + }, + tokens: { + anonymous: false, + persist: false, + }, + }, + ); + expect(persistentState).toEqual({ password: 'secret password', privateKey: '123' }); + }); +}); diff --git a/src/BaseControllerV2.ts b/src/BaseControllerV2.ts index d87507cfa34..b386bcdca4a 100644 --- a/src/BaseControllerV2.ts +++ b/src/BaseControllerV2.ts @@ -11,6 +11,15 @@ enablePatches(); */ export type Listener = (state: T, patches: Patch[]) => void; +export type Anonymizer = (value: T) => T; + +export type Schema = { + [P in keyof T]: { + persist: boolean; + anonymous: boolean | Anonymizer; + }; +}; + /** * Controller class that provides state management and subscriptions */ @@ -19,13 +28,18 @@ export class BaseController> { private internalListeners: Set> = new Set(); + public readonly schema: Schema; + /** * Creates a BaseController instance. * * @param state - Initial controller state + * @param schema - State schema, describing how to "anonymize" the state, + * and which parts should be persisted. */ - constructor(state: S) { + constructor(state: S, schema: Schema) { this.internalState = state; + this.schema = schema; } /** @@ -89,3 +103,31 @@ export class BaseController> { this.internalListeners.clear(); } } + +// This function acts as a type guard. Using a `typeof` conditional didn't seem to work. +function isAnonymizingFunction(x: boolean | Anonymizer): x is Anonymizer { + return typeof x === 'function'; +} + +export function getAnonymizedState>(state: S, schema: Schema) { + return Object.keys(state).reduce((anonymizedState, _key) => { + const key: keyof S = _key; // https://stackoverflow.com/questions/63893394/string-cannot-be-used-to-index-type-t + const schemaValue = schema[key].anonymous; + if (isAnonymizingFunction(schemaValue)) { + anonymizedState[key] = schemaValue(state[key]); + } else if (schemaValue) { + anonymizedState[key] = state[key]; + } + return anonymizedState; + }, {} as Partial); +} + +export function getPersistentState>(state: S, schema: Schema) { + return Object.keys(state).reduce((persistedState, _key) => { + const key: keyof S = _key; // https://stackoverflow.com/questions/63893394/string-cannot-be-used-to-index-type-t + if (schema[key].persist) { + persistedState[key] = state[key]; + } + return persistedState; + }, {} as Partial); +}