Pluggable persistence

src/Client.ts

@@ -32,6 +32,7 @@ import {
   NetworkKeystoreProvider,
   StaticKeystoreProvider,
 } from './keystore/providers'
+import { LocalStoragePersistence, Persistence } from './keystore/persistence'
 const { Compression } = proto
 
 // eslint-disable @typescript-eslint/explicit-module-boundary-types
@@ -143,6 +144,17 @@ export type KeyStoreOptions = {
    * A bundle can be retried using `Client.getKeys(...)`
    */
   privateKeyOverride?: Uint8Array
+
+  /**
+   * Override the base persistence provider.
+   * Defaults to LocalStoragePersistence, which is fine for most implementations
+   */
+  basePersistence: Persistence
+  /**
+   * Whether or not the persistence provider should encrypt the values.
+   * Only disable if you are using a secure datastore that already has encryption
+   */
+  disablePersistenceEncryption: boolean
 }
 
 export type LegacyOptions = {
@@ -194,9 +206,12 @@ export function defaultOptions(opts?: Partial<ClientOptions>): ClientOptions {
     maxContentSize: MaxContentSize,
     persistConversations: true,
     skipContactPublishing: false,
+    basePersistence: new LocalStoragePersistence(),
+    disablePersistenceEncryption: false,
     keystoreProviders: defaultKeystoreProviders(),
     apiClientFactory: createHttpApiClientFromOptions,
   }
+
   if (opts?.codecs) {
     opts.codecs = _defaultOptions.codecs.concat(opts.codecs)
   }

src/index.ts

@@ -73,10 +73,11 @@ export {
   UnsubscribeFn,
   OnConnectionLostCallback,
 } from './ApiClient'
-export { Authenticator, AuthCache } from './authn'
+export { Authenticator, AuthCache, LocalAuthenticator } from './authn'
 export {
   nsToDate,
   dateToNs,
+  retry,
   fromNanoString,
   toNanoString,
   mapPaginatedStream,

src/keystore/providers/helpers.ts

@@ -1,10 +1,6 @@
 import { PrivateKeyBundleV2 } from './../../crypto/PrivateKeyBundle'
 import { PrivateKeyBundleV1 } from '../../crypto/PrivateKeyBundle'
-import {
-  EncryptedPersistence,
-  LocalStoragePersistence,
-  PrefixedPersistence,
-} from '../persistence'
+import { EncryptedPersistence, PrefixedPersistence } from '../persistence'
 import { KeystoreProviderOptions } from './interfaces'
 
 export const buildPersistenceFromOptions = async (
@@ -16,9 +12,13 @@ export const buildPersistenceFromOptions = async (
   }
   const address = await keys.identityKey.publicKey.walletSignatureAddress()
   const prefix = `xmtp/${opts.env}/${address}/`
+  const basePersistence = opts.basePersistence
+  const shouldEncrypt = !opts.disablePersistenceEncryption
 
   return new PrefixedPersistence(
     prefix,
-    new EncryptedPersistence(new LocalStoragePersistence(), keys.identityKey)
+    shouldEncrypt
+      ? new EncryptedPersistence(basePersistence, keys.identityKey)
+      : basePersistence
   )
 }

src/keystore/providers/interfaces.ts

@@ -2,11 +2,14 @@ import type { XmtpEnv, PreEventCallbackOptions } from '../../Client'
 import type { Signer } from '../../types/Signer'
 import type { Keystore } from '../interfaces'
 import type { ApiClient } from '../../ApiClient'
+import { Persistence } from '../persistence'
 
 export type KeystoreProviderOptions = {
   env: XmtpEnv
   persistConversations: boolean
   privateKeyOverride?: Uint8Array
+  basePersistence: Persistence
+  disablePersistenceEncryption: boolean
 } & PreEventCallbackOptions
 
 /**

test/Client.test.ts

@@ -8,7 +8,13 @@ import {
 } from './helpers'
 import { buildUserContactTopic } from '../src/utils'
 import Client, { ClientOptions } from '../src/Client'
-import { ApiUrls, Compression, HttpApiClient, PublishParams } from '../src'
+import {
+  ApiUrls,
+  Compression,
+  HttpApiClient,
+  LocalStoragePersistence,
+  PublishParams,
+} from '../src'
 import NetworkKeyManager from '../src/keystore/providers/NetworkKeyManager'
 import TopicPersistence from '../src/keystore/persistence/TopicPersistence'
 import { PrivateKeyBundleV1 } from '../src/crypto'
@@ -336,4 +342,19 @@ describe('ClientOptions', () => {
       expect(c).rejects.toThrow(expectedError)
     })
   })
+
+  describe('pluggable persistence', () => {
+    it('allows for an override of the persistence engine', async () => {
+      class MyNewPersistence extends LocalStoragePersistence {
+        async getItem(key: string): Promise<Uint8Array | null> {
+          throw new Error('MyNewPersistence')
+        }
+      }
+
+      const c = newLocalHostClient({
+        basePersistence: new MyNewPersistence(),
+      })
+      expect(c).rejects.toThrow('MyNewPersistence')
+    })
+  })
 })

test/keystore/providers/KeyGeneratorKeystoreProvider.test.ts

@@ -18,7 +18,7 @@ describe('KeyGeneratorKeystoreProvider', () => {
   it('creates a key when wallet supplied', async () => {
     const provider = new KeyGeneratorKeystoreProvider()
     const keystore = await provider.newKeystore(
-      testProviderOptions(),
+      testProviderOptions({}),
       apiClient,
       wallet
     )
@@ -28,7 +28,7 @@ describe('KeyGeneratorKeystoreProvider', () => {
   it('throws KeystoreProviderUnavailableError when no wallet supplied', async () => {
     const provider = new KeyGeneratorKeystoreProvider()
     const prom = provider.newKeystore(
-      testProviderOptions(),
+      testProviderOptions({}),
       apiClient,
       undefined
     )
@@ -39,7 +39,7 @@ describe('KeyGeneratorKeystoreProvider', () => {
     const provider = new KeyGeneratorKeystoreProvider()
     const preCreateIdentityCallback = jest.fn()
     const keystore = await provider.newKeystore(
-      { ...testProviderOptions(), preCreateIdentityCallback },
+      { ...testProviderOptions({}), preCreateIdentityCallback },
       apiClient,
       wallet
     )
@@ -51,7 +51,7 @@ describe('KeyGeneratorKeystoreProvider', () => {
     const provider = new KeyGeneratorKeystoreProvider()
     const preEnableIdentityCallback = jest.fn()
     const keystore = await provider.newKeystore(
-      { ...testProviderOptions(), preEnableIdentityCallback },
+      { ...testProviderOptions({}), preEnableIdentityCallback },
       apiClient,
       wallet
     )

test/keystore/providers/NetworkKeyManager.test.ts

@@ -6,6 +6,7 @@ import { buildPersistenceFromOptions } from '../../../src/keystore/providers/hel
 import NetworkKeyManager from '../../../src/keystore/providers/NetworkKeyManager'
 import { Signer } from '../../../src/types/Signer'
 import { newWallet, pollFor, sleep, wrapAsLedgerWallet } from '../../helpers'
+import { testProviderOptions } from './helpers'
 
 describe('NetworkKeyManager', () => {
   let wallet: Signer
@@ -100,13 +101,13 @@ describe('NetworkKeyManager', () => {
   it('respects the options provided', async () => {
     const bundle = await PrivateKeyBundleV1.generate(wallet)
     const shouldBeUndefined = await buildPersistenceFromOptions(
-      { persistConversations: false, env: 'local' },
+      testProviderOptions({ persistConversations: false }),
       bundle
     )
     expect(shouldBeUndefined).toBeUndefined()
 
     const shouldBeDefined = await buildPersistenceFromOptions(
-      { persistConversations: true, env: 'local' },
+      testProviderOptions({ persistConversations: true }),
       bundle
     )
     expect(shouldBeDefined).toBeInstanceOf(PrefixedPersistence)

test/keystore/providers/NetworkKeystoreProvider.test.ts

@@ -28,7 +28,7 @@ describe('NetworkKeystoreProvider', () => {
   it('fails gracefully when no keys are found', async () => {
     const provider = new NetworkKeystoreProvider()
     expect(
-      provider.newKeystore(testProviderOptions(), apiClient, wallet)
+      provider.newKeystore(testProviderOptions({}), apiClient, wallet)
     ).rejects.toThrow(KeystoreProviderUnavailableError)
   })
 
@@ -41,7 +41,7 @@ describe('NetworkKeystoreProvider', () => {
 
     const provider = new NetworkKeystoreProvider()
     const keystore = await provider.newKeystore(
-      testProviderOptions(),
+      testProviderOptions({}),
       apiClient,
       wallet
     )
@@ -74,7 +74,7 @@ describe('NetworkKeystoreProvider', () => {
     // Now try and load it
     const provider = new NetworkKeystoreProvider()
     const keystore = await provider.newKeystore(
-      testProviderOptions(),
+      testProviderOptions({}),
       apiClient,
       wallet
     )
@@ -91,7 +91,7 @@ describe('NetworkKeystoreProvider', () => {
     const provider = new NetworkKeystoreProvider()
     const mockNotifier = jest.fn()
     await provider.newKeystore(
-      { ...testProviderOptions(), preEnableIdentityCallback: mockNotifier },
+      { ...testProviderOptions({}), preEnableIdentityCallback: mockNotifier },
       apiClient,
       wallet
     )

test/keystore/providers/StaticKeystoreProvider.test.ts

@@ -1,6 +1,7 @@
 import { privateKey } from '@xmtp/proto'
 import { PrivateKeyBundleV1 } from '../../../src/crypto'
 import { newWallet } from '../../helpers'
+import { testProviderOptions } from './helpers'
 import StaticKeystoreProvider from '../../../src/keystore/providers/StaticKeystoreProvider'
 import { KeystoreProviderUnavailableError } from '../../../src/keystore/providers/errors'
 
@@ -14,31 +15,37 @@ describe('StaticKeystoreProvider', () => {
       v2: undefined,
     }).finish()
     const provider = new StaticKeystoreProvider()
-    const keystore = await provider.newKeystore({
-      privateKeyOverride: keyBytes,
-      env: ENV,
-      persistConversations: false,
-    })
+    const keystore = await provider.newKeystore(
+      testProviderOptions({
+        privateKeyOverride: keyBytes,
+        env: ENV,
+        persistConversations: false,
+      })
+    )
 
     expect(keystore).not.toBeNull()
   })
 
   it('throws with an unset key', async () => {
     expect(
-      new StaticKeystoreProvider().newKeystore({
-        env: ENV,
-        persistConversations: false,
-      })
+      new StaticKeystoreProvider().newKeystore(
+        testProviderOptions({
+          env: ENV,
+          persistConversations: false,
+        })
+      )
     ).rejects.toThrow(KeystoreProviderUnavailableError)
   })
 
   it('fails with an invalid key', async () => {
     expect(
-      new StaticKeystoreProvider().newKeystore({
-        privateKeyOverride: Uint8Array.from([1, 2, 3]),
-        env: ENV,
-        persistConversations: false,
-      })
+      new StaticKeystoreProvider().newKeystore(
+        testProviderOptions({
+          privateKeyOverride: Uint8Array.from([1, 2, 3]),
+          env: ENV,
+          persistConversations: false,
+        })
+      )
     ).rejects.toThrow()
   })
 })

test/keystore/providers/helpers.ts

@@ -1,8 +1,15 @@
-export const testProviderOptions = (
+import { LocalStoragePersistence } from '../../../src'
+import { KeystoreProviderOptions } from '../../../src/keystore/providers'
+
+export const testProviderOptions = ({
   privateKeyOverride = undefined,
-  persistConversations = false
-) => ({
-  env: 'local' as const,
+  persistConversations = false,
+  basePersistence = new LocalStoragePersistence(),
+  env = 'local' as const,
+}: Partial<KeystoreProviderOptions>) => ({
+  env,
   persistConversations,
   privateKeyOverride,
+  basePersistence,
+  disablePersistenceEncryption: false,
 })