Skip to content

Commit 7bd5abf

Browse files
neekolasneekolas
authored
Merge pull request #162 from xmtp/key-manager
Key manager
2 parents 9cd2a81 + 2c70b9d commit 7bd5abf

File tree

2 files changed

+252
-0
lines changed

2 files changed

+252
-0
lines changed

‎src/TopicKeyManager.ts

+136
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,136 @@
1+
import PublicKeyBundle from './crypto/PublicKeyBundle'
2+
3+
export enum EncryptionAlgorithm {
4+
AES_256_GCM_HKDF_SHA_256,
5+
}
6+
7+
/**
8+
* TopicKeyRecord encapsulates the key, algorithm, and a list of allowed signers
9+
*/
10+
export type TopicKeyRecord = {
11+
keyMaterial: Uint8Array
12+
encryptionAlgorithm: EncryptionAlgorithm
13+
// Callers should validate that the signature comes from the list of allowed signers
14+
// Not strictly necessary, but it prevents against compromised topic keys being
15+
// used by third parties who would sign the message with a different key
16+
allowedSigners: PublicKeyBundle[]
17+
}
18+
19+
/**
20+
* TopicResult is the public interface for receiving a TopicKey
21+
*/
22+
export type TopicResult = {
23+
topicKey: TopicKeyRecord
24+
contentTopic: string
25+
}
26+
27+
// Internal data structure used to store the relationship between a topic and a wallet address
28+
type WalletTopicRecord = {
29+
contentTopic: string
30+
createdAt: Date
31+
}
32+
33+
type ContentTopic = string
34+
type WalletAddress = string
35+
36+
/**
37+
* Custom error type for cases where the caller attempted to add a second key to the same topic
38+
*/
39+
export class DuplicateTopicError extends Error {
40+
constructor(topic: string) {
41+
super(`Topic ${topic} has already been added`)
42+
this.name = 'DuplicateTopicError'
43+
Object.setPrototypeOf(this, DuplicateTopicError.prototype)
44+
}
45+
}
46+
47+
const findLatestTopic = (records: WalletTopicRecord[]): WalletTopicRecord => {
48+
let latestRecord: WalletTopicRecord | undefined
49+
for (const record of records) {
50+
if (!latestRecord || record.createdAt > latestRecord.createdAt) {
51+
latestRecord = record
52+
}
53+
}
54+
if (!latestRecord) {
55+
throw new Error('No record found')
56+
}
57+
return latestRecord
58+
}
59+
60+
/**
61+
* TopicKeyManager stores the mapping between topics -> keys and wallet addresses -> keys
62+
*/
63+
export default class TopicKeyManager {
64+
// Mapping of content topics to the keys used for decryption on that topic
65+
private topicKeys: Map<ContentTopic, TopicKeyRecord>
66+
// Mapping of wallet addresses and topics
67+
private dmTopics: Map<WalletAddress, WalletTopicRecord[]>
68+
69+
constructor() {
70+
this.topicKeys = new Map<ContentTopic, TopicKeyRecord>()
71+
this.dmTopics = new Map<WalletAddress, WalletTopicRecord[]>()
72+
}
73+
74+
/**
75+
* Create a TopicKeyRecord for the topic and store it for later access
76+
*
77+
* @param contentTopic The topic
78+
* @param key TopicKeyRecord that contains the topic key and encryption algorithm
79+
* @param counterparty The other user's PublicKeyBundle
80+
* @param createdAt Date
81+
*/
82+
addDirectMessageTopic(
83+
contentTopic: string,
84+
key: TopicKeyRecord,
85+
counterparty: PublicKeyBundle,
86+
createdAt: Date
87+
): void {
88+
if (this.topicKeys.has(contentTopic)) {
89+
throw new DuplicateTopicError(contentTopic)
90+
}
91+
this.topicKeys.set(contentTopic, key)
92+
93+
const walletAddress = counterparty.identityKey.walletSignatureAddress()
94+
const counterpartyTopicList = this.dmTopics.get(walletAddress) || []
95+
counterpartyTopicList.push({ contentTopic, createdAt })
96+
this.dmTopics.set(walletAddress, counterpartyTopicList)
97+
}
98+
99+
/**
100+
* Would be used to get all information required to decrypt/validate a given message
101+
*/
102+
getByTopic(contentTopic: string): TopicResult | undefined {
103+
const topicKey = this.topicKeys.get(contentTopic)
104+
if (!topicKey) {
105+
return undefined
106+
}
107+
return {
108+
topicKey,
109+
contentTopic,
110+
}
111+
}
112+
113+
/**
114+
* Used to know which topic/key to use to send to a given wallet address
115+
*/
116+
getLatestByWalletAddress(walletAddress: string): TopicResult | undefined {
117+
const walletTopics = this.dmTopics.get(walletAddress)
118+
if (!walletTopics || !walletTopics.length) {
119+
return undefined
120+
}
121+
const newestTopic = findLatestTopic(walletTopics)
122+
return this.getByTopic(newestTopic.contentTopic)
123+
}
124+
125+
/**
126+
* Used to get the topic list to listen for all messages from a given wallet address
127+
*/
128+
getAllByWalletAddress(walletAddress: string): TopicResult[] {
129+
const dmTopics = this.dmTopics
130+
.get(walletAddress)
131+
?.map(({ contentTopic }) => this.getByTopic(contentTopic))
132+
.filter((res) => !!res) as TopicResult[]
133+
134+
return dmTopics || []
135+
}
136+
}

