From bcfd3f20c262bc1e265228266f1bde34c407b603 Mon Sep 17 00:00:00 2001
From: Daniel Gultsch <daniel@gultsch.de>
Date: Wed, 5 Dec 2018 10:42:56 +0100
Subject: [PATCH] postpone prekey removal and republish to after mam

---
 .../crypto/axolotl/AxolotlService.java        | 17 +++++++++------
 .../crypto/axolotl/SQLiteAxolotlStore.java    | 21 ++++++++++++++++++-
 .../persistance/DatabaseBackend.java          |  4 ++--
 3 files changed, 33 insertions(+), 9 deletions(-)

diff --git a/src/main/java/eu/siacs/conversations/crypto/axolotl/AxolotlService.java b/src/main/java/eu/siacs/conversations/crypto/axolotl/AxolotlService.java
index f12c7a086..d2de25629 100644
--- a/src/main/java/eu/siacs/conversations/crypto/axolotl/AxolotlService.java
+++ b/src/main/java/eu/siacs/conversations/crypto/axolotl/AxolotlService.java
@@ -1444,7 +1444,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) {
@@ -1494,19 +1494,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()) {
@@ -1541,7 +1546,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/eu/siacs/conversations/crypto/axolotl/SQLiteAxolotlStore.java b/src/main/java/eu/siacs/conversations/crypto/axolotl/SQLiteAxolotlStore.java
index 5838f2661..d723bd262 100644
--- a/src/main/java/eu/siacs/conversations/crypto/axolotl/SQLiteAxolotlStore.java
+++ b/src/main/java/eu/siacs/conversations/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
@@ -385,7 +388,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;
 	}
 
 	// --------------------------------------
diff --git a/src/main/java/eu/siacs/conversations/persistance/DatabaseBackend.java b/src/main/java/eu/siacs/conversations/persistance/DatabaseBackend.java
index 113b99ded..2962d41c7 100644
--- a/src/main/java/eu/siacs/conversations/persistance/DatabaseBackend.java
+++ b/src/main/java/eu/siacs/conversations/persistance/DatabaseBackend.java
@@ -1178,10 +1178,10 @@ public class DatabaseBackend extends SQLiteOpenHelper {
 		db.insert(SQLiteAxolotlStore.PREKEY_TABLENAME, null, values);
 	}
 
-	public void deletePreKey(Account account, int preKeyId) {
+	public int deletePreKey(Account account, int preKeyId) {
 		SQLiteDatabase db = this.getWritableDatabase();
 		String[] args = {account.getUuid(), Integer.toString(preKeyId)};
-		db.delete(SQLiteAxolotlStore.PREKEY_TABLENAME,
+		return db.delete(SQLiteAxolotlStore.PREKEY_TABLENAME,
 				SQLiteAxolotlStore.ACCOUNT + "=? AND "
 						+ SQLiteAxolotlStore.ID + "=?",
 				args);