2 files changed, 25 insertions, 11 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 043cc91ff..94064a661 100644
--- a/src/main/java/de/pixart/messenger/crypto/axolotl/XmppAxolotlMessage.java
+++ b/src/main/java/de/pixart/messenger/crypto/axolotl/XmppAxolotlMessage.java
@@ -41,7 +41,7 @@ public class XmppAxolotlMessage {
     private byte[] innerKey;
     private byte[] ciphertext = null;
     private byte[] iv = null;
-    private final Map<Integer, byte[]> keys;
+    private final Map<Integer, XmppAxolotlSession.AxolotlKey> keys;
     private final Jid from;
     private final int sourceDeviceId;
 
@@ -104,7 +104,8 @@ public class XmppAxolotlMessage {
                     try {
                         Integer recipientId = Integer.parseInt(keyElement.getAttribute(REMOTEID));
                         byte[] key = Base64.decode(keyElement.getContent().trim(), Base64.DEFAULT);
-                        this.keys.put(recipientId, key);
+                        boolean isPreKey = keyElement.getAttributeAsBoolean("prekey");
+                        this.keys.put(recipientId, new XmppAxolotlSession.AxolotlKey(key, isPreKey));
                     } catch (NumberFormatException e) {
                         throw new IllegalArgumentException("invalid remote id");
                     }
@@ -199,7 +200,7 @@ public class XmppAxolotlMessage {
     }
 
     public void addDevice(XmppAxolotlSession session) {
-        byte[] key = session.processSending(innerKey);
+        XmppAxolotlSession.AxolotlKey key = session.processSending(innerKey);
         if (key != null) {
             keys.put(session.getRemoteAddress().getDeviceId(), key);
         }
@@ -217,10 +218,13 @@ public class XmppAxolotlMessage {
         Element encryptionElement = new Element(CONTAINERTAG, AxolotlService.PEP_PREFIX);
         Element headerElement = encryptionElement.addChild(HEADER);
         headerElement.setAttribute(SOURCEID, sourceDeviceId);
-        for (Map.Entry<Integer, byte[]> keyEntry : keys.entrySet()) {
+        for (Map.Entry<Integer, XmppAxolotlSession.AxolotlKey> keyEntry : keys.entrySet()) {
             Element keyElement = new Element(KEYTAG);
             keyElement.setAttribute(REMOTEID, keyEntry.getKey());
-            keyElement.setContent(Base64.encodeToString(keyEntry.getValue(), Base64.NO_WRAP));
+            if (keyEntry.getValue().prekey) {
+                keyElement.setAttribute("prekey", "true");
+            }
+            keyElement.setContent(Base64.encodeToString(keyEntry.getValue().key, Base64.NO_WRAP));
             headerElement.addChild(keyElement);
         }
         headerElement.addChild(IVTAG).setContent(Base64.encodeToString(iv, Base64.NO_WRAP));
@@ -232,7 +236,7 @@ public class XmppAxolotlMessage {
     }
 
     private byte[] unpackKey(XmppAxolotlSession session, Integer sourceDeviceId) {
-        byte[] encryptedKey = keys.get(sourceDeviceId);
+        XmppAxolotlSession.AxolotlKey encryptedKey = keys.get(sourceDeviceId);
         return (encryptedKey != null) ? session.processReceiving(encryptedKey) : null;
     }
 
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 d850e18e7..59a7fc71f 100644
--- a/src/main/java/de/pixart/messenger/crypto/axolotl/XmppAxolotlSession.java
+++ b/src/main/java/de/pixart/messenger/crypto/axolotl/XmppAxolotlSession.java
@@ -82,13 +82,13 @@ public class XmppAxolotlSession implements Comparable<XmppAxolotlSession> {
     }
 
     @Nullable
-    public byte[] processReceiving(byte[] encryptedKey) {
+    public byte[] processReceiving(AxolotlKey encryptedKey) {
         byte[] plaintext = null;
         FingerprintStatus status = getTrust();
         if (!status.isCompromised()) {
             try {
                 try {
-                    PreKeyWhisperMessage message = new PreKeyWhisperMessage(encryptedKey);
+                    PreKeyWhisperMessage message = new PreKeyWhisperMessage(encryptedKey.key);
                     if (!message.getPreKeyId().isPresent()) {
                         Log.w(Config.LOGTAG, AxolotlService.getLogprefix(account) + "PreKeyWhisperMessage did not contain a PreKeyId");
                         return null;
@@ -104,7 +104,7 @@ public class XmppAxolotlSession implements Comparable<XmppAxolotlSession> {
                     }
                 } catch (InvalidMessageException | InvalidVersionException e) {
                     Log.i(Config.LOGTAG, AxolotlService.getLogprefix(account) + "WhisperMessage received");
-                    WhisperMessage message = new WhisperMessage(encryptedKey);
+                    WhisperMessage message = new WhisperMessage(encryptedKey.key);
                     plaintext = cipher.decrypt(message);
                 } catch (InvalidKeyException | InvalidKeyIdException | UntrustedIdentityException e) {
                     Log.w(Config.LOGTAG, AxolotlService.getLogprefix(account) + "Error decrypting axolotl header, " + e.getClass().getName() + ": " + e.getMessage());
@@ -123,11 +123,11 @@ public class XmppAxolotlSession implements Comparable<XmppAxolotlSession> {
     }
 
     @Nullable
-    public byte[] processSending(@NonNull byte[] outgoingMessage) {
+    public AxolotlKey processSending(@NonNull byte[] outgoingMessage) {
         FingerprintStatus status = getTrust();
         if (status.isTrustedAndActive()) {
             CiphertextMessage ciphertextMessage = cipher.encrypt(outgoingMessage);
-            return ciphertextMessage.serialize();
+            return new AxolotlKey(ciphertextMessage.serialize(),ciphertextMessage.getType() == CiphertextMessage.PREKEY_TYPE);
         } else {
             return null;
         }
@@ -141,4 +141,14 @@ public class XmppAxolotlSession implements Comparable<XmppAxolotlSession> {
     public int compareTo(XmppAxolotlSession o) {
         return getTrust().compareTo(o.getTrust());
     }
+
+    public static class AxolotlKey {
+        public final byte[] key;
+        public final boolean prekey;
+
+        public AxolotlKey(byte[] key, boolean prekey) {
+            this.key = key;
+            this.prekey = prekey;
+        }
+    }
 }