Add BaseControllerV2 state metadata

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';
 
 type MockControllerState = {
   count: number;
 };
 
+const mockControllerStateMetadata = {
+  count: {
+    persist: true,
+    anonymous: true,
+  },
+};
+
 class MockController extends BaseController<MockControllerState> {
   update(callback: (state: Draft<MockControllerState>) => void | MockControllerState) {
     super.update(callback);
@@ -19,21 +26,27 @@ class MockController extends BaseController<MockControllerState> {
 
 describe('BaseController', () => {
   it('should set initial state', () => {
-    const controller = new MockController({ count: 0 });
+    const controller = new MockController({ count: 0 }, mockControllerStateMetadata);
 
     expect(controller.state).toEqual({ count: 0 });
   });
 
+  it('should set initial schema', () => {
+    const controller = new MockController({ count: 0 }, mockControllerStateMetadata);
+
+    expect(controller.metadata).toEqual(mockControllerStateMetadata);
+  });
+
   it('should not allow mutating state directly', () => {
-    const controller = new MockController({ count: 0 });
+    const controller = new MockController({ count: 0 }, mockControllerStateMetadata);
 
     expect(() => {
       controller.state = { count: 1 };
     }).toThrow();
   });
 
   it('should allow updating state by modifying draft', () => {
-    const controller = new MockController({ count: 0 });
+    const controller = new MockController({ count: 0 }, mockControllerStateMetadata);
 
     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 }, mockControllerStateMetadata);
 
     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 }, mockControllerStateMetadata);
 
     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 }, mockControllerStateMetadata);
     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 }, mockControllerStateMetadata);
     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 }, mockControllerStateMetadata);
     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 }, mockControllerStateMetadata);
     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 }, mockControllerStateMetadata);
     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 }, mockControllerStateMetadata);
     const listener1 = sinon.stub();
     const listener2 = sinon.stub();
 
@@ -146,3 +159,166 @@ 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' });
+  });
+
+  it('should allow returning a partial object from an anonymizing function', () => {
+    const anonymizeTransactionHash = (txMeta: { hash: string; value: number }) => {
+      return { value: txMeta.value };
+    };
+
+    const anonymizedState = getAnonymizedState(
+      {
+        txMeta: {
+          hash: '0x123',
+          value: 10,
+        },
+      },
+      {
+        txMeta: {
+          anonymous: anonymizeTransactionHash,
+          persist: false,
+        },
+      },
+    );
+
+    expect(anonymizedState).toEqual({ txMeta: { value: 10 } });
+  });
+
+  it('should allow returning a nested partial object from an anonymizing function', () => {
+    const anonymizeTransactionHash = (txMeta: {
+      hash: string;
+      value: number;
+      history: { hash: string; value: number }[];
+    }) => {
+      return {
+        history: txMeta.history.map((entry) => {
+          return { value: entry.value };
+        }),
+        value: txMeta.value,
+      };
+    };
+
+    const anonymizedState = getAnonymizedState(
+      {
+        txMeta: {
+          hash: '0x123',
+          history: [
+            {
+              hash: '0x123',
+              value: 9,
+            },
+          ],
+          value: 10,
+        },
+      },
+      {
+        txMeta: {
+          anonymous: anonymizeTransactionHash,
+          persist: false,
+        },
+      },
+    );
+
+    expect(anonymizedState).toEqual({ txMeta: { history: [{ value: 9 }], value: 10 } });
+  });
+});
+
+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' });
+  });
+});

src/BaseControllerV2.ts

@@ -7,25 +7,87 @@ import type { Draft, Patch } from 'immer';
 enablePatches();
 
 /**
- * State change callbacks
+ * A state change listener.
+ *
+ * This function will get called for each state change, and is given a copy of
+ * the new state along with a set of patches describing the changes since the
+ * last update.
+ *
+ * @param state - The new controller state
+ * @param patches - A list of patches describing any changes (see here for more
+ *   information: https://immerjs.github.io/immer/docs/patches)
  */
 export type Listener<T> = (state: T, patches: Patch[]) => void;
 
