Improve handling of pni change number messages

lib/src/main/java/org/asamk/signal/manager/helper/AccountHelper.java

@@ -324,10 +324,14 @@ public class AccountHelper {
             return;
         }
 
-        handlePniChangeNumberMessage(selfChangeNumber, updatePni);
+        handlePniChangeNumberMessage(selfChangeNumber, updatePni, false);
     }
 
-    public void handlePniChangeNumberMessage(final SyncMessage.PniChangeNumber pniChangeNumber, final PNI updatedPni) {
+    public boolean handlePniChangeNumberMessage(
+            final SyncMessage.PniChangeNumber pniChangeNumber,
+            final PNI updatedPni,
+            final boolean forcePniPreKeyRotation
+    ) {
         if (pniChangeNumber.identityKeyPair != null
                 && pniChangeNumber.registrationId != null
                 && pniChangeNumber.signedPreKey != null) {
@@ -341,10 +345,19 @@ public class AccountHelper {
                         pniChangeNumber.lastResortKyberPreKey != null
                                 ? new KyberPreKeyRecord(pniChangeNumber.lastResortKyberPreKey.toByteArray())
                                 : null);
+                if (forcePniPreKeyRotation) {
+                    try {
+                        context.getPreKeyHelper().forceRefreshPreKeys(ServiceIdType.PNI);
+                    } catch (IOException e) {
+                        logger.warn("Failed to force refresh PNI pre keys after PNI change sync", e);
+                    }
+                }
+                return true;
             } catch (Exception e) {
                 logger.warn("Failed to handle change number message", e);
             }
         }
+        return false;
     }
 
     public static final int USERNAME_MIN_LENGTH = 3;

lib/src/main/java/org/asamk/signal/manager/helper/IncomingMessageHandler.java

@@ -717,11 +717,30 @@ public final class IncomingMessageHandler {
         }
         if (syncMessage.getPniChangeNumber().isPresent()) {
             final var pniChangeNumber = syncMessage.getPniChangeNumber().get();
-            logger.debug("Received PNI change number sync message, applying.");
+            if (account.isPrimaryDevice()) {
-            final var updatedPniString = envelope.getUpdatedPni();
+                logger.warn("Received PNI change number sync message on primary device, ignoring.");
-            if (updatedPniString != null && !updatedPniString.isEmpty()) {
+            } else if (sender.deviceId() != SignalServiceAddress.DEFAULT_DEVICE_ID) {
-                final var updatedPni = ServiceId.PNI.parseOrThrow(updatedPniString);
+                logger.warn("Received PNI change number sync message from non-primary device {}, ignoring.",
-                context.getAccountHelper().handlePniChangeNumberMessage(pniChangeNumber, updatedPni);
+                        sender.deviceId());
+            } else {
+                final var envelopeServerTimestamp = envelope.getServerDeliveredTimestamp();
+                final var lastAppliedServerTimestamp = account.getLastAppliedPniChangeServerTimestamp();
+                if (envelopeServerTimestamp <= lastAppliedServerTimestamp) {
+                    logger.warn(
+                            "PNI change number sync server timestamp ({}) is not newer than last applied ({}), treating as replay.",
+                            envelopeServerTimestamp,
+                            lastAppliedServerTimestamp);
+                } else {
+                    final var updatedPniString = envelope.getUpdatedPni();
+                    if (updatedPniString != null && !updatedPniString.isEmpty()) {
+                        final var updatedPni = ServiceId.PNI.parseOrThrow(updatedPniString);
+                        final var applied = context.getAccountHelper()
+                                .handlePniChangeNumberMessage(pniChangeNumber, updatedPni, true);
+                        if (applied) {
+                            account.setLastAppliedPniChangeServerTimestamp(envelopeServerTimestamp);
+                        }
+                    }
+                }
             }
         }
         if (syncMessage.getDeviceNameChange().isPresent()) {

lib/src/main/java/org/asamk/signal/manager/storage/SignalAccount.java

@@ -153,6 +153,10 @@ public class SignalAccount implements Closeable {
     private final KeyValueEntry<Long> lastReceiveTimestamp = new KeyValueEntry<>("last-receive-timestamp",
             long.class,
             0L);
+    private final KeyValueEntry<Long> lastAppliedPniChangeServerTimestamp = new KeyValueEntry<>(
+            "last-applied-pni-change-server-timestamp",
+            long.class,
+            0L);
     private final KeyValueEntry<Boolean> needsToRetryFailedMessages = new KeyValueEntry<>("retry-failed-messages",
             Boolean.class,
             true);
@@ -313,6 +317,7 @@ public class SignalAccount implements Closeable {
         this.registered = false;
         this.isMultiDevice = true;
         setLastReceiveTimestamp(0L);
+        setLastAppliedPniChangeServerTimestamp(0L);
         if (accountEntropyPool != null) {
             this.pinMasterKey = null;
             this.accountEntropyPool = accountEntropyPool;
@@ -368,6 +373,7 @@ public class SignalAccount implements Closeable {
         init();
         this.registrationLockPin = pin;
         setLastReceiveTimestamp(0L);
+        setLastAppliedPniChangeServerTimestamp(0L);
         save();
 
         setPreKeys(ServiceIdType.ACI, aciPreKeys);
@@ -1753,6 +1759,14 @@ public class SignalAccount implements Closeable {
         getKeyValueStore().storeEntry(lastReceiveTimestamp, value);
     }
 
+    public long getLastAppliedPniChangeServerTimestamp() {
+        return getKeyValueStore().getEntry(lastAppliedPniChangeServerTimestamp);
+    }
+
+    public void setLastAppliedPniChangeServerTimestamp(final long value) {
+        getKeyValueStore().storeEntry(lastAppliedPniChangeServerTimestamp, value);
+    }
+
     public void setNeedsToRetryFailedMessages(final boolean value) {
         getKeyValueStore().storeEntry(needsToRetryFailedMessages, value);
     }