postpone prekey removal and republish to after mam

2 files changed, 34 insertions, 10 deletions

diff --git a/src/main/java/de/pixart/messenger/crypto/axolotl/AxolotlService.java b/src/main/java/de/pixart/messenger/crypto/axolotl/AxolotlService.java
index 34047295a..09b3b1320 100644
--- a/src/main/java/de/pixart/messenger/crypto/axolotl/AxolotlService.java
+++ b/src/main/java/de/pixart/messenger/crypto/axolotl/AxolotlService.java
@@ -1445,7 +1445,7 @@ public class AxolotlService implements OnAdvancedStreamFeaturesLoaded {
             plaintextMessage = message.decrypt(session, ownDeviceId);
             Integer preKeyId = session.getPreKeyIdAndReset();
             if (preKeyId != null) {
-                postPreKeyMessageHandling(session, preKeyId, postponePreKeyMessageHandling);
+                postPreKeyMessageHandling(session, postponePreKeyMessageHandling);
             }
         } catch (NotEncryptedForThisDeviceException e) {
             if (account.getJid().asBareJid().equals(message.getFrom().asBareJid()) && message.getSenderDeviceId() == ownDeviceId) {
@@ -1495,19 +1495,24 @@ public class AxolotlService implements OnAdvancedStreamFeaturesLoaded {
         }
     }
 
-    private void postPreKeyMessageHandling(final XmppAxolotlSession session, int preKeyId, final boolean postpone) {
+    private void postPreKeyMessageHandling(final XmppAxolotlSession session, final boolean postpone) {
         if (postpone) {
             postponedSessions.add(session);
         } else {
-            //TODO: do not republish if we already removed this preKeyId
-            publishBundlesIfNeeded(false, false);
+            if (axolotlStore.flushPreKeys()) {
+                publishBundlesIfNeeded(false, false);
+            } else {
+                Log.d(Config.LOGTAG, account.getJid().asBareJid() + ": nothing to flush. Not republishing key");
+            }
             completeSession(session);
         }
     }
 
     public void processPostponed() {
         if (postponedSessions.size() > 0) {
-            publishBundlesIfNeeded(false, false);
+            if (axolotlStore.flushPreKeys()) {
+                publishBundlesIfNeeded(false, false);
+            }
         }
         Iterator<XmppAxolotlSession> iterator = postponedSessions.iterator();
         while (iterator.hasNext()) {
@@ -1542,7 +1547,7 @@ public class AxolotlService implements OnAdvancedStreamFeaturesLoaded {
             keyTransportMessage = message.getParameters(session, getOwnDeviceId());
             Integer preKeyId = session.getPreKeyIdAndReset();
             if (preKeyId != null) {
-                postPreKeyMessageHandling(session, preKeyId, postponePreKeyMessageHandling);
+                postPreKeyMessageHandling(session, postponePreKeyMessageHandling);
             }
         } catch (CryptoFailedException e) {
             Log.d(Config.LOGTAG, "could not decrypt keyTransport message " + e.getMessage());
diff --git a/src/main/java/de/pixart/messenger/crypto/axolotl/SQLiteAxolotlStore.java b/src/main/java/de/pixart/messenger/crypto/axolotl/SQLiteAxolotlStore.java
index 0f7823d2d..6b8468af2 100644
--- a/src/main/java/de/pixart/messenger/crypto/axolotl/SQLiteAxolotlStore.java
+++ b/src/main/java/de/pixart/messenger/crypto/axolotl/SQLiteAxolotlStore.java
@@ -16,6 +16,7 @@ import org.whispersystems.libsignal.state.SignedPreKeyRecord;
 import org.whispersystems.libsignal.util.KeyHelper;
 
 import java.security.cert.X509Certificate;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
 
@@ -55,6 +56,8 @@ public class SQLiteAxolotlStore implements SignalProtocolStore {
     private int localRegistrationId;
     private int currentPreKeyId = 0;
 
+    private final HashSet<Integer> preKeysMarkedForRemoval = new HashSet<>();
+
     private final LruCache<String, FingerprintStatus> trustCache =
             new LruCache<String, FingerprintStatus>(NUM_TRUSTS_TO_CACHE) {
                 @Override
@@ -180,7 +183,7 @@ public class SQLiteAxolotlStore implements SignalProtocolStore {
      * <p/>
      * Store a remote client's identity key as trusted.
      *
-     * @param address        The address of the remote client.
+     * @param address     The address of the remote client.
      * @param identityKey The remote client's identity key.
      * @return true on success
      */
@@ -215,7 +218,6 @@ public class SQLiteAxolotlStore implements SignalProtocolStore {
      * store.  Only if it mismatches an entry in the local store is it considered
      * 'untrusted.'
      *
-     * @param name        The name of the remote client.
      * @param identityKey The identity key to verify.
      * @return true if trusted, false if untrusted.
      */
@@ -284,6 +286,7 @@ public class SQLiteAxolotlStore implements SignalProtocolStore {
                 new SignalProtocolAddress(name, 0));
     }
 
+
     public List<String> getKnownAddresses() {
         return mXmppConnectionService.databaseBackend.getKnownSignalAddresses(account);
     }
@@ -386,7 +389,23 @@ public class SQLiteAxolotlStore implements SignalProtocolStore {
      */
     @Override
     public void removePreKey(int preKeyId) {
-        mXmppConnectionService.databaseBackend.deletePreKey(account, preKeyId);
+        Log.d(Config.LOGTAG, "mark prekey for removal " + preKeyId);
+        synchronized (preKeysMarkedForRemoval) {
+            preKeysMarkedForRemoval.add(preKeyId);
+        }
+    }
+
+
+    public boolean flushPreKeys() {
+        Log.d(Config.LOGTAG, "flushing pre keys");
+        int count = 0;
+        synchronized (preKeysMarkedForRemoval) {
+            for (Integer preKeyId : preKeysMarkedForRemoval) {
+                count += mXmppConnectionService.databaseBackend.deletePreKey(account, preKeyId);
+            }
+            preKeysMarkedForRemoval.clear();
+        }
+        return count > 0;
     }
 
     // --------------------------------------
@@ -456,4 +475,4 @@ public class SQLiteAxolotlStore implements SignalProtocolStore {
     public void preVerifyFingerprint(Account account, String name, String fingerprint) {
         mXmppConnectionService.databaseBackend.storePreVerification(account, name, fingerprint, FingerprintStatus.createInactiveVerified());
     }
-}
+}
\ No newline at end of file