From 1bcbf27c83fa7e6f629f1fbeafc1adb83e1ebcd2 Mon Sep 17 00:00:00 2001
From: "Alexander \"PapaTutuWawa" <papatutuwawa@polynom.me>
Date: Mon, 8 Aug 2022 15:48:09 +0200
Subject: [PATCH] feat: Make any SPK survive exactly one rotation

---
 lib/src/errors.dart               |  6 ++++++
 lib/src/omemo/sessionmanager.dart | 21 +++++++++++++++----
 test/omemo_test.dart              | 35 +++++++++++++++++++++++++++++++
 3 files changed, 58 insertions(+), 4 deletions(-)

diff --git a/lib/src/errors.dart b/lib/src/errors.dart
index 8398381..67cf159 100644
--- a/lib/src/errors.dart
+++ b/lib/src/errors.dart
@@ -24,3 +24,9 @@ class NotEncryptedForDeviceException implements Exception {
 class NoDecryptionKeyException implements Exception {
   String errMsg() => 'No key available for decrypting the message';
 }
+
+/// Triggered by the Session Manager when the identifier of the used Signed Prekey
+/// is neither the current SPK's identifier nor the old one's.
+class UnknownSignedPrekeyException implements Exception {
+  String errMsg() => 'Unknown Signed Prekey used.';
+}
diff --git a/lib/src/omemo/sessionmanager.dart b/lib/src/omemo/sessionmanager.dart
index 8cba9ca..a76b77e 100644
--- a/lib/src/omemo/sessionmanager.dart
+++ b/lib/src/omemo/sessionmanager.dart
@@ -156,27 +156,40 @@ class OmemoSessionManager {
 
     return OmemoKeyExchange()
       ..pkId = kexResult.opkId
-      ..spkId = 0
+      ..spkId = bundle.spkId
       ..ik = await device.ik.pk.getBytes()
       ..ek = await kexResult.ek.pk.getBytes();
   }
 
   /// Build a new session with the user at [jid] with the device [deviceId] using data
-  /// from the key exchange [kex].
+  /// from the key exchange [kex]. In case [kex] contains an unknown Signed Prekey
+  /// identifier an UnknownSignedPrekeyException will be thrown.
   Future<void> _addSessionFromKeyExchange(String jid, int deviceId, OmemoKeyExchange kex) async {
+    // Pick the correct SPK
     final device = await getDevice();
+    OmemoKeyPair? spk;
+    if (kex.spkId == device.spkId) {
+      spk = device.spk;
+    } else if (kex.spkId == device.oldSpkId) {
+      spk = device.oldSpk;
+    } else {
+      throw UnknownSignedPrekeyException();
+    }
+
+    assert(spk != null, 'The used SPK must be found');
+    
     final kexResult = await x3dhFromInitialMessage(
       X3DHMessage(
         OmemoPublicKey.fromBytes(kex.ik!, KeyPairType.ed25519),
         OmemoPublicKey.fromBytes(kex.ek!, KeyPairType.x25519),
         kex.pkId!,
       ),
-      device.spk,
+      spk!,
       device.opks.values.elementAt(kex.pkId!),
       device.ik,
     );
     final ratchet = await OmemoDoubleRatchet.acceptNewSession(
-      device.spk,
+      spk,
       OmemoPublicKey.fromBytes(kex.ik!, KeyPairType.ed25519),
       kexResult.sk,
       kexResult.ad,
diff --git a/test/omemo_test.dart b/test/omemo_test.dart
index 0689b4e..a66cc8d 100644
--- a/test/omemo_test.dart
+++ b/test/omemo_test.dart
@@ -249,4 +249,39 @@ void main() {
     expect(newDevice!.oldSpkId, oldDevice.spkId);
     expect(listsEqual(newDevice!.oldSpkSignature!, oldDevice.spkSignature), true);
   });
+
+  test('Test accepting a session with an old SPK', () async {
+    const aliceJid = 'alice@server.example';
+    const bobJid = 'bob@other.server.example';
+      
+    // Alice and Bob generate their sessions
+    final aliceSession = await OmemoSessionManager.generateNewIdentity(aliceJid, opkAmount: 1);
+    final bobSession = await OmemoSessionManager.generateNewIdentity(bobJid, opkAmount: 1);
+
+    // Alice encrypts a message for Bob
+    const messagePlaintext = 'Hello Bob!';
+    final aliceMessage = await aliceSession.encryptToJid(
+      bobJid,
+      messagePlaintext,
+      newSessions: [
+        await (await bobSession.getDevice()).toBundle(),
+      ],
+    );
+    expect(aliceMessage.encryptedKeys.length, 1);
+
+    // Alice loses her Internet connection. Bob rotates his SPK.
+    await bobSession.rotateSignedPrekey();
+    
+    // Alice regains her Internet connection and sends the message to Bob
+    // ...
+
+    // Bob decrypts it
+    final bobMessage = await bobSession.decryptMessage(
+      aliceMessage.ciphertext,
+      aliceJid,
+      (await aliceSession.getDevice()).id,
+      aliceMessage.encryptedKeys,
+    );
+    expect(messagePlaintext, bobMessage);
+  });
 }