Only remove unsigned prekey if bundled message decrypts properly.

libaxolotl/src/androidTest/java/org/whispersystems/test/InMemorySignedPreKeyStore.java

@@ -18,7 +18,7 @@ public class InMemorySignedPreKeyStore implements SignedPreKeyStore {
   public SignedPreKeyRecord loadSignedPreKey(int signedPreKeyId) throws InvalidKeyIdException {
     try {
       if (!store.containsKey(signedPreKeyId)) {
-        throw new InvalidKeyIdException("No such signedprekeyrecord!");
+        throw new InvalidKeyIdException("No such signedprekeyrecord! " + signedPreKeyId);
       }
 
       return new SignedPreKeyRecord(store.get(signedPreKeyId));

libaxolotl/src/androidTest/java/org/whispersystems/test/SessionBuilderTest.java

@@ -19,13 +19,13 @@ import org.whispersystems.libaxolotl.protocol.CiphertextMessage;
 import org.whispersystems.libaxolotl.protocol.KeyExchangeMessage;
 import org.whispersystems.libaxolotl.protocol.PreKeyWhisperMessage;
 import org.whispersystems.libaxolotl.protocol.WhisperMessage;
-import org.whispersystems.libaxolotl.state.SignedPreKeyRecord;
-import org.whispersystems.libaxolotl.state.SignedPreKeyStore;
 import org.whispersystems.libaxolotl.state.IdentityKeyStore;
 import org.whispersystems.libaxolotl.state.PreKeyBundle;
 import org.whispersystems.libaxolotl.state.PreKeyRecord;
 import org.whispersystems.libaxolotl.state.PreKeyStore;
 import org.whispersystems.libaxolotl.state.SessionStore;
+import org.whispersystems.libaxolotl.state.SignedPreKeyRecord;
+import org.whispersystems.libaxolotl.state.SignedPreKeyStore;
 import org.whispersystems.libaxolotl.util.Pair;
 
 import java.util.HashSet;
@@ -406,6 +406,68 @@ public class SessionBuilderTest extends AndroidTestCase {
 
   }
 
+  public void testBadMessageBundle() throws InvalidKeyException, UntrustedIdentityException, InvalidVersionException, InvalidMessageException, DuplicateMessageException, LegacyMessageException, InvalidKeyIdException {
+    SessionStore      aliceSessionStore      = new InMemorySessionStore();
+    SignedPreKeyStore aliceSignedPreKeyStore = new InMemorySignedPreKeyStore();
+    PreKeyStore       alicePreKeyStore       = new InMemoryPreKeyStore();
+    IdentityKeyStore  aliceIdentityKeyStore  = new InMemoryIdentityKeyStore();
+    SessionBuilder    aliceSessionBuilder    = new SessionBuilder(aliceSessionStore, alicePreKeyStore,
+                                                                  aliceSignedPreKeyStore,
+                                                                  aliceIdentityKeyStore,
+                                                                  BOB_RECIPIENT_ID, 1);
+
+    SessionStore      bobSessionStore      = new InMemorySessionStore();
+    PreKeyStore       bobPreKeyStore       = new InMemoryPreKeyStore();
+    SignedPreKeyStore bobSignedPreKeyStore = new InMemorySignedPreKeyStore();
+    IdentityKeyStore  bobIdentityKeyStore  = new InMemoryIdentityKeyStore();
+
+    ECKeyPair bobPreKeyPair            = Curve.generateKeyPair();
+    ECKeyPair bobSignedPreKeyPair      = Curve.generateKeyPair();
+    byte[]    bobSignedPreKeySignature = Curve.calculateSignature(bobIdentityKeyStore.getIdentityKeyPair().getPrivateKey(),
+                                                                  bobSignedPreKeyPair.getPublicKey().serialize());
+
+    PreKeyBundle bobPreKey = new PreKeyBundle(bobIdentityKeyStore.getLocalRegistrationId(), 1,
+                                              31337, bobPreKeyPair.getPublicKey(),
+                                              22, bobSignedPreKeyPair.getPublicKey(), bobSignedPreKeySignature,
+                                              bobIdentityKeyStore.getIdentityKeyPair().getPublicKey());
+
+    bobPreKeyStore.storePreKey(31337, new PreKeyRecord(bobPreKey.getPreKeyId(), bobPreKeyPair));
+    bobSignedPreKeyStore.storeSignedPreKey(22, new SignedPreKeyRecord(22, System.currentTimeMillis(), bobSignedPreKeyPair, bobSignedPreKeySignature));
+
+    aliceSessionBuilder.process(bobPreKey);
+
+    String            originalMessage    = "L'homme est condamné à être libre";
+    SessionCipher     aliceSessionCipher = new SessionCipher(aliceSessionStore, alicePreKeyStore, aliceSignedPreKeyStore, aliceIdentityKeyStore, BOB_RECIPIENT_ID, 1);
+    CiphertextMessage outgoingMessageOne = aliceSessionCipher.encrypt(originalMessage.getBytes());
+
+    assertTrue(outgoingMessageOne.getType() == CiphertextMessage.PREKEY_TYPE);
+
+    byte[] goodMessage = outgoingMessageOne.serialize();
+    byte[] badMessage  = new byte[goodMessage.length];
+    System.arraycopy(goodMessage, 0, badMessage, 0, badMessage.length);
+
+    badMessage[badMessage.length-10] ^= 0x01;
+
+    PreKeyWhisperMessage incomingMessage  = new PreKeyWhisperMessage(badMessage);
+    SessionCipher        bobSessionCipher = new SessionCipher(bobSessionStore, bobPreKeyStore, bobSignedPreKeyStore, bobIdentityKeyStore, ALICE_RECIPIENT_ID, 1);
+
+    byte[] plaintext = new byte[0];
+
+    try {
+      plaintext = bobSessionCipher.decrypt(incomingMessage);
+      throw new AssertionError("Decrypt should have failed!");
+    } catch (InvalidMessageException e) {
+      // good.
+    }
+
+    assertTrue(bobPreKeyStore.containsPreKey(31337));
+
+    plaintext = bobSessionCipher.decrypt(new PreKeyWhisperMessage(goodMessage));
+
+    assertTrue(originalMessage.equals(new String(plaintext)));
+    assertTrue(!bobPreKeyStore.containsPreKey(31337));
+  }
+
   public void testBasicKeyExchange() throws InvalidKeyException, LegacyMessageException, InvalidMessageException, DuplicateMessageException, UntrustedIdentityException, StaleKeyExchangeException, InvalidVersionException, NoSessionException {
     SessionStore     aliceSessionStore     = new InMemorySessionStore();
     PreKeyStore      alicePreKeyStore      = new InMemoryPreKeyStore();

libaxolotl/src/main/java/org/whispersystems/libaxolotl/SessionBuilder.java

@@ -88,32 +88,34 @@ public class SessionBuilder {
    * @throws org.whispersystems.libaxolotl.InvalidKeyException when the message is formatted incorrectly.
    * @throws org.whispersystems.libaxolotl.UntrustedIdentityException when the {@link IdentityKey} of the sender is untrusted.
    */
-  /*package*/ void process(SessionRecord sessionRecord, PreKeyWhisperMessage message)
+  /*package*/ int process(SessionRecord sessionRecord, PreKeyWhisperMessage message)
       throws InvalidKeyIdException, InvalidKeyException, UntrustedIdentityException
   {
     int         messageVersion   = message.getMessageVersion();
     IdentityKey theirIdentityKey = message.getIdentityKey();
+    int         unsignedPreKeyId;
 
     if (!identityKeyStore.isTrustedIdentity(recipientId, theirIdentityKey)) {
       throw new UntrustedIdentityException();
     }
 
     switch (messageVersion) {
-      case 2:  processV2(sessionRecord, message); break;
+      case 2:  unsignedPreKeyId = processV2(sessionRecord, message); break;
-      case 3:  processV3(sessionRecord, message); break;
+      case 3:  unsignedPreKeyId = processV3(sessionRecord, message); break;
       default: throw new AssertionError("Unknown version: " + messageVersion);
     }
 
     identityKeyStore.saveIdentity(recipientId, theirIdentityKey);
+    return unsignedPreKeyId;
   }
 
-  private void processV3(SessionRecord sessionRecord, PreKeyWhisperMessage message)
+  private int processV3(SessionRecord sessionRecord, PreKeyWhisperMessage message)
       throws UntrustedIdentityException, InvalidKeyIdException, InvalidKeyException
   {
 
     if (sessionRecord.hasSessionState(message.getMessageVersion(), message.getBaseKey().serialize())) {
       Log.w(TAG, "We've already setup a session for this V3 message, letting bundled message fall through...");
-      return;
+      return -1;
     }
 
     boolean   simultaneousInitiate = sessionRecord.getSessionState().hasUnacknowledgedPreKeyMessage();
@@ -145,11 +147,13 @@ public class SessionBuilder {
     if (simultaneousInitiate) sessionRecord.getSessionState().setNeedsRefresh(true);
 
     if (message.getPreKeyId() >= 0 && message.getPreKeyId() != Medium.MAX_VALUE) {
-      preKeyStore.removePreKey(message.getPreKeyId());
+      return message.getPreKeyId();
+    } else {
+      return -1;
     }
   }
 
-  private void processV2(SessionRecord sessionRecord, PreKeyWhisperMessage message)
+  private int processV2(SessionRecord sessionRecord, PreKeyWhisperMessage message)
       throws UntrustedIdentityException, InvalidKeyIdException, InvalidKeyException
   {
 
@@ -157,7 +161,7 @@ public class SessionBuilder {
         sessionStore.containsSession(recipientId, deviceId))
     {
       Log.w(TAG, "We've already processed the prekey part of this V2 session, letting bundled message fall through...");
-      return;
+      return -1;
     }
 
     ECKeyPair     ourPreKey            = preKeyStore.loadPreKey(message.getPreKeyId()).getKeyPair();
@@ -183,7 +187,9 @@ public class SessionBuilder {
     if (simultaneousInitiate) sessionRecord.getSessionState().setNeedsRefresh(true);
 
     if (message.getPreKeyId() != Medium.MAX_VALUE) {
-      preKeyStore.removePreKey(message.getPreKeyId());
+      return message.getPreKeyId();
+    } else {
+      return -1;
     }
   }
 

libaxolotl/src/main/java/org/whispersystems/libaxolotl/SessionCipher.java

@@ -63,6 +63,7 @@ public class SessionCipher {
 
   private final SessionStore   sessionStore;
   private final SessionBuilder sessionBuilder;
+  private final PreKeyStore    preKeyStore;
   private final long           recipientId;
   private final int            deviceId;
 
@@ -82,6 +83,7 @@ public class SessionCipher {
     this.sessionStore   = sessionStore;
     this.recipientId    = recipientId;
     this.deviceId       = deviceId;
+    this.preKeyStore    = preKeyStore;
     this.sessionBuilder = new SessionBuilder(sessionStore, preKeyStore, signedPreKeyStore,
                                              identityKeyStore, recipientId, deviceId);
   }
@@ -145,12 +147,16 @@ public class SessionCipher {
              InvalidKeyIdException, InvalidKeyException, UntrustedIdentityException
   {
     synchronized (SESSION_LOCK) {
-      SessionRecord sessionRecord = sessionStore.loadSession(recipientId, deviceId);
+      SessionRecord sessionRecord    = sessionStore.loadSession(recipientId, deviceId);
-
+      int           unsignedPreKeyId = sessionBuilder.process(sessionRecord, ciphertext);
-      sessionBuilder.process(sessionRecord, ciphertext);
+      byte[]        plaintext        = decrypt(sessionRecord, ciphertext.getWhisperMessage());
-      byte[] plaintext = decrypt(sessionRecord, ciphertext.getWhisperMessage());
 
       sessionStore.storeSession(recipientId, deviceId, sessionRecord);
+
+      if (unsignedPreKeyId >=0) {
+        preKeyStore.removePreKey(unsignedPreKeyId);
+      }
+
       return plaintext;
     }
   }