Add types for store and fix type discrepancies (#247)

* fix: add types for persistent and memory store

* fix: lint

* refactor: remove return from void function

* refactor: edit keyringBuilders type

* fix: make keyringBuilders arg optional

* rollback setLocked change

* Update src/KeyringController.test.ts

Co-authored-by: Elliot Winkler <[email protected]>

* refactor: rollback unlockKeyrings

---------

Co-authored-by: Elliot Winkler <[email protected]>

Author: mikesposito

package.json

@@ -44,8 +44,8 @@
     "@metamask/eth-hd-keyring": "^6.0.0",
     "@metamask/eth-sig-util": "^6.0.0",
     "@metamask/eth-simple-keyring": "^5.0.0",
-    "@metamask/utils": "^6.2.0",
-    "obs-store": "^4.0.3"
+    "@metamask/obs-store": "^8.1.0",
+    "@metamask/utils": "^6.2.0"
   },
   "devDependencies": {
     "@lavamoat/allow-scripts": "^2.3.1",

src/KeyringController.test.ts

@@ -122,6 +122,7 @@ describe('KeyringController', () => {
       await keyringController.persistAllKeyrings();
 
       const { vault } = keyringController.store.getState();
+      assert(vault, 'Vault is not set');
       const keyrings = await mockEncryptor.decrypt(PASSWORD, vault);
       expect(keyrings).toContain(unsupportedKeyring);
       expect(keyrings).toHaveLength(2);
@@ -136,7 +137,9 @@ describe('KeyringController', () => {
         const vault = JSON.stringify({ salt: vaultEncryptionSalt });
         keyringController.store.updateState({ vault });
 
-        expect(keyringController.memStore.getState().encryptionKey).toBeNull();
+        expect(
+          keyringController.memStore.getState().encryptionKey,
+        ).toBeUndefined();
         expect(
           keyringController.memStore.getState().encryptionSalt,
         ).toBeUndefined();
@@ -201,7 +204,7 @@ describe('KeyringController', () => {
 
   describe('createNewVaultAndKeychain', () => {
     it('should create a new vault', async () => {
-      keyringController.store.updateState({ vault: null });
+      keyringController.store.putState({});
       assert(!keyringController.store.getState().vault, 'no previous vault');
 
       const newVault = await keyringController.createNewVaultAndKeychain(
@@ -213,7 +216,7 @@ describe('KeyringController', () => {
     });
 
     it('should unlock the vault', async () => {
-      keyringController.store.updateState({ vault: null });
+      keyringController.store.putState({});
       assert(!keyringController.store.getState().vault, 'no previous vault');
 
       await keyringController.createNewVaultAndKeychain(PASSWORD);
@@ -222,7 +225,7 @@ describe('KeyringController', () => {
     });
 
     it('should encrypt keyrings with the correct password each time they are persisted', async () => {
-      keyringController.store.updateState({ vault: null });
+      keyringController.store.putState({});
       assert(!keyringController.store.getState().vault, 'no previous vault');
 
       await keyringController.createNewVaultAndKeychain(PASSWORD);
@@ -254,7 +257,7 @@ describe('KeyringController', () => {
         await keyringController.createNewVaultAndKeychain(PASSWORD);
         const finalMemStore = keyringController.memStore.getState();
 
-        expect(initialMemStore.encryptionKey).toBeNull();
+        expect(initialMemStore.encryptionKey).toBeUndefined();
         expect(initialMemStore.encryptionSalt).toBeUndefined();
 
         expect(finalMemStore.encryptionKey).toBe(MOCK_HARDCODED_KEY);
@@ -347,7 +350,7 @@ describe('KeyringController', () => {
         );
         const finalMemStore = keyringController.memStore.getState();
 
-        expect(initialMemStore.encryptionKey).toBeNull();
+        expect(initialMemStore.encryptionKey).toBeUndefined();
         expect(initialMemStore.encryptionSalt).toBeUndefined();
 
         expect(finalMemStore.encryptionKey).toBe(MOCK_HARDCODED_KEY);
@@ -888,9 +891,13 @@ describe('KeyringController', () => {
 
       await keyringController.setLocked();
 
-      expect(keyringController.memStore.getState().encryptionSalt).toBeNull();
+      expect(
+        keyringController.memStore.getState().encryptionSalt,
+      ).toBeUndefined();
       expect(keyringController.password).toBeUndefined();
-      expect(keyringController.memStore.getState().encryptionKey).toBeNull();
+      expect(
+        keyringController.memStore.getState().encryptionKey,
+      ).toBeUndefined();
     });
   });
 

src/KeyringController.ts

@@ -3,6 +3,7 @@ import * as encryptorUtils from '@metamask/browser-passworder';
 import HDKeyring from '@metamask/eth-hd-keyring';
 import { normalize as normalizeToHex } from '@metamask/eth-sig-util';
 import SimpleKeyring from '@metamask/eth-simple-keyring';
+import { ObservableStore } from '@metamask/obs-store';
 import { remove0x, isValidHexAddress } from '@metamask/utils';
 import type {
   Hex,
@@ -14,13 +15,13 @@ import type {
 // TODO: Stop using `events`, and remove the notice about this from the README
 // eslint-disable-next-line import/no-nodejs-modules
 import { EventEmitter } from 'events';
-import ObservableStore from 'obs-store';
 
 import { KeyringType, KeyringControllerError } from './constants';
 import {
   SerializedKeyring,
   KeyringControllerArgs,
   KeyringControllerState,
+  KeyringControllerPersistentState,
 } from './types';
 
 const defaultKeyringBuilders = [
@@ -31,9 +32,9 @@ const defaultKeyringBuilders = [
 class KeyringController extends EventEmitter {
   keyringBuilders: { (): Keyring<Json>; type: string }[];
 
-  public store: typeof ObservableStore;
+  public store: ObservableStore<KeyringControllerPersistentState>;
 
-  public memStore: typeof ObservableStore;
+  public memStore: ObservableStore<KeyringControllerState>;
 
   public encryptor: any;
 
@@ -62,7 +63,6 @@ class KeyringController extends EventEmitter {
         (keyringBuilder) => keyringBuilder.type,
       ),
       keyrings: [],
-      encryptionKey: null,
     });
 
     this.encryptor = encryptor;
@@ -167,10 +167,9 @@ class KeyringController extends EventEmitter {
     delete this.password;
 
     // set locked
-    this.memStore.updateState({
+    this.memStore.putState({
       isUnlocked: false,
-      encryptionKey: null,
-      encryptionSalt: null,
+      keyrings: [],
     });
 
     // remove keyrings
@@ -362,9 +361,9 @@ class KeyringController extends EventEmitter {
    *
    * Updates the in-memory keyrings, without persisting.
    */
-  async updateMemStoreKeyrings(): Promise<Json> {
+  async updateMemStoreKeyrings(): Promise<void> {
     const keyrings = await Promise.all(this.keyrings.map(displayForKeyring));
-    return this.memStore.updateState({ keyrings });
+    this.memStore.updateState({ keyrings });
   }
 
   /**
@@ -900,7 +899,10 @@ class KeyringController extends EventEmitter {
         // is not yet inside the memStore
         this.memStore.updateState({
           encryptionKey,
-          encryptionSalt,
+          // we can safely assume that encryptionSalt is defined here
+          // because we compare it with the salt from the vault
+          // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+          encryptionSalt: encryptionSalt!,
         });
       }
     } else {

src/types.ts

@@ -1,32 +1,26 @@
 import type { Json, Keyring } from '@metamask/utils';
-import ObservableStore from 'obs-store';
 
 export type KeyringControllerArgs = {
-  keyringBuilders:
-    | { (): Keyring<Json>; type: string }
-    | ConcatArray<{ (): Keyring<Json>; type: string }>;
-
+  keyringBuilders?: { (): Keyring<Json>; type: string }[];
   cacheEncryptionKey: boolean;
-  initState?: KeyringControllerState;
+  initState?: KeyringControllerPersistentState;
   encryptor?: any;
 };
 
-export type KeyringControllerState = {
-  keyringBuilders?: { (): Keyring<Json>; type: string }[];
-
-  store?: typeof ObservableStore;
-
-  memStore?: typeof ObservableStore;
-
-  keyrings?: Keyring<Json>[];
+export type KeyringObject = {
+  type: string;
+  accounts: string[];
+};
 
-  isUnlocked?: boolean;
+export type KeyringControllerPersistentState = {
+  vault?: string;
+};
 
+export type KeyringControllerState = {
+  keyrings: KeyringObject[];
+  isUnlocked: boolean;
   encryptionKey?: string;
-
   encryptionSalt?: string;
-
-  password?: string;
 };
 
 export type SerializedKeyring = {

yarn.lock

@@ -1040,6 +1040,7 @@ __metadata:
     "@metamask/eth-hd-keyring": ^6.0.0
     "@metamask/eth-sig-util": ^6.0.0
     "@metamask/eth-simple-keyring": ^5.0.0
+    "@metamask/obs-store": ^8.1.0
     "@metamask/utils": ^6.2.0
     "@types/jest": ^29.4.0
     "@types/sinon": ^10.0.13
@@ -1056,7 +1057,6 @@ __metadata:
     ethereumjs-wallet: ^1.0.2
     jest: ^29.0.0
     jest-it-up: ^2.0.2
-    obs-store: ^4.0.3
     prettier: ^2.8.1
     prettier-plugin-packagejson: ^2.3.0
     rimraf: ^3.0.2
@@ -1108,6 +1108,23 @@ __metadata:
   languageName: node
   linkType: hard
 
+"@metamask/obs-store@npm:^8.1.0":
+  version: 8.1.0
+  resolution: "@metamask/obs-store@npm:8.1.0"
+  dependencies:
+    "@metamask/safe-event-emitter": ^2.0.0
+    through2: ^2.0.3
+  checksum: 92356067fa3517526d656f2f0bdfbc4d39f65e27fb30d84240cfc9c1aa9cd5d743498952df18ed8efbb8887b6cc1bc1fab37bde3fb0fc059539e0dfcc67ff86f
+  languageName: node
+  linkType: hard
+
+"@metamask/safe-event-emitter@npm:^2.0.0":
+  version: 2.0.0
+  resolution: "@metamask/safe-event-emitter@npm:2.0.0"
+  checksum: 8b717ac5d53df0027c05509f03d0534700b5898dd1c3a53fb2dc4c0499ca5971b14aae67f522d09eb9f509e77f50afa95fdb3eda1afbff8b071c18a3d2905e93
+  languageName: node
+  linkType: hard
+
 "@metamask/scure-bip39@npm:^2.0.3":
   version: 2.1.0
   resolution: "@metamask/scure-bip39@npm:2.1.0"
@@ -3453,13 +3470,6 @@ __metadata:
   languageName: node
   linkType: hard
 
-"events@npm:^3.0.0":
-  version: 3.3.0
-  resolution: "events@npm:3.3.0"
-  checksum: f6f487ad2198aa41d878fa31452f1a3c00958f46e9019286ff4787c84aac329332ab45c9cdc8c445928fc6d7ded294b9e005a7fce9426488518017831b272780
-  languageName: node
-  linkType: hard
-
 "evp_bytestokey@npm:^1.0.3":
   version: 1.0.3
   resolution: "evp_bytestokey@npm:1.0.3"
@@ -5689,18 +5699,6 @@ __metadata:
   languageName: node
   linkType: hard
 
-"obs-store@npm:^4.0.3":
-  version: 4.0.3
-  resolution: "obs-store@npm:4.0.3"
-  dependencies:
-    readable-stream: ^2.2.2
-    safe-event-emitter: ^1.0.1
-    through2: ^2.0.3
-    xtend: ^4.0.1
-  checksum: a3c05dad7483489f2c083059256f24838b106e10012dd296c7d3e8066869bbc7313dc90775354b519cbeb7aa4c230b0f66cc87ef1414189bad6d03adb0b00b75
-  languageName: node
-  linkType: hard
-
 "once@npm:^1.3.0":
   version: 1.4.0
   resolution: "once@npm:1.4.0"
@@ -6090,7 +6088,18 @@ __metadata:
   languageName: node
   linkType: hard
 
-"readable-stream@npm:^2.2.2, readable-stream@npm:~2.3.6":
+"readable-stream@npm:^3.6.0":
+  version: 3.6.2
+  resolution: "readable-stream@npm:3.6.2"
+  dependencies:
+    inherits: ^2.0.3
+    string_decoder: ^1.1.1
+    util-deprecate: ^1.0.1
+  checksum: bdcbe6c22e846b6af075e32cf8f4751c2576238c5043169a1c221c92ee2878458a816a4ea33f4c67623c0b6827c8a400409bfb3cf0bf3381392d0b1dfb52ac8d
+  languageName: node
+  linkType: hard
+
+"readable-stream@npm:~2.3.6":
   version: 2.3.8
   resolution: "readable-stream@npm:2.3.8"
   dependencies:
@@ -6105,17 +6114,6 @@ __metadata:
   languageName: node
   linkType: hard
 
-"readable-stream@npm:^3.6.0":
-  version: 3.6.2
-  resolution: "readable-stream@npm:3.6.2"
-  dependencies:
-    inherits: ^2.0.3
-    string_decoder: ^1.1.1
-    util-deprecate: ^1.0.1
-  checksum: bdcbe6c22e846b6af075e32cf8f4751c2576238c5043169a1c221c92ee2878458a816a4ea33f4c67623c0b6827c8a400409bfb3cf0bf3381392d0b1dfb52ac8d
-  languageName: node
-  linkType: hard
-
 "readdirp@npm:^3.5.0, readdirp@npm:~3.6.0":
   version: 3.6.0
   resolution: "readdirp@npm:3.6.0"
@@ -6298,15 +6296,6 @@ __metadata:
   languageName: node
   linkType: hard
 
-"safe-event-emitter@npm:^1.0.1":
-  version: 1.0.1
-  resolution: "safe-event-emitter@npm:1.0.1"
-  dependencies:
-    events: ^3.0.0
-  checksum: 2a15094bd28b0966571693f219b5a846949ae24f7ba87c6024f0ed552bef63ebe72970a784b85b77b1f03f1c95e78fabe19306d44538dbc4a3a685bed31c18c4
-  languageName: node
-  linkType: hard
-
 "safe-regex-test@npm:^1.0.0":
   version: 1.0.0
   resolution: "safe-regex-test@npm:1.0.0"
@@ -7306,7 +7295,7 @@ __metadata:
   languageName: node
   linkType: hard
 
-"xtend@npm:^4.0.1, xtend@npm:~4.0.1":
+"xtend@npm:~4.0.1":
   version: 4.0.2
   resolution: "xtend@npm:4.0.2"
   checksum: ac5dfa738b21f6e7f0dd6e65e1b3155036d68104e67e5d5d1bde74892e327d7e5636a076f625599dc394330a731861e87343ff184b0047fef1360a7ec0a5a36a