+type Primitive = boolean | string | number | null;
+
+// Based upon this StackOverflow answer: https://stackoverflow.com/a/64060332
+type RecursivePartial<T> = {
+  [P in keyof T]?: T[P] extends (infer U)[]
+    ? RecursivePartial<U>[]
+    : T[P] extends Primitive
+    ? T[P]
+    : RecursivePartial<T[P]>;
+};
+
 /**
- * Controller class that provides state management and subscriptions
+ * An anonymizing function
+ *
+ * This function will accept one piece of the controller state (one property),
+ * and will return an anonymized representation of this state. By "anonymized",
+ * we mean that it should not contain any information that could be personally
+ * identifiable.
+ *
+ * @param value - A piece of controller state
+ * @returns An anonymized representation of the given state
+ */
+export type Anonymizer<T> = (value: T) => T extends Primitive ? T : RecursivePartial<T>;
+
+/**
+ * State metadata.
+ *
+ * This metadata describes which parts of state should be persisted, and how to
+ * get an anonymized representation of the state.
+ */
+export type StateMetadata<T> = {
+  [P in keyof T]: StatePropertyMetadata<T[P]>;
+};
+
+/**
+ * Metadata for a single state property
+ *
+ * @property persist - Indicates whether this property should be persisted
+ *   (`true` for persistent, `false` for transient)
+ * @property anonymous - Indicates whether this property is already anonymous,
+ *   (`true` for anonymous, `false` if it has potential to be personally
+ *   identifiable), or is set to a function that returns an anonymized
+ *   representation of this state.
+ */
+export interface StatePropertyMetadata<T> {
+  persist: boolean;
+  anonymous: boolean | Anonymizer<T>;
+}
+
+/**
+ * Controller class that provides state management, subscriptions, and state metadata
  */
 export class BaseController<S extends Record<string, unknown>> {
   private internalState: S;
 
   private internalListeners: Set<Listener<S>> = new Set();
 
+  public readonly metadata: StateMetadata<S>;
+
   /**
    * Creates a BaseController instance.
    *
    * @param state - Initial controller state
+   * @param metadata - State metadata, describing how to "anonymize" the state,
+   *   and which parts should be persisted.
    */
-  constructor(state: S) {
+  constructor(state: S, metadata: StateMetadata<S>) {
     this.internalState = state;
+    this.metadata = metadata;
   }
 
   /**
@@ -89,3 +151,52 @@ export class BaseController<S extends Record<string, unknown>> {
     this.internalListeners.clear();
   }
 }
+
+// This function acts as a type guard. Using a `typeof` conditional didn't seem to work.
+function isAnonymizingFunction<T>(x: boolean | Anonymizer<T>): x is Anonymizer<T> {
+  return typeof x === 'function';
+}
+
+/**
+ * Returns an anonymized representation of the controller state.
+ *
+ * By "anonymized" we mean that it should not contain any information that could be personally
+ * identifiable.
+ *
+ * @param state - The controller state
+ * @param metadata - The controller state metadata, which describes how to derive the
+ *   anonymized state
+ * @returns The anonymized controller state
+ */
+export function getAnonymizedState<S extends Record<string, any>>(
+  state: S,
+  metadata: StateMetadata<S>,
+): RecursivePartial<S> {
+  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 metadataValue = metadata[key].anonymous;
+    if (isAnonymizingFunction(metadataValue)) {
+      anonymizedState[key] = metadataValue(state[key]);
+    } else if (metadataValue) {
+      anonymizedState[key] = state[key];
+    }
+    return anonymizedState;
+  }, {} as RecursivePartial<S>);
+}
+
+/**
+ * Returns the subset of state that should be persisted
+ *
+ * @param state - The controller state
+ * @param metadata - The controller state metadata, which describes which pieces of state should be persisted
+ * @returns The subset of controller state that should be persisted
+ */
+export function getPersistentState<S extends Record<string, any>>(state: S, metadata: StateMetadata<S>): Partial<S> {
+  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 (metadata[key].persist) {
+      persistedState[key] = state[key];
+    }
+    return persistedState;
+  }, {} as Partial<S>);
+}