‎test/TopicKeyManager.test.ts

+116
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
import {
2+
TopicKeyRecord,
3+
EncryptionAlgorithm,
4+
DuplicateTopicError,
5+
} from './../src/TopicKeyManager'
6+
import KeyManager from '../src/TopicKeyManager'
7+
import { crypto } from '../src/crypto/encryption'
8+
import PrivateKeyBundle from '../src/crypto/PrivateKeyBundle'
9+
import PublicKeyBundle from '../src/crypto/PublicKeyBundle'
10+
import { newWallet } from './helpers'
11+
12+
const TOPICS = ['topic1', 'topic2']
13+
14+
const createTopicKeyRecord = (
15+
allowedSigners: PublicKeyBundle[] = []
16+
): TopicKeyRecord => {
17+
return {
18+
keyMaterial: crypto.getRandomValues(new Uint8Array(32)),
19+
encryptionAlgorithm: EncryptionAlgorithm.AES_256_GCM_HKDF_SHA_256,
20+
allowedSigners: allowedSigners,
21+
}
22+
}
23+
24+
describe('TopicKeyManager', () => {
25+
let keyManager: KeyManager
26+
beforeEach(async () => {
27+
keyManager = new KeyManager()
28+
})
29+
30+
it('can add and retrieve a topic key', async () => {
31+
const senderWallet = newWallet()
32+
const sender = (
33+
await PrivateKeyBundle.generate(senderWallet)
34+
).getPublicKeyBundle()
35+
const sentAt = new Date()
36+
const record = createTopicKeyRecord([sender])
37+
38+
keyManager.addDirectMessageTopic(TOPICS[0], record, sender, sentAt)
39+
40+
// Lookup latest result by address
41+
const topicResultByAddress = keyManager.getLatestByWalletAddress(
42+
senderWallet.address
43+
)
44+
expect(topicResultByAddress?.contentTopic).toEqual(TOPICS[0])
45+
expect(topicResultByAddress?.topicKey.keyMaterial).toEqual(
46+
record.keyMaterial
47+
)
48+
49+
// Lookup all results by address
50+
const allResults = keyManager.getAllByWalletAddress(senderWallet.address)
51+
expect(allResults).toHaveLength(1)
52+
expect(allResults[0]).toEqual(topicResultByAddress)
53+
54+
// Lookup result by topic
55+
const topicResultByTopic = keyManager.getByTopic(TOPICS[0])
56+
expect(topicResultByTopic).toEqual(topicResultByAddress)
57+
})
58+
59+
it('returns undefined when no topic key has been added', async () => {
60+
expect(keyManager.getByTopic(TOPICS[0])).toBeUndefined()
61+
expect(keyManager.getLatestByWalletAddress(TOPICS[0])).toBeUndefined()
62+
expect(keyManager.getAllByWalletAddress(TOPICS[0])).toHaveLength(0)
63+
})
64+
65+
it('can add multiple topic keys for a wallet', async () => {
66+
const senderWallet = newWallet()
67+
const sender = (
68+
await PrivateKeyBundle.generate(senderWallet)
69+
).getPublicKeyBundle()
70+
const record1 = createTopicKeyRecord([sender])
71+
const record2 = createTopicKeyRecord([sender])
72+
const d1 = new Date(+new Date() - 100)
73+
const d2 = new Date()
74+
75+
keyManager.addDirectMessageTopic(TOPICS[0], record1, sender, d1)
76+
keyManager.addDirectMessageTopic(TOPICS[1], record2, sender, d2)
77+
78+
// Should use the record with the latest date
79+
expect(
80+
keyManager.getLatestByWalletAddress(senderWallet.address)?.contentTopic
81+
).toEqual(TOPICS[1])
82+
83+
// Should return both results
84+
expect(keyManager.getAllByWalletAddress(senderWallet.address)).toHaveLength(
85+
2
86+
)
87+
88+
// It should still be possible to look up the older topic using the topic name directly
89+
expect(keyManager.getByTopic(TOPICS[0])?.topicKey.keyMaterial).toEqual(
90+
record1.keyMaterial
91+
)
92+
})
93+
94+
it('cannot add multiple records for the same topic', async () => {
95+
const senderWallet = newWallet()
96+
const sender = (
97+
await PrivateKeyBundle.generate(senderWallet)
98+
).getPublicKeyBundle()
99+
100+
keyManager.addDirectMessageTopic(
101+
TOPICS[0],
102+
createTopicKeyRecord([sender]),
103+
sender,
104+
new Date()
105+
)
106+
107+
expect(() =>
108+
keyManager.addDirectMessageTopic(
109+
TOPICS[0],
110+
createTopicKeyRecord([sender]),
111+
sender,
112+
new Date()
113+
)
114+
).toThrow(DuplicateTopicError)
115+
})
116+
})

0 commit comments

Comments
 (0)