feat: multi-seed memoized keychain

features/keychain/index.js

@@ -1,11 +1,16 @@
 import moduleDefinition from './module/keychain'
+import memoizedModuleDefinition from './module/memoized-keychain'
 import apiDefinition from './api'
 
 const keychain = ({ cachePublicKeys }) => {
-  // TODO: support caching
   return {
     id: 'keychain',
-    definitions: [{ definition: moduleDefinition }, { definition: apiDefinition }],
+    definitions: [
+      cachePublicKeys
+        ? { definition: memoizedModuleDefinition, storage: { namespace: 'keychain' } }
+        : { definition: moduleDefinition },
+      { definition: apiDefinition },
+    ],
   }
 }
 

features/keychain/module/__tests__/keychain.integration-test.js

@@ -1,3 +1,4 @@
+import createInMemoryStorage from '@exodus/storage-memory'
 import {
   getPublicKey as getCardanoPublicKey,
   getShelleyAddress as getCardanoAddress,
@@ -12,6 +13,7 @@ import simpleTx from './fixtures/simple-tx'
 
 import KeyIdentifier from '@exodus/key-identifier'
 import keychainDefinition, { Keychain } from '../keychain'
+import memoizedKeychainDefinition from '../memoized-keychain'
 import { getSeedId } from '../crypto/seed-id'
 
 const { factory: createMultiSeedKeychain } = keychainDefinition
@@ -24,7 +26,7 @@ const secondSeed = mnemonicToSeed(
   'wine system mean beyond filter human meat rubber episode wash stomach aunt'
 )
 
-describe.each([
+const seedOpts = [
   // reduce fixtures by switching seeds
   {
     primarySeed: seed,
@@ -36,11 +38,33 @@ describe.each([
     secondarySeed: seed,
     seedId: getSeedId(seed),
   },
-])('multi-seed-keychain', ({ primarySeed, secondarySeed, seedId }) => {
+]
+
+const keychainImplementations = [
+  {
+    module: 'keychain',
+    factory: () => createMultiSeedKeychain(),
+    ...seedOpts,
+  },
+  {
+    module: 'memoizedKeychain',
+    factory: () => memoizedKeychainDefinition.factory({ storage: createInMemoryStorage() }),
+    ...seedOpts,
+  },
+]
+
+describe.each(
+  seedOpts.flatMap((seedOptsVariant) =>
+    keychainImplementations.map((implementationVariant) => ({
+      ...implementationVariant,
+      ...seedOptsVariant,
+    }))
+  )
+)('multi-seed-keychain: $module', ({ module, factory, primarySeed, secondarySeed, seedId }) => {
   let keychain
 
   it('should construct correctly', () => {
-    expect(keychainDefinition.factory()).toBeInstanceOf(Keychain)
+    expect(factory()).toBeInstanceOf(Keychain)
   })
 
   beforeEach(() => {

features/keychain/module/__tests__/memoized-keychain.test.js

@@ -0,0 +1,106 @@
+import stableStringify from 'json-stable-stringify'
+import { mnemonicToSeed } from 'bip39'
+import BJSON from 'buffer-json'
+
+import KeyIdentifier from '@exodus/key-identifier'
+import memoizedKeychainDefinition from '../memoized-keychain'
+import { getSeedId } from '../crypto/seed-id'
+
+const mockStorage = {
+  get: jest.fn(),
+  set: jest.fn(),
+}
+
+const seed = mnemonicToSeed(
+  'menu memory fury language physical wonder dog valid smart edge decrease worth'
+)
+
+const seedId = getSeedId(seed)
+
+const keyIdentifierEthereum0 = new KeyIdentifier({
+  derivationAlgorithm: 'BIP32',
+  derivationPath: "m/44'/60'/0'/0/0",
+  assetName: 'ethereum',
+})
+
+const storage = mockStorage
+
+describe('memoizedKeychain', () => {
+  let memoizedKeychain
+
+  beforeEach(() => {
+    mockStorage.get.mockResolvedValue(null)
+    memoizedKeychain = memoizedKeychainDefinition.factory({ storage })
+    memoizedKeychain.addSeed(seed)
+    jest.clearAllMocks()
+  })
+
+  it('should construct correctly', () => {
+    expect(() => memoizedKeychainDefinition.factory({ storage })).not.toThrow()
+  })
+
+  it('should clone correctly', () => {
+    expect(() => memoizedKeychainDefinition.factory({ storage }).clone()).not.toThrow()
+  })
+
+  describe.each([
+    {
+      name: 'memoized',
+      factory: () => {
+        const keychain = memoizedKeychainDefinition.factory({ storage })
+        keychain.addSeed(seed)
+        return keychain
+      },
+    },
+    {
+      name: 'memoizedClone',
+      factory: () => {
+        const clone = memoizedKeychain.clone()
+        clone.addSeed(seed)
+        return clone
+      },
+    },
+  ])('memoization: $name', ({ factory }) => {
+    const storageKey = stableStringify({
+      seedId,
+      keyId: {
+        assetName: 'ethereum',
+        derivationAlgorithm: 'BIP32',
+        derivationPath: "m/44'/60'/0'/0/0",
+        keyType: 'secp256k1',
+      },
+    })
+
+    const storageData = {
+      [storageKey]: {
+        xpub: 'xpub6H8P7xAy1nvvefMui3rD6yK3cdkBSAhukKRcxeqydPqdm8L8FAvxu33Hgoajcr8PW1oBPDm7sRdPSoUz55kcCF9LNd5RatNgExPrn8Pvd5P',
+        publicKey: Buffer.from('AgyNoIyo7xR+pkMj5TnXkQE/8NtK6dMck+KUghia9w3l', 'base64'),
+      },
+    }
+
+    it('should retrieve publicKeys from storage at construction', async () => {
+      mockStorage.get.mockResolvedValueOnce(BJSON.stringify(storageData))
+      factory()
+      expect(mockStorage.get).toHaveBeenCalledTimes(1)
+    })
+
+    it('should write publicKeys to storage if not found', async () => {
+      mockStorage.get.mockResolvedValueOnce(null)
+      const keychain = factory()
+      const exported = await keychain.exportKey({ seedId, keyId: keyIdentifierEthereum0 })
+      expect(mockStorage.set).toHaveBeenCalledTimes(1)
+      expect(mockStorage.set).toHaveBeenCalledWith('data1', BJSON.stringify(storageData))
+      const exportedFromCache = await keychain.exportKey({ seedId, keyId: keyIdentifierEthereum0 })
+      expect(mockStorage.set).toHaveBeenCalledTimes(1)
+      expect(exported).toStrictEqual({ privateKey: null, xpriv: null, ...exportedFromCache })
+    })
+
+    it('should still allow private keys', async () => {
+      mockStorage.get.mockResolvedValueOnce(null)
+      const keychain = factory()
+      await expect(
+        keychain.exportKey({ seedId, keyId: keyIdentifierEthereum0, exportPrivate: true })
+      ).resolves.not.toThrow()
+    })
+  })
+})

features/keychain/module/memoized-keychain.js

@@ -1,11 +1,20 @@
+import typeforce from '@exodus/typeforce'
 import BJSON from 'buffer-json'
 import stableStringify from 'json-stable-stringify'
 
 import { Keychain } from './keychain'
 
-const keyIdToCacheKey = stableStringify
+const cacheParamsType = {
+  seedId: 'String',
+  keyId: 'Object',
+}
+
+const keyIdToCacheKey = (opts) => {
+  typeforce(cacheParamsType, opts, true)
+  return stableStringify(opts)
+}
 
-const CACHE_KEY = 'data'
+const CACHE_KEY = 'data1'
 
 const getPublicKeyData = ({ xpub, publicKey }) => ({ xpub, publicKey })
 
@@ -22,29 +31,33 @@ class MemoizedKeychain extends Keychain {
       this.#publicKeys = data ? BJSON.parse(data) : Object.create(null)
     })
 
-    this.#cloneOpts = { storage }
+    this.#cloneOpts = { storage, legacyPrivToPub }
   }
 
-  #getCachedKey = async (keyId) => {
-    const cacheKey = keyIdToCacheKey(keyId)
+  #getCachedKey = async (cacheParams) => {
+    const cacheKey = keyIdToCacheKey(cacheParams)
     return this.#publicKeys[cacheKey]
   }
 
-  #setCachedKey = async (keyId, value) => {
-    this.#publicKeys[keyIdToCacheKey(keyId)] = value
+  #setCachedKey = async (cacheParams, value) => {
+    this.#publicKeys[keyIdToCacheKey(cacheParams)] = value
     await this.#storage.set(CACHE_KEY, BJSON.stringify(this.#publicKeys))
   }
 
-  exportKey = async (keyId, opts) => {
-    if (!opts?.exportPrivate) {
+  async exportKey(opts) {
+    typeforce({ ...cacheParamsType, exportPrivate: '?Boolean' }, opts, true)
+
+    const { seedId, keyId, exportPrivate } = opts
+    const cacheParams = { seedId, keyId }
+    if (!exportPrivate) {
       // take advantage of public key cache
-      const cached = await this.#getCachedKey(keyId)
+      const cached = await this.#getCachedKey(cacheParams)
       if (cached) return cached
     }
 
-    const result = await super.exportKey(keyId, opts)
+    const result = await super.exportKey(opts)
     // don't wait for this to finish
-    this.#setCachedKey(keyId, getPublicKeyData(result))
+    this.#setCachedKey(cacheParams, getPublicKeyData(result))
     return result
   }
 
@@ -54,6 +67,11 @@ class MemoizedKeychain extends Keychain {
     await super.clear()
     await this.#storage.delete(CACHE_KEY)
   }
+
+  removeAllSeeds = () => {
+    super.removeAllSeeds()
+    this.#publicKeys = Object.create(null)
+  }
 }
 
 const memoizedKeychainDefinition = {

features/keychain/package.json

@@ -35,6 +35,7 @@
     "@exodus/key-utils": "^3.0.0",
     "@exodus/slip10": "^1.0.0",
     "@exodus/sodium-crypto": "^3.1.0",
+    "@exodus/typeforce": "^1.18.1",
     "buffer-json": "^2.0.0",
     "create-hash": "^1.2.0",
     "json-stable-stringify": "^1.0.1",

yarn.lock

@@ -2033,6 +2033,7 @@ __metadata:
     "@exodus/key-utils": ^3.0.0
     "@exodus/slip10": ^1.0.0
     "@exodus/sodium-crypto": ^3.1.0
+    "@exodus/typeforce": ^1.18.1
     bip39: 2.6.0
     buffer-json: ^2.0.0
     create-hash: ^1.2.0
@@ -2121,6 +2122,13 @@ __metadata:
   languageName: node
   linkType: hard
 
+"@exodus/typeforce@npm:^1.18.1":
+  version: 1.18.1
+  resolution: "@exodus/typeforce@npm:1.18.1"
+  checksum: 9acdc0bb4fd3bd377871e1632233c45b032498a05b597ccbf37ed5cf3085f713d8b55332b07d5039caf4173cf8b58ab8c44abe85c623d0e6abc311a1a523534b
+  languageName: node
+  linkType: hard
+
 "@fastify/busboy@npm:^2.0.0":
   version: 2.0.0
   resolution: "@fastify/busboy@npm:2.0.0"