add export and import of session keys (#257)

packages/encryption/src/cryptoStore.ts

@@ -1,4 +1,9 @@
-import { AccountRecord, GroupSessionRecord, UserDeviceRecord } from './storeTypes'
+import {
+    AccountRecord,
+    ExtendedInboundGroupSessionData,
+    GroupSessionRecord,
+    UserDeviceRecord,
+} from './storeTypes'
 import Dexie, { Table } from 'dexie'
 
 import { InboundGroupSessionData } from './encryptionDevice'
@@ -11,7 +16,7 @@ const DEFAULT_USER_DEVICE_EXPIRATION_TIME_MS = 15 * 60 * 1000
 export class CryptoStore extends Dexie {
     account!: Table<AccountRecord>
     outboundGroupSessions!: Table<GroupSessionRecord>
-    inboundGroupSessions!: Table<InboundGroupSessionData & { streamId: string; sessionId: string }>
+    inboundGroupSessions!: Table<ExtendedInboundGroupSessionData>
     devices!: Table<UserDeviceRecord>
     userId: string
 
@@ -77,6 +82,10 @@ export class CryptoStore extends Dexie {
         return await this.inboundGroupSessions.get({ sessionId, streamId })
     }
 
+    async getAllEndToEndInboundGroupSessions(): Promise<ExtendedInboundGroupSessionData[]> {
+        return await this.inboundGroupSessions.toArray()
+    }
+
     async storeEndToEndInboundGroupSession(
         streamId: string,
         sessionId: string,

packages/encryption/src/encryptionDevice.ts

@@ -8,7 +8,6 @@ import {
     IOutboundGroupSessionKey,
     OutboundGroupSession,
     Utility,
-    Session,
 } from './encryptionTypes'
 import { EncryptionDelegate } from './encryptionDelegate'
 import { GROUP_ENCRYPTION_ALGORITHM, GroupEncryptionSession } from './olmLib'
@@ -635,7 +634,6 @@ export class EncryptionDevice {
      *
      * @param streamId - streamId of session
      * @param sessionId  - session identifier
-     * @param sessionData - the session object from the store
      */
     public async exportInboundGroupSession(
         streamId: string,
@@ -661,4 +659,36 @@ export class EncryptionDevice {
             algorithm: GROUP_ENCRYPTION_ALGORITHM,
         }
     }
+
+    /**
+     * Get a list containing all of the room keys
+     *
+     * @returns a list of session export objects
+     */
+    public async exportRoomKeys(): Promise<GroupEncryptionSession[] | undefined> {
+        const exportedSessions: GroupEncryptionSession[] = []
+
+        await this.cryptoStore.withGroupSessions(async () => {
+            const sessions = await this.cryptoStore.getAllEndToEndInboundGroupSessions()
+            for (const sessionData of sessions) {
+                if (!sessionData) {
+                    continue
+                }
+
+                const session = this.unpickleInboundGroupSession(sessionData)
+                const messageIndex = session.first_known_index()
+                const sessionKey = session.export_session(messageIndex)
+                session.free()
+
+                exportedSessions.push({
+                    streamId: sessionData.streamId,
+                    sessionId: sessionData.sessionId,
+                    sessionKey: sessionKey,
+                    algorithm: GROUP_ENCRYPTION_ALGORITHM,
+                })
+            }
+        })
+
+        return exportedSessions
+    }
 }

packages/encryption/src/groupEncryptionCrypto.ts

@@ -16,6 +16,24 @@ import { check, dlog } from '@river-build/dlog'
 
 const log = dlog('csb:encryption:groupEncryptionCrypto')
 
+export interface ImportRoomKeysOpts {
+    /** Reports ongoing progress of the import process. Can be used for feedback. */
+    progressCallback?: (stage: ImportRoomKeyProgressData) => void
+}
+
+/**
+ * Room key import progress report.
+ * Used when calling {@link GroupEncryptionCrypto#importRoomKeys} or
+ * {@link GroupEncryptionCrypto#importRoomKeysAsJson} as the parameter of
+ * the progressCallback. Used to display feedback.
+ */
+export interface ImportRoomKeyProgressData {
+    stage: string // TODO: Enum
+    successes?: number
+    failures?: number
+    total?: number
+}
+
 export class GroupEncryptionCrypto {
     private delegate: EncryptionDelegate | undefined
 
@@ -160,4 +178,70 @@ export class GroupEncryptionCrypto {
             ),
         )
     }
+
+    /**
+     * Import a list of room keys previously exported by exportRoomKeys
+     *
+     * @param keys - a list of session export objects
+     * @returns a promise which resolves once the keys have been imported
+     */
+    public importRoomKeys(
+        keys: GroupEncryptionSession[],
+        opts: ImportRoomKeysOpts = {},
+    ): Promise<void> {
+        let successes = 0
+        let failures = 0
+        const total = keys.length
+
+        function updateProgress(): void {
+            opts.progressCallback?.({
+                stage: 'load_keys',
+                successes,
+                failures,
+                total,
+            })
+        }
+
+        return Promise.all(
+            keys.map(async (key) => {
+                if (!key.streamId || !key.algorithm) {
+                    log('ignoring room key entry with missing fields', key)
+                    failures++
+                    if (opts.progressCallback) {
+                        updateProgress()
+                    }
+                    return
+                }
+
+                try {
+                    await this.groupDecryption.importStreamKey(key.streamId, key)
+                    successes++
+                    if (opts.progressCallback) {
+                        updateProgress()
+                    }
+                } catch (error) {
+                    log('failed to import key', error)
+                    failures++
+                    if (opts.progressCallback) {
+                        updateProgress()
+                    }
+                }
+                return
+            }),
+        ).then()
+    }
+
+    /**
+     * Import a JSON string encoding a list of room keys previously
+     * exported by exportRoomKeysAsJson
+     *
+     * @param keys - a JSON string encoding a list of session export
+     *    objects, each of which is an GroupEncryptionSession
+     * @param opts - options object
+     * @returns a promise which resolves once the keys have been imported
+     */
+    public async importRoomKeysAsJson(keys: string): Promise<void> {
+        // eslint-disable-next-line @typescript-eslint/no-unsafe-argument
+        return await this.importRoomKeys(JSON.parse(keys))
+    }
 }

packages/encryption/src/storeTypes.ts

@@ -1,3 +1,5 @@
+import { InboundGroupSessionData } from './encryptionDevice'
+
 export interface AccountRecord {
     id: string
     accountPickle: string
@@ -15,3 +17,8 @@ export interface UserDeviceRecord {
     fallbackKey: string
     expirationTimestamp: number
 }
+
+export interface ExtendedInboundGroupSessionData extends InboundGroupSessionData {
+    streamId: string
+    sessionId: string
+}

packages/encryption/src/tests/decryptionExtensions.test.ts

@@ -85,6 +85,47 @@ describe('TestDecryptionExtensions', () => {
         expect(bobDex.seenStates).toContain(DecryptionStatus.respondingToKeyRequests)
         expect(aliceDex.seenStates).toContain(DecryptionStatus.processingNewGroupSessions)
     })
+
+    test('should be able to export/import stream room key', async () => {
+        // arrange
+        const clientDiscoveryService: ClientDiscoveryService = {}
+        const streamId = genStreamId()
+        const alice = genUserId('Alice')
+        const bob = genUserId('Bob')
+        const bobsPlaintext = "bob's plaintext"
+        const { decryptionExtension: aliceDex } = await createCryptoMocks(
+            alice,
+            clientDiscoveryService,
+        )
+        const { crypto: bobCrypto, decryptionExtension: bobDex } = await createCryptoMocks(
+            bob,
+            clientDiscoveryService,
+        )
+
+        // act
+        aliceDex.start()
+        // bob starts the decryption extension
+        bobDex.start()
+        // bob encrypts a message
+        const encryptedData = await bobCrypto.encryptGroupEvent(streamId, bobsPlaintext)
+        // alice doesn't have the session key
+        // alice imports the keys exported by bob
+        const roomKeys = await bobDex.crypto.encryptionDevice.exportRoomKeys()
+        if (roomKeys) {
+            await aliceDex.crypto.importRoomKeys(roomKeys)
+        }
+
+        // after alice gets the session key,
+        // try to decrypt the message
+        const decrypted = await aliceDex.crypto.decryptGroupEvent(streamId, encryptedData)
+
+        // stop the decryption extensions
+        await bobDex.stop()
+        await aliceDex.stop()
+
+        // assert
+        expect(decrypted).toBe(bobsPlaintext)
+    })
 })
 
 type ReleaseFunction = (value: void | PromiseLike<void>) => void