implement self healing omemo

after receiving a SignalMessage that can’t be decrypted because of broken sessions
Conversations will attempt to grab a new pre key bundle and send a new PreKeySignalMessage
wrapped in a key transport message.

4 files changed, 105 insertions, 14 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 d16aa7afe..34047295a 100644
--- a/src/main/java/de/pixart/messenger/crypto/axolotl/AxolotlService.java
+++ b/src/main/java/de/pixart/messenger/crypto/axolotl/AxolotlService.java
@@ -82,9 +82,11 @@ public class AxolotlService implements OnAdvancedStreamFeaturesLoaded {
     private final SerialSingleThreadExecutor executor;
     private int numPublishTriesOnEmptyPep = 0;
     private boolean pepBroken = false;
+    private final Set<SignalProtocolAddress> healingAttempts = new HashSet<>();
     private int lastDeviceListNotificationHash = 0;
     private final HashSet<Integer> cleanedOwnDeviceIds = new HashSet<>();
     private Set<XmppAxolotlSession> postponedSessions = new HashSet<>(); //sessions stored here will receive after mam catchup treatment
+    private Set<SignalProtocolAddress> postponedHealing = new HashSet<>(); //addresses stored here will need a healing notification after mam catchup
 
     private AtomicBoolean changeAccessMode = new AtomicBoolean(false);
 
@@ -390,6 +392,7 @@ public class AxolotlService implements OnAdvancedStreamFeaturesLoaded {
         this.pepBroken = false;
         this.numPublishTriesOnEmptyPep = 0;
         this.lastDeviceListNotificationHash = 0;
+        this.healingAttempts.clear();
     }
 
     public void clearErrorsInFetchStatusMap(Jid jid) {
@@ -1071,7 +1074,17 @@ public class AxolotlService implements OnAdvancedStreamFeaturesLoaded {
         }
     }
 
+    interface OnSessionBuildFromPep {
+        void onSessionBuildSuccessful();
+
+        void onSessionBuildFailed();
+    }
+
     private void buildSessionFromPEP(final SignalProtocolAddress address) {
+        buildSessionFromPEP(address, null);
+    }
+
+    private void buildSessionFromPEP(final SignalProtocolAddress address, OnSessionBuildFromPep callback) {
         Log.i(Config.LOGTAG, AxolotlService.getLogprefix(account) + "Building new session for " + address.toString());
         if (address.equals(getOwnAxolotlAddress())) {
             throw new AssertionError("We should NEVER build a session with ourselves. What happened here?!");
@@ -1092,6 +1105,9 @@ public class AxolotlService implements OnAdvancedStreamFeaturesLoaded {
                     Log.e(Config.LOGTAG, AxolotlService.getLogprefix(account) + "preKey IQ packet invalid: " + packet);
                     fetchStatusMap.put(address, FetchStatus.ERROR);
                     finishBuildingSessionsFromPEP(address);
+                    if (callback != null) {
+                        callback.onSessionBuildFailed();
+                    }
                     return;
                 }
                 Random random = new Random();
@@ -1100,6 +1116,9 @@ public class AxolotlService implements OnAdvancedStreamFeaturesLoaded {
                     //should never happen
                     fetchStatusMap.put(address, FetchStatus.ERROR);
                     finishBuildingSessionsFromPEP(address);
+                    if (callback != null) {
+                        callback.onSessionBuildFailed();
+                    }
                     return;
                 }
 
@@ -1114,7 +1133,7 @@ public class AxolotlService implements OnAdvancedStreamFeaturesLoaded {
                     XmppAxolotlSession session = new XmppAxolotlSession(account, axolotlStore, address, bundle.getIdentityKey());
                     sessions.put(address, session);
                     if (Config.X509_VERIFICATION) {
-                        verifySessionWithPEP(session);
+                        verifySessionWithPEP(session); //TODO; maybe inject callback in here too
                     } else {
                         FingerprintStatus status = getFingerprintTrust(CryptoHelper.bytesToHex(bundle.getIdentityKey().getPublicKey().serialize()));
                         FetchStatus fetchStatus;
@@ -1127,6 +1146,9 @@ public class AxolotlService implements OnAdvancedStreamFeaturesLoaded {
                         }
                         fetchStatusMap.put(address, fetchStatus);
                         finishBuildingSessionsFromPEP(address);
+                        if (callback != null) {
+                            callback.onSessionBuildSuccessful();
+                        }
                     }
                 } catch (UntrustedIdentityException | InvalidKeyException e) {
                     Log.e(Config.LOGTAG, AxolotlService.getLogprefix(account) + "Error building session for " + address + ": "
@@ -1136,6 +1158,9 @@ public class AxolotlService implements OnAdvancedStreamFeaturesLoaded {
                     if (oneOfOurs && cleanedOwnDeviceIds.add(address.getDeviceId())) {
                         removeFromDeviceAnnouncement(address.getDeviceId());
                     }
+                    if (callback != null) {
+                        callback.onSessionBuildFailed();
+                    }
                 }
             } else {
                 fetchStatusMap.put(address, FetchStatus.ERROR);
@@ -1146,6 +1171,9 @@ public class AxolotlService implements OnAdvancedStreamFeaturesLoaded {
                 if (oneOfOurs && itemNotFound && cleanedOwnDeviceIds.add(address.getDeviceId())) {
                     removeFromDeviceAnnouncement(address.getDeviceId());
                 }
+                if (callback != null) {
+                    callback.onSessionBuildFailed();
+                }
             }
         });
     }
@@ -1391,11 +1419,15 @@ public class AxolotlService implements OnAdvancedStreamFeaturesLoaded {
     }
 
     private XmppAxolotlSession getReceivingSession(XmppAxolotlMessage message) {
-        SignalProtocolAddress senderAddress = new SignalProtocolAddress(message.getFrom().toString(),
-                message.getSenderDeviceId());
+        SignalProtocolAddress senderAddress = new SignalProtocolAddress(message.getFrom().toString(), message.getSenderDeviceId());
+        return getReceivingSession(senderAddress);
+
+    }
+
+    private XmppAxolotlSession getReceivingSession(SignalProtocolAddress senderAddress) {
         XmppAxolotlSession session = sessions.get(senderAddress);
         if (session == null) {
-            Log.d(Config.LOGTAG, AxolotlService.getLogprefix(account) + "Account: " + account.getJid() + " No axolotl session found while parsing received message " + message);
+            //Log.d(Config.LOGTAG, AxolotlService.getLogprefix(account) + "Account: " + account.getJid() + " No axolotl session found while parsing received message " + message);
             session = recreateUncachedSession(senderAddress);
             if (session == null) {
                 session = new XmppAxolotlSession(account, axolotlStore, senderAddress);
@@ -1404,7 +1436,7 @@ public class AxolotlService implements OnAdvancedStreamFeaturesLoaded {
         return session;
     }
 
-    public XmppAxolotlMessage.XmppAxolotlPlaintextMessage processReceivingPayloadMessage(XmppAxolotlMessage message, boolean postponePreKeyMessageHandling) throws NotEncryptedForThisDeviceException {
+    public XmppAxolotlMessage.XmppAxolotlPlaintextMessage processReceivingPayloadMessage(XmppAxolotlMessage message, boolean postponePreKeyMessageHandling) throws NotEncryptedForThisDeviceException, BrokenSessionException {
         XmppAxolotlMessage.XmppAxolotlPlaintextMessage plaintextMessage = null;
 
         XmppAxolotlSession session = getReceivingSession(message);
@@ -1421,8 +1453,10 @@ public class AxolotlService implements OnAdvancedStreamFeaturesLoaded {
             } else {
                 throw e;
             }
+        } catch (final BrokenSessionException e) {
+            throw e;
         } catch (CryptoFailedException e) {
-            Log.w(Config.LOGTAG, getLogprefix(account) + "Failed to decrypt message from " + message.getFrom() + ": " + e.getMessage());
+            Log.w(Config.LOGTAG, getLogprefix(account) + "Failed to decrypt message from " + message.getFrom(), e);
         }
 
         if (session.isFresh() && plaintextMessage != null) {
@@ -1432,6 +1466,35 @@ public class AxolotlService implements OnAdvancedStreamFeaturesLoaded {
         return plaintextMessage;
     }
 
+    public void reportBrokenSessionException(BrokenSessionException e, boolean postpone) {
+        Log.e(Config.LOGTAG, account.getJid().asBareJid() + ": broken session with " + e.getSignalProtocolAddress().toString() + " detected", e);
+        if (postpone) {
+            postponedHealing.add(e.getSignalProtocolAddress());
+        } else {
+            notifyRequiresHealing(e.getSignalProtocolAddress());
+        }
+    }
+
+    private void notifyRequiresHealing(final SignalProtocolAddress signalProtocolAddress) {
+        if (healingAttempts.add(signalProtocolAddress)) {
+            Log.d(Config.LOGTAG, account.getJid().asBareJid() + ": attempt to heal " + signalProtocolAddress);
+            buildSessionFromPEP(signalProtocolAddress, new OnSessionBuildFromPep() {
+                @Override
+                public void onSessionBuildSuccessful() {
+                    Log.d(Config.LOGTAG, "successfully build new session from pep after detecting broken session");
+                    completeSession(getReceivingSession(signalProtocolAddress));
+                }
+
+                @Override
+                public void onSessionBuildFailed() {
+                    Log.d(Config.LOGTAG, account.getJid().asBareJid() + ": unable to build new session from pep after detecting broken session");
+                }
+            });
+        } else {
+            Log.d(Config.LOGTAG, account.getJid().asBareJid() + ": do not attempt to heal " + signalProtocolAddress + " again");
+        }
+    }
+
     private void postPreKeyMessageHandling(final XmppAxolotlSession session, int preKeyId, final boolean postpone) {
         if (postpone) {
             postponedSessions.add(session);
@@ -1451,6 +1514,11 @@ public class AxolotlService implements OnAdvancedStreamFeaturesLoaded {
             completeSession(iterator.next());
             iterator.remove();
         }
+        Iterator<SignalProtocolAddress> postponedHealingAttemptsIterator = postponedHealing.iterator();
+        while (postponedHealingAttemptsIterator.hasNext()) {
+            notifyRequiresHealing(postponedHealingAttemptsIterator.next());
+            postponedHealingAttemptsIterator.remove();
+        }
     }
 
     private void completeSession(XmppAxolotlSession session) {
diff --git a/src/main/java/de/pixart/messenger/crypto/axolotl/BrokenSessionException.java b/src/main/java/de/pixart/messenger/crypto/axolotl/BrokenSessionException.java
new file mode 100644
index 000000000..1193ad8e4
--- /dev/null
+++ b/src/main/java/de/pixart/messenger/crypto/axolotl/BrokenSessionException.java
@@ -0,0 +1,18 @@
+package de.pixart.messenger.crypto.axolotl;
+
+import org.whispersystems.libsignal.SignalProtocolAddress;
+
+public class BrokenSessionException extends CryptoFailedException {
+
+    private final SignalProtocolAddress signalProtocolAddress;
+
+    public BrokenSessionException(SignalProtocolAddress address, Exception e) {
+        super(e);
+        this.signalProtocolAddress = address;
+
+    }
+
+    public SignalProtocolAddress getSignalProtocolAddress() {
+        return signalProtocolAddress;
+    }
+}
\ No newline at end of file
diff --git a/src/main/java/de/pixart/messenger/crypto/axolotl/CryptoFailedException.java b/src/main/java/de/pixart/messenger/crypto/axolotl/CryptoFailedException.java
index e6f5e2a65..3933df006 100644
--- a/src/main/java/de/pixart/messenger/crypto/axolotl/CryptoFailedException.java
+++ b/src/main/java/de/pixart/messenger/crypto/axolotl/CryptoFailedException.java
@@ -6,6 +6,10 @@ public class CryptoFailedException extends Exception {
         super(msg);
     }
 
+    public CryptoFailedException(String msg, Exception e) {
+        super(msg, e);
+    }
+
     public CryptoFailedException(Exception e) {
         super(e);
     }
diff --git a/src/main/java/de/pixart/messenger/crypto/axolotl/XmppAxolotlSession.java b/src/main/java/de/pixart/messenger/crypto/axolotl/XmppAxolotlSession.java
index caa7f9c23..cbfdaf28e 100644
--- a/src/main/java/de/pixart/messenger/crypto/axolotl/XmppAxolotlSession.java
+++ b/src/main/java/de/pixart/messenger/crypto/axolotl/XmppAxolotlSession.java
@@ -79,7 +79,7 @@ public class XmppAxolotlSession implements Comparable<XmppAxolotlSession> {
     }
 
     @Nullable
-    public byte[] processReceiving(AxolotlKey encryptedKey) throws CryptoFailedException {
+    byte[] processReceiving(AxolotlKey encryptedKey) throws CryptoFailedException {
         byte[] plaintext;
         FingerprintStatus status = getTrust();
         if (!status.isCompromised()) {
@@ -99,21 +99,22 @@ public class XmppAxolotlSession implements Comparable<XmppAxolotlSession> {
                     plaintext = cipher.decrypt(preKeySignalMessage);
                 } else {
                     SignalMessage signalMessage = new SignalMessage(encryptedKey.key);
-                    plaintext = cipher.decrypt(signalMessage);
+                    try {
+                        plaintext = cipher.decrypt(signalMessage);
+                    } catch (InvalidMessageException | NoSessionException e) {
+                        throw new BrokenSessionException(this.remoteAddress, e);
+                    }
                     preKeyId = null; //better safe than sorry because we use that to do special after prekey handling
                 }
-            } catch (InvalidVersionException | InvalidKeyException | LegacyMessageException | InvalidMessageException | DuplicateMessageException | NoSessionException | InvalidKeyIdException | UntrustedIdentityException e) {
-                if (!(e instanceof DuplicateMessageException)) {
-                    e.printStackTrace();
-                }
-                throw new CryptoFailedException("Error decrypting WhisperMessage " + e.getClass().getSimpleName() + ": " + e.getMessage());
+            } catch (InvalidVersionException | InvalidKeyException | LegacyMessageException | InvalidMessageException | DuplicateMessageException | InvalidKeyIdException | UntrustedIdentityException e) {
+                throw new CryptoFailedException("Error decrypting SignalMessage", e);
             }
             if (!status.isActive()) {
                 setTrust(status.toActive());
                 //TODO: also (re)add to device list?
             }
         } else {
-            throw new CryptoFailedException("not encrypting omemo message from fingerprint "+getFingerprint()+" because it was marked as compromised");
+            throw new CryptoFailedException("not encrypting omemo message from fingerprint " + getFingerprint() + " because it was marked as compromised");
         }
         return plaintext;
     }