From 3d9c94428898f6e633b67082237b620c7877f5c1 Mon Sep 17 00:00:00 2001
From: Moxie Marlinspike <moxie@thoughtcrime.org>
Date: Mon, 6 Apr 2015 12:21:54 -0700
Subject: [PATCH] Add decryption callback to group decrypt operations.

---
 .../libaxolotl/DecryptionCallback.java        |  5 ++++
 .../libaxolotl/SessionCipher.java             |  4 ---
 .../libaxolotl/groups/GroupCipher.java        | 29 +++++++++++++++++++
 .../libaxolotl/SessionBuilderTest.java        |  2 +-
 4 files changed, 35 insertions(+), 5 deletions(-)
 create mode 100644 java/src/main/java/org/whispersystems/libaxolotl/DecryptionCallback.java

diff --git a/java/src/main/java/org/whispersystems/libaxolotl/DecryptionCallback.java b/java/src/main/java/org/whispersystems/libaxolotl/DecryptionCallback.java
new file mode 100644
index 000000000..1c70d6be2
--- /dev/null
+++ b/java/src/main/java/org/whispersystems/libaxolotl/DecryptionCallback.java
@@ -0,0 +1,5 @@
+package org.whispersystems.libaxolotl;
+
+public interface DecryptionCallback {
+  public void handlePlaintext(byte[] plaintext);
+}
\ No newline at end of file
diff --git a/java/src/main/java/org/whispersystems/libaxolotl/SessionCipher.java b/java/src/main/java/org/whispersystems/libaxolotl/SessionCipher.java
index b9cc6ac1a..071e0905e 100644
--- a/java/src/main/java/org/whispersystems/libaxolotl/SessionCipher.java
+++ b/java/src/main/java/org/whispersystems/libaxolotl/SessionCipher.java
@@ -455,10 +455,6 @@ public class SessionCipher {
     }
   }
 
-  public static interface DecryptionCallback {
-    public void handlePlaintext(byte[] plaintext);
-  }
-
   private static class NullDecryptionCallback implements DecryptionCallback {
     @Override
     public void handlePlaintext(byte[] plaintext) {}
diff --git a/java/src/main/java/org/whispersystems/libaxolotl/groups/GroupCipher.java b/java/src/main/java/org/whispersystems/libaxolotl/groups/GroupCipher.java
index 778ca916a..55e261ac6 100644
--- a/java/src/main/java/org/whispersystems/libaxolotl/groups/GroupCipher.java
+++ b/java/src/main/java/org/whispersystems/libaxolotl/groups/GroupCipher.java
@@ -16,6 +16,7 @@
  */
 package org.whispersystems.libaxolotl.groups;
 
+import org.whispersystems.libaxolotl.DecryptionCallback;
 import org.whispersystems.libaxolotl.DuplicateMessageException;
 import org.whispersystems.libaxolotl.InvalidKeyIdException;
 import org.whispersystems.libaxolotl.InvalidMessageException;
@@ -101,6 +102,27 @@ public class GroupCipher {
    * @throws DuplicateMessageException
    */
   public byte[] decrypt(byte[] senderKeyMessageBytes)
+      throws LegacyMessageException, DuplicateMessageException, InvalidMessageException
+  {
+    return decrypt(senderKeyMessageBytes, new NullDecryptionCallback());
+  }
+
+  /**
+   * Decrypt a SenderKey group message.
+   *
+   * @param senderKeyMessageBytes The received ciphertext.
+   * @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 Plaintext
+   * @throws LegacyMessageException
+   * @throws InvalidMessageException
+   * @throws DuplicateMessageException
+   */
+  public byte[] decrypt(byte[] senderKeyMessageBytes, DecryptionCallback callback)
       throws LegacyMessageException, InvalidMessageException, DuplicateMessageException
   {
     synchronized (LOCK) {
@@ -115,6 +137,8 @@ public class GroupCipher {
 
         byte[] plaintext = getPlainText(senderKey.getIv(), senderKey.getCipherKey(), senderKeyMessage.getCipherText());
 
+        callback.handlePlaintext(plaintext);
+
         senderKeyStore.storeSenderKey(senderKeyId, record);
 
         return plaintext;
@@ -185,4 +209,9 @@ public class GroupCipher {
     }
   }
 
+  private static class NullDecryptionCallback implements DecryptionCallback {
+    @Override
+    public void handlePlaintext(byte[] plaintext) {}
+  }
+
 }
diff --git a/tests/src/test/java/org/whispersystems/libaxolotl/SessionBuilderTest.java b/tests/src/test/java/org/whispersystems/libaxolotl/SessionBuilderTest.java
index 1105ffc62..9bf8fcb31 100644
--- a/tests/src/test/java/org/whispersystems/libaxolotl/SessionBuilderTest.java
+++ b/tests/src/test/java/org/whispersystems/libaxolotl/SessionBuilderTest.java
@@ -135,7 +135,7 @@ public class SessionBuilderTest extends TestCase {
     bobStore.storeSignedPreKey(22, new SignedPreKeyRecord(22, System.currentTimeMillis(), bobSignedPreKeyPair, bobSignedPreKeySignature));
 
     SessionCipher bobSessionCipher = new SessionCipher(bobStore, ALICE_ADDRESS);
-    byte[] plaintext = bobSessionCipher.decrypt(incomingMessage, new SessionCipher.DecryptionCallback() {
+    byte[] plaintext = bobSessionCipher.decrypt(incomingMessage, new DecryptionCallback() {
       @Override
       public void handlePlaintext(byte[] plaintext) {
         assertTrue(originalMessage.equals(new String(plaintext)));