refactor: Simplify key exchange subscriber by removing buffering & pending promise de-duplication.

Author: timoxley

src/stream/encryption/KeyExchangeSubscriber.ts

@@ -3,8 +3,6 @@ import {
 } from 'streamr-client-protocol'
 import pMemoize from 'p-memoize'
 import { uuid, Defer } from '../../utils'
-import Scaffold from '../../utils/Scaffold'
-import mem from 'mem'
 
 import { validateOptions } from '../utils'
 import EncryptionUtil, { GroupKey } from './Encryption'
@@ -22,7 +20,7 @@ import {
 type MessageMatch = (content: any, streamMessage: StreamMessage) => boolean
 
 function waitForSubMessage(sub: Subscription, matchFn: MessageMatch) {
-    const task = Defer()
+    const task = Defer<StreamMessage | undefined>()
     const onMessage = (content: any, streamMessage: StreamMessage) => {
         try {
             if (matchFn(content, streamMessage)) {
@@ -85,12 +83,7 @@ export class SubscriberKeyExchange {
     client
     initialGroupKeys
     encryptionUtil
-    pending = new Map<GroupKeyId, ReturnType<typeof Defer>>()
-    getBuffer = mem<(groupKeyId: GroupKeyId) => GroupKeyId[], [string]>(() => [])
-    timeouts: Record<string, ReturnType<typeof setTimeout>> = Object.create(null)
-    next
     enabled = true
-    sub?: Subscription
 
     constructor(client: StreamrClient, { groupKeys = {} }: KeyExhangeOptions = {}) {
         this.client = client
@@ -102,80 +95,51 @@ export class SubscriberKeyExchange {
             }
         })
         this.encryptionUtil = new EncryptionUtil(client.options.keyExchange)
-        this.next = this.initNext()
+    }
+
+    async getSubscription() {
+        return SubscriberKeyExhangeSubscription(this.client, this.getGroupKeyStore, this.encryptionUtil)
     }
 
     async requestKeys({ streamId, publisherId, groupKeyIds }: {
         streamId: string,
         publisherId: string,
         groupKeyIds: GroupKeyId[]
     }) {
-        let done = false
         const requestId = uuid('GroupKeyRequest')
         const rsaPublicKey = this.encryptionUtil.getPublicKey()
         const keyExchangeStreamId = getKeyExchangeStreamId(publisherId)
-        let responseTask: ReturnType<typeof Defer>
-        let cancelTask: ReturnType<typeof Defer>
-        let receivedGroupKeys: GroupKey[] = []
-        let response: any
-
-        this.requestKeysStep = Scaffold([
-            async () => {
-                if (!this.sub) { throw new Error('no subscription') }
-                cancelTask = Defer()
-                responseTask = waitForSubMessage(this.sub, (content, streamMessage) => {
-                    const { messageType } = streamMessage
-                    const matchesMessageType = (
-                        messageType === StreamMessage.MESSAGE_TYPES.GROUP_KEY_RESPONSE
-                        || messageType === StreamMessage.MESSAGE_TYPES.GROUP_KEY_ERROR_RESPONSE
-                    )
-
-                    if (!matchesMessageType) {
-                        return false
-                    }
-
-                    const groupKeyResponse = GroupKeyResponse.fromArray(content)
-                    return groupKeyResponse.requestId === requestId
-                })
-
-                cancelTask.then(responseTask.resolve).catch(responseTask.reject)
-                return () => {
-                    cancelTask.resolve(undefined)
-                }
-            }, async () => {
-                const msg = new GroupKeyRequest({
-                    streamId,
-                    requestId,
-                    rsaPublicKey,
-                    groupKeyIds,
-                })
-                await this.client.publish(keyExchangeStreamId, msg)
-            }, async () => {
-                response = await responseTask
-                return () => {
-                    response = undefined
+        let sub!: Subscription
+        try {
+            sub = await this.getSubscription()
+            const responseTask = waitForSubMessage(sub, (content, streamMessage) => {
+                const { messageType } = streamMessage
+                const matchesMessageType = (
+                    messageType === StreamMessage.MESSAGE_TYPES.GROUP_KEY_RESPONSE
+                    || messageType === StreamMessage.MESSAGE_TYPES.GROUP_KEY_ERROR_RESPONSE
+                )
+
+                if (!matchesMessageType) {
+                    return false
                 }
-            }, async () => {
-                receivedGroupKeys = response ? await getGroupKeysFromStreamMessage(response, this.encryptionUtil) : []
 
-                return () => {
-                    receivedGroupKeys = []
-                }
-            },
-        ], async () => this.enabled && !done, {
-            id: `requestKeys.${requestId}`,
-            onChange(isGoingUp) {
-                if (!isGoingUp && cancelTask) {
-                    cancelTask.resolve(undefined)
-                }
+                const groupKeyResponse = GroupKeyResponse.fromArray(content)
+                return groupKeyResponse.requestId === requestId
+            })
+            const msg = new GroupKeyRequest({
+                streamId,
+                requestId,
+                rsaPublicKey,
+                groupKeyIds,
+            })
+            await this.client.publish(keyExchangeStreamId, msg)
+            const response = await responseTask
+            return response ? getGroupKeysFromStreamMessage(response, this.encryptionUtil) : []
+        } finally {
+            if (sub) {
+                await sub.unsubscribe()
             }
-        })
-
-        await this.requestKeysStep()
-        const keys = receivedGroupKeys.slice()
-        done = true
-        await this.requestKeysStep()
-        return keys
+        }
     }
 
     async getGroupKeyStore(streamId: string) {
@@ -187,46 +151,6 @@ export class SubscriberKeyExchange {
         })
     }
 
-    private async processBuffer({ streamId, publisherId }: { streamId: string, publisherId: string }) {
-        if (!this.enabled) { return }
-        const key = `${streamId}.${publisherId}`
-        const currentBuffer = this.getBuffer(key)
-        const groupKeyIds = currentBuffer.slice()
-        const groupKeyStore = await this.getGroupKeyStore(streamId)
-        currentBuffer.length = 0
-        try {
-            const receivedGroupKeys = await this.requestKeys({
-                streamId,
-                publisherId,
-                groupKeyIds,
-            })
-            if (!this.enabled) { return }
-            await Promise.all(receivedGroupKeys.map(async (groupKey) => (
-                groupKeyStore.add(groupKey)
-            )))
-            if (!this.enabled) { return }
-            await Promise.all(groupKeyIds.map(async (id) => {
-                if (!this.pending.has(id)) { return }
-                const groupKeyTask = groupKeyStore.get(id)
-                const task = this.pending.get(id)
-                this.pending.delete(id)
-                const groupKey = await groupKeyTask
-                if (task) {
-                    task.resolve(groupKey)
-                }
-            }))
-        } catch (err) {
-            groupKeyIds.forEach((id) => {
-                if (!this.pending.has(id)) { return }
-                const task = this.pending.get(id)
-                this.pending.delete(id)
-                if (task) {
-                    task.reject(err)
-                }
-            })
-        }
-    }
-
     async getKey(streamMessage: StreamMessage) {
         const streamId = streamMessage.getStreamId()
         const publisherId = streamMessage.getPublisherId()
@@ -245,79 +169,32 @@ export class SubscriberKeyExchange {
             return existingGroupKey
         }
 
-        if (this.pending.has(groupKeyId)) {
-            return this.pending.get(groupKeyId)
-        }
+        const receivedGroupKeys = await this.requestKeys({
+            streamId,
+            publisherId,
+            groupKeyIds: [groupKeyId],
+        })
 
-        const key = `${streamId}.${publisherId}`
-        const buffer = this.getBuffer(key)
-        buffer.push(groupKeyId)
-        this.pending.set(groupKeyId, Defer())
+        await Promise.all(receivedGroupKeys.map(async (groupKey: GroupKey) => (
+            groupKeyStore.add(groupKey)
+        )))
 
-        if (!this.timeouts[key]) {
-            this.timeouts[key] = setTimeout(() => {
-                delete this.timeouts[key]
-                this.processBuffer({ streamId, publisherId })
-            }, 1000)
-        }
-
-        return this.pending.get(groupKeyId)
+        return groupKeyStore.get(groupKeyId)
     }
 
     cleanupPending() {
-        Array.from(Object.entries(this.timeouts)).forEach(([key, value]) => {
-            clearTimeout(value)
-            delete this.timeouts[key]
-        })
-        const pendingValues = Array.from(this.pending.values())
-        this.pending.clear()
-        pendingValues.forEach((value) => {
-            value.resolve(undefined)
-        })
         pMemoize.clear(this.getGroupKeyStore)
     }
 
-    initNext() {
-        return Scaffold([
-            async () => {
-                await this.encryptionUtil.onReady()
-            },
-            async () => {
-                this.sub = await SubscriberKeyExhangeSubscription(this.client, this.getGroupKeyStore, this.encryptionUtil)
-                return async () => {
-                    if (!this.sub) { return }
-                    const cancelTask = this.sub.cancel()
-                    this.sub = undefined
-                    await cancelTask
-                }
-            }
-        ], async () => this.enabled, {
-            id: `SubscriberKeyExhangeSubscription.${this.client.id}`,
-            onChange: (shouldUp) => {
-                if (!shouldUp) {
-                    this.cleanupPending()
-                }
-            },
-            onDone: async () => {
-                // clean up requestKey
-                if (this.requestKeysStep) {
-                    await this.requestKeysStep()
-                }
-            }
-        })
-    }
-
     async getGroupKey(streamMessage: StreamMessage) {
         if (!streamMessage.groupKeyId) { return [] }
-        await this.next()
-        if (!this.enabled) { return [] }
+        await this.encryptionUtil.onReady()
 
         return this.getKey(streamMessage)
     }
 
     async start() {
         this.enabled = true
-        return this.next()
     }
 
     async addNewKey(streamMessage: StreamMessage) {
@@ -331,6 +208,5 @@ export class SubscriberKeyExchange {
 
     async stop() {
         this.enabled = false
-        return this.next()
     }
 }