From 228713ebb5476ef2f2e66d3670004aaa82c0eafb Mon Sep 17 00:00:00 2001
From: AsamK <asamk@gmx.de>
Date: Tue, 18 Oct 2022 17:55:51 +0200
Subject: [PATCH] Reset pre key offset if it somehow gets corrupted

Fixes #1055
---
 .../signal/manager/helper/PreKeyHelper.java     | 12 ++++++++++++
 .../signal/manager/storage/SignalAccount.java   | 17 +++++++++++++----
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/lib/src/main/java/org/asamk/signal/manager/helper/PreKeyHelper.java b/lib/src/main/java/org/asamk/signal/manager/helper/PreKeyHelper.java
index 3fc1a5bf..30ab919e 100644
--- a/lib/src/main/java/org/asamk/signal/manager/helper/PreKeyHelper.java
+++ b/lib/src/main/java/org/asamk/signal/manager/helper/PreKeyHelper.java
@@ -53,6 +53,18 @@ public class PreKeyHelper {
         if (accountId == null) {
             return;
         }
+        try {
+            refreshPreKeys(serviceIdType, identityKeyPair);
+        } catch (Exception e) {
+            logger.warn("Failed to store new pre keys, resetting preKey id offset", e);
+            account.resetPreKeyOffsets(serviceIdType);
+            refreshPreKeys(serviceIdType, identityKeyPair);
+        }
+    }
+
+    private void refreshPreKeys(
+            final ServiceIdType serviceIdType, final IdentityKeyPair identityKeyPair
+    ) throws IOException {
         final var oneTimePreKeys = generatePreKeys(serviceIdType);
         final var signedPreKeyRecord = generateSignedPreKey(serviceIdType, identityKeyPair);
 
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 730fe85a..db919269 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
@@ -294,10 +294,8 @@ public class SignalAccount implements Closeable {
     }
 
     private void clearAllPreKeys() {
-        this.aciPreKeyIdOffset = new SecureRandom().nextInt(Medium.MAX_VALUE);
-        this.aciNextSignedPreKeyId = new SecureRandom().nextInt(Medium.MAX_VALUE);
-        this.pniPreKeyIdOffset = new SecureRandom().nextInt(Medium.MAX_VALUE);
-        this.pniNextSignedPreKeyId = new SecureRandom().nextInt(Medium.MAX_VALUE);
+        resetPreKeyOffsets(ServiceIdType.ACI);
+        resetPreKeyOffsets(ServiceIdType.PNI);
         this.getAciPreKeyStore().removeAllPreKeys();
         this.getAciSignedPreKeyStore().removeAllSignedPreKeys();
         this.getPniPreKeyStore().removeAllPreKeys();
@@ -998,6 +996,17 @@ public class SignalAccount implements Closeable {
         }
     }
 
+    public void resetPreKeyOffsets(final ServiceIdType serviceIdType) {
+        if (serviceIdType.equals(ServiceIdType.ACI)) {
+            this.aciPreKeyIdOffset = new SecureRandom().nextInt(Medium.MAX_VALUE);
+            this.aciNextSignedPreKeyId = new SecureRandom().nextInt(Medium.MAX_VALUE);
+        } else {
+            this.pniPreKeyIdOffset = new SecureRandom().nextInt(Medium.MAX_VALUE);
+            this.pniNextSignedPreKeyId = new SecureRandom().nextInt(Medium.MAX_VALUE);
+        }
+        save();
+    }
+
     public void addPreKeys(ServiceIdType serviceIdType, List<PreKeyRecord> records) {
         if (serviceIdType.equals(ServiceIdType.ACI)) {
             addAciPreKeys(records);