handle decrypting/encrypting of omemo messages with duplicate device ids

2 files changed, 77 insertions, 43 deletions

diff --git a/src/main/java/de/pixart/messenger/crypto/axolotl/XmppAxolotlMessage.java b/src/main/java/de/pixart/messenger/crypto/axolotl/XmppAxolotlMessage.java
index 71fbd40db..100ad1c43 100644
--- a/src/main/java/de/pixart/messenger/crypto/axolotl/XmppAxolotlMessage.java
+++ b/src/main/java/de/pixart/messenger/crypto/axolotl/XmppAxolotlMessage.java
@@ -9,6 +9,7 @@ import java.security.InvalidKeyException;
 import java.security.NoSuchAlgorithmException;
 import java.security.NoSuchProviderException;
 import java.security.SecureRandom;
+import java.util.ArrayList;
 import java.util.List;
 
 import javax.crypto.BadPaddingException;
@@ -42,7 +43,7 @@ public class XmppAxolotlMessage {
     private byte[] ciphertext = null;
     private byte[] authtagPlusInnerKey = null;
     private byte[] iv = null;
-    private final SparseArray<XmppAxolotlSession.AxolotlKey> keys;
+    private final List<XmppAxolotlSession.AxolotlKey> keys;
     private final Jid from;
     private final int sourceDeviceId;
 
@@ -110,7 +111,7 @@ public class XmppAxolotlMessage {
             throw new IllegalArgumentException("invalid source id");
         }
         List<Element> keyElements = header.getChildren();
-        this.keys = new SparseArray<>();
+        this.keys = new ArrayList<>();
         for (Element keyElement : keyElements) {
             switch (keyElement.getName()) {
                 case KEYTAG:
@@ -118,7 +119,7 @@ public class XmppAxolotlMessage {
                         Integer recipientId = Integer.parseInt(keyElement.getAttribute(REMOTEID));
                         byte[] key = Base64.decode(keyElement.getContent().trim(), Base64.DEFAULT);
                         boolean isPreKey = keyElement.getAttributeAsBoolean("prekey");
-                        this.keys.put(recipientId, new XmppAxolotlSession.AxolotlKey(key, isPreKey));
+                        this.keys.add(new XmppAxolotlSession.AxolotlKey(recipientId, key, isPreKey));
                     } catch (NumberFormatException e) {
                         throw new IllegalArgumentException("invalid remote id");
                     }
@@ -143,7 +144,7 @@ public class XmppAxolotlMessage {
     XmppAxolotlMessage(Jid from, int sourceDeviceId) {
         this.from = from;
         this.sourceDeviceId = sourceDeviceId;
-        this.keys = new SparseArray<>();
+        this.keys = new ArrayList<>();
         this.iv = generateIv();
         this.innerKey = generateKey();
     }
@@ -198,15 +199,15 @@ public class XmppAxolotlMessage {
 
     private static byte[] getPaddedBytes(String plaintext) {
         int plainLength = plaintext.getBytes().length;
-        int pad = Math.max(64,(plainLength / 32 + 1) * 32) - plainLength;
+        int pad = Math.max(64, (plainLength / 32 + 1) * 32) - plainLength;
         SecureRandom random = new SecureRandom();
         int left = random.nextInt(pad);
         int right = pad - left;
         StringBuilder builder = new StringBuilder(plaintext);
-        for(int i = 0; i < left; ++i) {
-            builder.insert(0,random.nextBoolean() ? "\t" : " ");
+        for (int i = 0; i < left; ++i) {
+            builder.insert(0, random.nextBoolean() ? "\t" : " ");
         }
-        for(int i = 0; i < right; ++i) {
+        for (int i = 0; i < right; ++i) {
             builder.append(random.nextBoolean() ? "\t" : " ");
         }
         return builder.toString().getBytes();
@@ -232,7 +233,7 @@ public class XmppAxolotlMessage {
             key = session.processSending(innerKey, ignoreSessionTrust);
         }
         if (key != null) {
-            keys.put(session.getRemoteAddress().getDeviceId(), key);
+            keys.add(key);
         }
     }
 
@@ -248,13 +249,13 @@ public class XmppAxolotlMessage {
         Element encryptionElement = new Element(CONTAINERTAG, AxolotlService.PEP_PREFIX);
         Element headerElement = encryptionElement.addChild(HEADER);
         headerElement.setAttribute(SOURCEID, sourceDeviceId);
-        for(int i = 0; i < keys.size(); ++i) {
+        for (XmppAxolotlSession.AxolotlKey key : keys) {
             Element keyElement = new Element(KEYTAG);
-            keyElement.setAttribute(REMOTEID, keys.keyAt(i));
-            if (keys.valueAt(i).prekey) {
+            keyElement.setAttribute(REMOTEID, key.deviceId);
+            if (key.prekey) {
                 keyElement.setAttribute("prekey", "true");
             }
-            keyElement.setContent(Base64.encodeToString(keys.valueAt(i).key, Base64.NO_WRAP));
+            keyElement.setContent(Base64.encodeToString(key.key, Base64.NO_WRAP));
             headerElement.addChild(keyElement);
         }
         headerElement.addChild(IVTAG).setContent(Base64.encodeToString(iv, Base64.NO_WRAP));
@@ -266,11 +267,16 @@ public class XmppAxolotlMessage {
     }
 
     private byte[] unpackKey(XmppAxolotlSession session, Integer sourceDeviceId) throws CryptoFailedException {
-        XmppAxolotlSession.AxolotlKey encryptedKey = keys.get(sourceDeviceId);
-        if (encryptedKey == null) {
+        ArrayList<XmppAxolotlSession.AxolotlKey> possibleKeys = new ArrayList<>();
+        for (XmppAxolotlSession.AxolotlKey key : keys) {
+            if (key.deviceId == sourceDeviceId) {
+                possibleKeys.add(key);
+            }
+        }
+        if (possibleKeys.size() == 0) {
             throw new NotEncryptedForThisDeviceException();
         }
-        return session.processReceiving(encryptedKey);
+        return session.processReceiving(possibleKeys);
     }
 
     XmppAxolotlKeyTransportMessage getParameters(XmppAxolotlSession session, Integer sourceDeviceId) throws CryptoFailedException {
@@ -311,4 +317,4 @@ public class XmppAxolotlMessage {
         }
         return plaintextMessage;
     }
-}
+}
\ No newline at end of file
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 cbfdaf28e..71e20676e 100644
--- a/src/main/java/de/pixart/messenger/crypto/axolotl/XmppAxolotlSession.java
+++ b/src/main/java/de/pixart/messenger/crypto/axolotl/XmppAxolotlSession.java
@@ -2,6 +2,7 @@ package de.pixart.messenger.crypto.axolotl;
 
 import android.support.annotation.NonNull;
 import android.support.annotation.Nullable;
+import android.util.Log;
 
 import org.whispersystems.libsignal.DuplicateMessageException;
 import org.whispersystems.libsignal.IdentityKey;
@@ -19,6 +20,10 @@ import org.whispersystems.libsignal.protocol.PreKeySignalMessage;
 import org.whispersystems.libsignal.protocol.SignalMessage;
 import org.whispersystems.libsignal.util.guava.Optional;
 
+import java.util.Iterator;
+import java.util.List;
+
+import de.pixart.messenger.Config;
 import de.pixart.messenger.entities.Account;
 import de.pixart.messenger.utils.CryptoHelper;
 
@@ -79,35 +84,56 @@ public class XmppAxolotlSession implements Comparable<XmppAxolotlSession> {
     }
 
     @Nullable
-    byte[] processReceiving(AxolotlKey encryptedKey) throws CryptoFailedException {
-        byte[] plaintext;
+    byte[] processReceiving(List<AxolotlKey> possibleKeys) throws CryptoFailedException {
+        byte[] plaintext = null;
         FingerprintStatus status = getTrust();
         if (!status.isCompromised()) {
-            try {
-                if (encryptedKey.prekey) {
-                    PreKeySignalMessage preKeySignalMessage = new PreKeySignalMessage(encryptedKey.key);
-                    Optional<Integer> optionalPreKeyId = preKeySignalMessage.getPreKeyId();
-                    IdentityKey identityKey = preKeySignalMessage.getIdentityKey();
-                    if (!optionalPreKeyId.isPresent()) {
-                        throw new CryptoFailedException("PreKeyWhisperMessage did not contain a PreKeyId");
+            Iterator<AxolotlKey> iterator = possibleKeys.iterator();
+            while (iterator.hasNext()) {
+                AxolotlKey encryptedKey = iterator.next();
+                try {
+                    if (encryptedKey.prekey) {
+                        PreKeySignalMessage preKeySignalMessage = new PreKeySignalMessage(encryptedKey.key);
+                        Optional<Integer> optionalPreKeyId = preKeySignalMessage.getPreKeyId();
+                        IdentityKey identityKey = preKeySignalMessage.getIdentityKey();
+                        if (!optionalPreKeyId.isPresent()) {
+                            if (iterator.hasNext()) {
+                                continue;
+                            }
+                            throw new CryptoFailedException("PreKeyWhisperMessage did not contain a PreKeyId");
+                        }
+                        preKeyId = optionalPreKeyId.get();
+                        if (this.identityKey != null && !this.identityKey.equals(identityKey)) {
+                            if (iterator.hasNext()) {
+                                continue;
+                            }
+                            throw new CryptoFailedException("Received PreKeyWhisperMessage but preexisting identity key changed.");
+                        }
+                        this.identityKey = identityKey;
+                        plaintext = cipher.decrypt(preKeySignalMessage);
+                    } else {
+                        SignalMessage signalMessage = new SignalMessage(encryptedKey.key);
+                        try {
+                            plaintext = cipher.decrypt(signalMessage);
+                        } catch (InvalidMessageException | NoSessionException e) {
+                            if (iterator.hasNext()) {
+                                Log.w(Config.LOGTAG, account.getJid().asBareJid() + ": ignoring crypto exception because possible keys left to try", e);
+                                continue;
+                            }
+                            throw new BrokenSessionException(this.remoteAddress, e);
+                        }
+                        preKeyId = null; //better safe than sorry because we use that to do special after prekey handling
                     }
-                    preKeyId = optionalPreKeyId.get();
-                    if (this.identityKey != null && !this.identityKey.equals(identityKey)) {
-                        throw new CryptoFailedException("Received PreKeyWhisperMessage but preexisting identity key changed.");
+                } catch (InvalidVersionException | InvalidKeyException | LegacyMessageException | InvalidMessageException | DuplicateMessageException | InvalidKeyIdException | UntrustedIdentityException e) {
+                    if (iterator.hasNext()) {
+                        Log.w(Config.LOGTAG, account.getJid().asBareJid() + ": ignoring crypto exception because possible keys left to try", e);
+                        continue;
                     }
-                    this.identityKey = identityKey;
-                    plaintext = cipher.decrypt(preKeySignalMessage);
-                } else {
-                    SignalMessage signalMessage = new SignalMessage(encryptedKey.key);
-                    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
+                    throw new CryptoFailedException("Error decrypting SignalMessage", e);
+                }
+                if (iterator.hasNext()) {
+                    break;
                 }
-            } catch (InvalidVersionException | InvalidKeyException | LegacyMessageException | InvalidMessageException | DuplicateMessageException | InvalidKeyIdException | UntrustedIdentityException e) {
-                throw new CryptoFailedException("Error decrypting SignalMessage", e);
             }
             if (!status.isActive()) {
                 setTrust(status.toActive());
@@ -125,7 +151,7 @@ public class XmppAxolotlSession implements Comparable<XmppAxolotlSession> {
         if (ignoreSessionTrust || status.isTrustedAndActive()) {
             try {
                 CiphertextMessage ciphertextMessage = cipher.encrypt(outgoingMessage);
-                return new AxolotlKey(ciphertextMessage.serialize(), ciphertextMessage.getType() == CiphertextMessage.PREKEY_TYPE);
+                return new AxolotlKey(getRemoteAddress().getDeviceId(), ciphertextMessage.serialize(), ciphertextMessage.getType() == CiphertextMessage.PREKEY_TYPE);
             } catch (UntrustedIdentityException e) {
                 return null;
             }
@@ -148,8 +174,10 @@ public class XmppAxolotlSession implements Comparable<XmppAxolotlSession> {
 
         public final byte[] key;
         public final boolean prekey;
+        public final int deviceId;
 
-        public AxolotlKey(byte[] key, boolean prekey) {
+        public AxolotlKey(int deviceId, byte[] key, boolean prekey) {
+            this.deviceId = deviceId;
             this.key = key;
             this.prekey = prekey;
         }