Fix some tests broken by kyber prekeys

[skip cd]
mautrix · Dec 28, 2023 · 175fde6 · 175fde6
1 parent fb18dce
commit 175fde6
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 4 deletions.
diff --git a/pkg/libsignalgo/inmemorystore_test.go b/pkg/libsignalgo/inmemorystore_test.go
@@ -49,6 +49,7 @@ type InMemorySignalProtocolStore struct {
 	senderKeyMap    map[SenderKeyName]*libsignalgo.SenderKeyRecord
 	sessionMap      map[AddressKey]*libsignalgo.SessionRecord
 	signedPreKeyMap map[uint32]*libsignalgo.SignedPreKeyRecord
+	kyberPreKeyMap  map[uint32]*libsignalgo.KyberPreKeyRecord
 }
 
 func NewInMemorySignalProtocolStore() *InMemorySignalProtocolStore {
@@ -71,6 +72,7 @@ func NewInMemorySignalProtocolStore() *InMemorySignalProtocolStore {
 		senderKeyMap:    make(map[SenderKeyName]*libsignalgo.SenderKeyRecord),
 		sessionMap:      make(map[AddressKey]*libsignalgo.SessionRecord),
 		signedPreKeyMap: make(map[uint32]*libsignalgo.SignedPreKeyRecord),
+		kyberPreKeyMap:  make(map[uint32]*libsignalgo.KyberPreKeyRecord),
 	}
 }
 
@@ -258,19 +260,23 @@ func (ps *BadInMemorySignalProtocolStore) LoadPreKey(id uint32, ctx context.Cont
 	return nil, errors.New("Test error")
 }
 
+func (ps *BadInMemorySignalProtocolStore) LoadKyberPreKey(id uint32, ctx context.Context) (*libsignalgo.KyberPreKeyRecord, error) {
+	return nil, errors.New("Test error")
+}
+
 // Implementation of the KyberPreKeyStore interface
 // TODO: this is just stubs, not implemented yet
 
 func (ps *InMemorySignalProtocolStore) LoadKyberPreKey(id uint32, ctx context.Context) (*libsignalgo.KyberPreKeyRecord, error) {
-	return nil, nil //ps.kyberPreKeyMap[id], nil
+	return ps.kyberPreKeyMap[id], nil
 }
 
 func (ps *InMemorySignalProtocolStore) StoreKyberPreKey(id uint32, kyberPreKeyRecord *libsignalgo.KyberPreKeyRecord, ctx context.Context) error {
-	//ps.kyberPreKeyMap[id] = kyberPreKeyRecord
+	ps.kyberPreKeyMap[id] = kyberPreKeyRecord
 	return nil
 }
 
 func (ps *InMemorySignalProtocolStore) MarkKyberPreKeyUsed(id uint32, ctx context.Context) error {
-	//ps.kyberPreKeyMap[id].MarkUsed()
+	//delete(ps.kyberPreKeyMap, id)
 	return nil
 }
diff --git a/pkg/libsignalgo/session_test.go b/pkg/libsignalgo/session_test.go
@@ -108,7 +108,12 @@ func initializeSessions(t *testing.T, aliceStore, bobStore *InMemorySignalProtoc
 	signedPreKeyRecord, err := libsignalgo.NewSignedPreKeyRecordFromPrivateKey(signedPreKeyID, time.UnixMilli(42000), bobSignedPreKey, bobSignedPreKeySignature)
 	assert.NoError(t, err)
 	err = bobStore.StoreSignedPreKey(signedPreKeyID, signedPreKeyRecord, ctx)
-	assert.NoError(t, err)
+	require.NoError(t, err)
+
+	kyberPreKeyRecord, err := libsignalgo.NewKyberPreKeyRecord(kyberPreKeyId, time.UnixMilli(42000), bobKyberPreKey, bobKyberPreKeySignature)
+	require.NoError(t, err)
+	err = bobStore.StoreKyberPreKey(kyberPreKeyId, kyberPreKeyRecord, ctx)
+	require.NoError(t, err)
 }
 
 // From SessionTests.swift:testSessionCipher