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