From bbd7010bed9a5cadaac981c999cf139f0f8b3a12 Mon Sep 17 00:00:00 2001
From: Ry Racherbaumer <ry@xmtp.com>
Date: Wed, 24 Jan 2024 15:25:20 -0600
Subject: [PATCH] fix: allow specific topics when getting HMAC keys

---
 src/keystore/InMemoryKeystore.ts       |  19 ++++-
 src/keystore/rpcDefinitions.ts         |   6 +-
 test/keystore/InMemoryKeystore.test.ts | 107 ++++++++++++++++++++++++-
 3 files changed, 127 insertions(+), 5 deletions(-)

diff --git a/src/keystore/InMemoryKeystore.ts b/src/keystore/InMemoryKeystore.ts
index bcd8c1c4..a2ee3077 100644
--- a/src/keystore/InMemoryKeystore.ts
+++ b/src/keystore/InMemoryKeystore.ts
@@ -34,12 +34,12 @@ import {
   userPreferencesEncrypt,
   generateUserPreferencesTopic,
 } from '../crypto/selfEncryption'
-import type { KeystoreInterface } from '..'
 import {
   exportHmacKey,
   generateHmacSignature,
   hkdfHmacKey,
 } from '../crypto/encryption'
+import type { KeystoreInterface } from './rpcDefinitions'
 
 const { ErrorCode } = keystore
 
@@ -595,15 +595,28 @@ export default class InMemoryKeystore implements KeystoreInterface {
     return this.v2Store.lookup(topic)
   }
 
-  async getV2ConversationHmacKeys(): Promise<keystore.GetConversationHmacKeysResponse> {
+  async getV2ConversationHmacKeys(
+    req?: keystore.GetConversationHmacKeysRequest
+  ): Promise<keystore.GetConversationHmacKeysResponse> {
     const thirtyDayPeriodsSinceEpoch = Math.floor(
       Date.now() / 1000 / 60 / 60 / 24 / 30
     )
 
     const hmacKeys: keystore.GetConversationHmacKeysResponse['hmacKeys'] = {}
 
+    let topics = this.v2Store.topics
+
+    // if specific topics are requested, only include those topics
+    if (req?.topics) {
+      topics = topics.filter(
+        (topicData) =>
+          topicData.invitation !== undefined &&
+          req.topics.includes(topicData.invitation.topic)
+      )
+    }
+
     await Promise.all(
-      this.v2Store.topics.map(async (topicData) => {
+      topics.map(async (topicData) => {
         if (topicData.invitation?.topic) {
           const keyMaterial = getKeyMaterial(topicData.invitation)
           const values = await Promise.all(
diff --git a/src/keystore/rpcDefinitions.ts b/src/keystore/rpcDefinitions.ts
index 798f5436..db73e08c 100644
--- a/src/keystore/rpcDefinitions.ts
+++ b/src/keystore/rpcDefinitions.ts
@@ -199,8 +199,12 @@ export const apiDefs = {
     req: null,
     res: keystore.GetPrivatePreferencesTopicIdentifierResponse,
   },
+  /**
+   * Returns the conversation HMAC keys for the current, previous, and next
+   * 30 day periods since the epoch
+   */
   getV2ConversationHmacKeys: {
-    req: null,
+    req: keystore.GetConversationHmacKeysRequest,
     res: keystore.GetConversationHmacKeysResponse,
   },
 }
diff --git a/test/keystore/InMemoryKeystore.test.ts b/test/keystore/InMemoryKeystore.test.ts
index 41f2aacb..c1a03736 100644
--- a/test/keystore/InMemoryKeystore.test.ts
+++ b/test/keystore/InMemoryKeystore.test.ts
@@ -868,7 +868,7 @@ describe('InMemoryKeystore', () => {
   })
 
   describe('getV2ConversationHmacKeys', () => {
-    it('returns conversation HMAC keys', async () => {
+    it('returns all conversation HMAC keys', async () => {
       const baseTime = new Date()
       const timestamps = Array.from(
         { length: 5 },
@@ -967,5 +967,110 @@ describe('InMemoryKeystore', () => {
         })
       )
     })
+
+    it('returns specific conversation HMAC keys', async () => {
+      const baseTime = new Date()
+      const timestamps = Array.from(
+        { length: 10 },
+        (_, i) => new Date(baseTime.getTime() + i)
+      )
+
+      const invites = await Promise.all(
+        [...timestamps].map(async (createdAt) => {
+          let keys = await PrivateKeyBundleV1.generate(newWallet())
+
+          const recipient = SignedPublicKeyBundle.fromLegacyBundle(
+            keys.getPublicKeyBundle()
+          )
+
+          return aliceKeystore.createInvite({
+            recipient,
+            createdNs: dateToNs(createdAt),
+            context: undefined,
+          })
+        })
+      )
+
+      const thirtyDayPeriodsSinceEpoch = Math.floor(
+        Date.now() / 1000 / 60 / 60 / 24 / 30
+      )
+
+      const periods = [
+        thirtyDayPeriodsSinceEpoch - 1,
+        thirtyDayPeriodsSinceEpoch,
+        thirtyDayPeriodsSinceEpoch + 1,
+      ]
+
+      const randomInvites = invites.slice(3, 8)
+
+      const { hmacKeys } = await aliceKeystore.getV2ConversationHmacKeys({
+        topics: randomInvites.map((invite) => invite.conversation!.topic),
+      })
+
+      const topics = Object.keys(hmacKeys)
+      expect(topics.length).toBe(randomInvites.length)
+      randomInvites.forEach((invite) => {
+        expect(topics.includes(invite.conversation!.topic)).toBeTruthy()
+      })
+
+      const topicHmacs: {
+        [topic: string]: Uint8Array
+      } = {}
+      const headerBytes = new Uint8Array(10)
+
+      await Promise.all(
+        randomInvites.map(async (invite) => {
+          const topic = invite.conversation!.topic
+          const payload = new TextEncoder().encode('Hello, world!')
+
+          const {
+            responses: [encrypted],
+          } = await aliceKeystore.encryptV2({
+            requests: [
+              {
+                contentTopic: topic,
+                payload,
+                headerBytes,
+              },
+            ],
+          })
+
+          if (encrypted.error) {
+            throw encrypted.error
+          }
+
+          const topicData = aliceKeystore.lookupTopic(topic)
+          const keyMaterial = getKeyMaterial(topicData!.invitation)
+          const info = `${thirtyDayPeriodsSinceEpoch}-${aliceKeystore.walletAddress}`
+          const hmac = await generateHmacSignature(
+            keyMaterial,
+            new TextEncoder().encode(info),
+            headerBytes
+          )
+
+          topicHmacs[topic] = hmac
+        })
+      )
+
+      await Promise.all(
+        Object.keys(hmacKeys).map(async (topic) => {
+          const hmacData = hmacKeys[topic]
+
+          await Promise.all(
+            hmacData.values.map(
+              async ({ hmacKey, thirtyDayPeriodsSinceEpoch }, idx) => {
+                expect(thirtyDayPeriodsSinceEpoch).toBe(periods[idx])
+                const valid = await verifyHmacSignature(
+                  await importHmacKey(hmacKey),
+                  topicHmacs[topic],
+                  headerBytes
+                )
+                expect(valid).toBe(idx === 1 ? true : false)
+              }
+            )
+          )
+        })
+      )
+    })
   })
 })