Add callback method parameter to axolotl decrypt functions.

// FREEBIE

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

@@ -3,7 +3,6 @@ package org.whispersystems.test;
 import android.test.AndroidTestCase;
 
 import org.whispersystems.libaxolotl.DuplicateMessageException;
-import org.whispersystems.libaxolotl.IdentityKey;
 import org.whispersystems.libaxolotl.InvalidKeyException;
 import org.whispersystems.libaxolotl.InvalidKeyIdException;
 import org.whispersystems.libaxolotl.InvalidMessageException;
@@ -24,10 +23,7 @@ import org.whispersystems.libaxolotl.state.AxolotlStore;
 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;
@@ -122,7 +118,7 @@ public class SessionBuilderTest extends AndroidTestCase {
     AxolotlStore   aliceStore          = new InMemoryAxolotlStore();
     SessionBuilder aliceSessionBuilder = new SessionBuilder(aliceStore, BOB_RECIPIENT_ID, 1);
 
-    AxolotlStore bobStore                 = new InMemoryAxolotlStore();
+    final AxolotlStore bobStore                 = new InMemoryAxolotlStore();
           ECKeyPair    bobPreKeyPair            = Curve.generateKeyPair();
           ECKeyPair    bobSignedPreKeyPair      = Curve.generateKeyPair();
           byte[]       bobSignedPreKeySignature = Curve.calculateSignature(bobStore.getIdentityKeyPair().getPrivateKey(),
@@ -139,7 +135,7 @@ public class SessionBuilderTest extends AndroidTestCase {
     assertTrue(aliceStore.containsSession(BOB_RECIPIENT_ID, 1));
     assertTrue(aliceStore.loadSession(BOB_RECIPIENT_ID, 1).getSessionState().getSessionVersion() == 3);
 
-    String            originalMessage    = "L'homme est condamné à être libre";
+    final String            originalMessage    = "L'homme est condamné à être libre";
           SessionCipher     aliceSessionCipher = new SessionCipher(aliceStore, BOB_RECIPIENT_ID, 1);
           CiphertextMessage outgoingMessage    = aliceSessionCipher.encrypt(originalMessage.getBytes());
 
@@ -150,8 +146,13 @@ public class SessionBuilderTest extends AndroidTestCase {
     bobStore.storeSignedPreKey(22, new SignedPreKeyRecord(22, System.currentTimeMillis(), bobSignedPreKeyPair, bobSignedPreKeySignature));
 
     SessionCipher bobSessionCipher = new SessionCipher(bobStore, ALICE_RECIPIENT_ID, 1);
-    byte[]        plaintext        = bobSessionCipher.decrypt(incomingMessage);
+    byte[] plaintext = bobSessionCipher.decrypt(incomingMessage, new SessionCipher.DecryptionCallback() {
-
+      @Override
+      public void handlePlaintext(byte[] plaintext) {
+        assertTrue(originalMessage.equals(new String(plaintext)));
+        assertFalse(bobStore.containsSession(ALICE_RECIPIENT_ID, 1));
+      }
+    });
 
     assertTrue(bobStore.containsSession(ALICE_RECIPIENT_ID, 1));
     assertTrue(bobStore.loadSession(ALICE_RECIPIENT_ID, 1).getSessionState().getSessionVersion() == 3);

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

@@ -19,12 +19,8 @@ import org.whispersystems.libaxolotl.ratchet.AliceAxolotlParameters;
 import org.whispersystems.libaxolotl.ratchet.BobAxolotlParameters;
 import org.whispersystems.libaxolotl.ratchet.RatchetingSession;
 import org.whispersystems.libaxolotl.state.AxolotlStore;
-import org.whispersystems.libaxolotl.state.IdentityKeyStore;
-import org.whispersystems.libaxolotl.state.PreKeyStore;
 import org.whispersystems.libaxolotl.state.SessionRecord;
 import org.whispersystems.libaxolotl.state.SessionState;
-import org.whispersystems.libaxolotl.state.SessionStore;
-import org.whispersystems.libaxolotl.state.SignedPreKeyStore;
 import org.whispersystems.libaxolotl.util.guava.Optional;
 
 import java.security.NoSuchAlgorithmException;

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

@@ -138,6 +138,7 @@ public class SessionCipher {
    * Decrypt a message.
    *
    * @param  ciphertext The {@link PreKeyWhisperMessage} to decrypt.
+   *
    * @return The plaintext.
    * @throws InvalidMessageException if the input is not valid ciphertext.
    * @throws DuplicateMessageException if the input is a message that has already been received.
@@ -147,17 +148,46 @@ public class SessionCipher {
    *                               that corresponds to the PreKey ID in the message.
    * @throws InvalidKeyException when the message is formatted incorrectly.
    * @throws UntrustedIdentityException when the {@link IdentityKey} of the sender is untrusted.
-
    */
   public byte[] decrypt(PreKeyWhisperMessage ciphertext)
       throws DuplicateMessageException, LegacyMessageException, InvalidMessageException,
              InvalidKeyIdException, InvalidKeyException, UntrustedIdentityException
+  {
+    return decrypt(ciphertext, new NullDecryptionCallback());
+  }
+
+  /**
+   * Decrypt a message.
+   *
+   * @param  ciphertext The {@link PreKeyWhisperMessage} to decrypt.
+   * @param  callback   A callback that is triggered after decryption is complete,
+   *                    but before the updated session state has been committed to the session
+   *                    DB.  This allows some implementations to store the committed plaintext
+   *                    to a DB first, in case they are concerned with a crash happening between
+   *                    the time the session state is updated but before they're able to store
+   *                    the plaintext to disk.
+   *
+   * @return The plaintext.
+   * @throws InvalidMessageException if the input is not valid ciphertext.
+   * @throws DuplicateMessageException if the input is a message that has already been received.
+   * @throws LegacyMessageException if the input is a message formatted by a protocol version that
+   *                                is no longer supported.
+   * @throws InvalidKeyIdException when there is no local {@link org.whispersystems.libaxolotl.state.PreKeyRecord}
+   *                               that corresponds to the PreKey ID in the message.
+   * @throws InvalidKeyException when the message is formatted incorrectly.
+   * @throws UntrustedIdentityException when the {@link IdentityKey} of the sender is untrusted.
+   */
+  public byte[] decrypt(PreKeyWhisperMessage ciphertext, DecryptionCallback callback)
+      throws DuplicateMessageException, LegacyMessageException, InvalidMessageException,
+             InvalidKeyIdException, InvalidKeyException, UntrustedIdentityException
   {
     synchronized (SESSION_LOCK) {
       SessionRecord     sessionRecord    = sessionStore.loadSession(recipientId, deviceId);
       Optional<Integer> unsignedPreKeyId = sessionBuilder.process(sessionRecord, ciphertext);
       byte[]            plaintext        = decrypt(sessionRecord, ciphertext.getWhisperMessage());
 
+      callback.handlePlaintext(plaintext);
+
       sessionStore.storeSession(recipientId, deviceId, sessionRecord);
 
       if (unsignedPreKeyId.isPresent()) {
@@ -168,7 +198,6 @@ public class SessionCipher {
     }
   }
 
-
   /**
    * Decrypt a message.
    *
@@ -184,6 +213,31 @@ public class SessionCipher {
   public byte[] decrypt(WhisperMessage ciphertext)
       throws InvalidMessageException, DuplicateMessageException, LegacyMessageException,
       NoSessionException
+  {
+    return decrypt(ciphertext, new NullDecryptionCallback());
+  }
+
+  /**
+   * Decrypt a message.
+   *
+   * @param  ciphertext The {@link WhisperMessage} to decrypt.
+   * @param  callback   A callback that is triggered after decryption is complete,
+   *                    but before the updated session state has been committed to the session
+   *                    DB.  This allows some implementations to store the committed plaintext
+   *                    to a DB first, in case they are concerned with a crash happening between
+   *                    the time the session state is updated but before they're able to store
+   *                    the plaintext to disk.
+   *
+   * @return The plaintext.
+   * @throws InvalidMessageException if the input is not valid ciphertext.
+   * @throws DuplicateMessageException if the input is a message that has already been received.
+   * @throws LegacyMessageException if the input is a message formatted by a protocol version that
+   *                                is no longer supported.
+   * @throws NoSessionException if there is no established session for this contact.
+   */
+  public byte[] decrypt(WhisperMessage ciphertext, DecryptionCallback callback)
+      throws InvalidMessageException, DuplicateMessageException, LegacyMessageException,
+             NoSessionException
   {
     synchronized (SESSION_LOCK) {
 
@@ -194,6 +248,8 @@ public class SessionCipher {
       SessionRecord sessionRecord = sessionStore.loadSession(recipientId, deviceId);
       byte[]        plaintext     = decrypt(sessionRecord, ciphertext);
 
+      callback.handlePlaintext(plaintext);
+
       sessionStore.storeSession(recipientId, deviceId, sessionRecord);
 
       return plaintext;
@@ -401,4 +457,13 @@ public class SessionCipher {
       throw new AssertionError(e);
     }
   }
+
+  public static interface DecryptionCallback {
+    public void handlePlaintext(byte[] plaintext);
+  }
+
+  private static class NullDecryptionCallback implements DecryptionCallback {
+    @Override
+    public void handlePlaintext(byte[] plaintext) {}
+  }
 }