From: AsamK <asamk@gmx.de>
Date: Sat, 1 Jul 2023 12:06:35 +0000 (+0200)
Subject: Store last resort kyber pre key from PniChangeNumber message
X-Git-Tag: v0.12.0~21
X-Git-Url: https://git.nmode.ca/signal-cli/commitdiff_plain/ffeae1a95adf542d8b263101091952841da7ae24

Store last resort kyber pre key from PniChangeNumber message
---

diff --git a/lib/src/main/java/org/asamk/signal/manager/helper/AccountHelper.java b/lib/src/main/java/org/asamk/signal/manager/helper/AccountHelper.java
index 801e7987..4841bdbd 100644
--- a/lib/src/main/java/org/asamk/signal/manager/helper/AccountHelper.java
+++ b/lib/src/main/java/org/asamk/signal/manager/helper/AccountHelper.java
@@ -14,6 +14,7 @@ import org.asamk.signal.manager.util.NumberVerificationUtils;
 import org.asamk.signal.manager.util.Utils;
 import org.signal.libsignal.protocol.IdentityKeyPair;
 import org.signal.libsignal.protocol.InvalidKeyException;
+import org.signal.libsignal.protocol.state.KyberPreKeyRecord;
 import org.signal.libsignal.protocol.state.SignedPreKeyRecord;
 import org.signal.libsignal.usernames.BaseUsernameException;
 import org.signal.libsignal.usernames.Username;
@@ -120,20 +121,21 @@ public class AccountHelper {
         // TODO check and update remote storage
         context.getUnidentifiedAccessHelper().rotateSenderCertificates();
         dependencies.resetAfterAddressChange();
+        context.getGroupV2Helper().clearAuthCredentialCache();
         context.getAccountFileUpdater().updateAccountIdentifiers(account.getNumber(), account.getAci());
     }
 
     public void setPni(
             final PNI updatedPni,
             final IdentityKeyPair pniIdentityKeyPair,
+            final String number,
+            final int localPniRegistrationId,
             final SignedPreKeyRecord pniSignedPreKey,
-            final int localPniRegistrationId
+            final KyberPreKeyRecord lastResortKyberPreKey
     ) throws IOException {
-        account.setPni(updatedPni, pniIdentityKeyPair, pniSignedPreKey, localPniRegistrationId);
+        updateSelfIdentifiers(number != null ? number : account.getNumber(), account.getAci(), updatedPni);
+        account.setNewPniIdentity(pniIdentityKeyPair, pniSignedPreKey, lastResortKyberPreKey, localPniRegistrationId);
         context.getPreKeyHelper().refreshPreKeysIfNecessary(ServiceIdType.PNI);
-        if (account.getPni() == null || !account.getPni().equals(updatedPni)) {
-            context.getGroupV2Helper().clearAuthCredentialCache();
-        }
     }
 
     public void startChangeNumber(
diff --git a/lib/src/main/java/org/asamk/signal/manager/helper/IncomingMessageHandler.java b/lib/src/main/java/org/asamk/signal/manager/helper/IncomingMessageHandler.java
index 7cc3d09d..0692a06e 100644
--- a/lib/src/main/java/org/asamk/signal/manager/helper/IncomingMessageHandler.java
+++ b/lib/src/main/java/org/asamk/signal/manager/helper/IncomingMessageHandler.java
@@ -44,6 +44,7 @@ import org.signal.libsignal.protocol.IdentityKeyPair;
 import org.signal.libsignal.protocol.InvalidMessageException;
 import org.signal.libsignal.protocol.groups.GroupSessionBuilder;
 import org.signal.libsignal.protocol.message.DecryptionErrorMessage;
+import org.signal.libsignal.protocol.state.KyberPreKeyRecord;
 import org.signal.libsignal.protocol.state.SignedPreKeyRecord;
 import org.signal.libsignal.zkgroup.InvalidInputException;
 import org.signal.libsignal.zkgroup.profiles.ProfileKey;
@@ -646,8 +647,13 @@ public final class IncomingMessageHandler {
                     context.getAccountHelper()
                             .setPni(updatedPni,
                                     new IdentityKeyPair(pniChangeNumber.getIdentityKeyPair().toByteArray()),
+                                    pniChangeNumber.hasNewE164() ? pniChangeNumber.getNewE164() : null,
+                                    pniChangeNumber.getRegistrationId(),
                                     new SignedPreKeyRecord(pniChangeNumber.getSignedPreKey().toByteArray()),
-                                    pniChangeNumber.getRegistrationId());
+                                    pniChangeNumber.hasLastResortKyberPreKey()
+                                            ? new KyberPreKeyRecord(pniChangeNumber.getLastResortKyberPreKey()
+                                            .toByteArray())
+                                            : null);
                 } catch (Exception e) {
                     logger.warn("Failed to handle change number message", e);
                 }
diff --git a/lib/src/main/java/org/asamk/signal/manager/storage/SignalAccount.java b/lib/src/main/java/org/asamk/signal/manager/storage/SignalAccount.java
index 8cb5d7ed..c45e772f 100644
--- a/lib/src/main/java/org/asamk/signal/manager/storage/SignalAccount.java
+++ b/lib/src/main/java/org/asamk/signal/manager/storage/SignalAccount.java
@@ -1379,17 +1379,22 @@ public class SignalAccount implements Closeable {
         save();
     }
 
-    public void setPni(
-            final PNI updatedPni,
+    public void setNewPniIdentity(
             final IdentityKeyPair pniIdentityKeyPair,
             final SignedPreKeyRecord pniSignedPreKey,
+            final KyberPreKeyRecord lastResortKyberPreKey,
             final int localPniRegistrationId
     ) {
-        setPni(updatedPni);
-
         setPniIdentityKeyPair(pniIdentityKeyPair);
-        addSignedPreKey(ServiceIdType.PNI, pniSignedPreKey);
         setLocalPniRegistrationId(localPniRegistrationId);
+
+        final var preKeyMetadata = getAccountData(ServiceIdType.PNI).getPreKeyMetadata();
+        preKeyMetadata.nextSignedPreKeyId = pniSignedPreKey.getId();
+        addSignedPreKey(ServiceIdType.PNI, pniSignedPreKey);
+        if (lastResortKyberPreKey != null) {
+            preKeyMetadata.kyberPreKeyIdOffset = lastResortKyberPreKey.getId();
+            addLastResortKyberPreKey(ServiceIdType.PNI, lastResortKyberPreKey);
+        }
     }
 
     public SignalServiceAddress getSelfAddress() {