From 30dd13e8e399f70bccf9bf86e27b3e1979e12d0c Mon Sep 17 00:00:00 2001
From: Matt Corallo <git@bluematt.me>
Date: Sun, 18 Jan 2015 11:10:39 -1000
Subject: [PATCH] Return a function instead of a session in
 decryptWhisperMessage

---
 libaxolotl/protocol.js           |  8 ++++++--
 libtextsecure/axolotl_wrapper.js | 22 ++++++++++++++--------
 2 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/libaxolotl/protocol.js b/libaxolotl/protocol.js
index 625250e5..7e26cf0b 100644
--- a/libaxolotl/protocol.js
+++ b/libaxolotl/protocol.js
@@ -457,7 +457,7 @@ window.axolotl.protocol = function() {
     /*************************
     *** Public crypto API ***
     *************************/
-    // returns decrypted plaintext
+    // returns decrypted plaintext and a function that MUST BE CALLED SYNCHRONOUSLY if the message indicates session close
     self.decryptWhisperMessage = function(encodedNumber, messageBytes, session, registrationId) {
         if (messageBytes[0] != String.fromCharCode((3 << 4) | 3))
             throw new Error("Bad version number on WhisperMessage");
@@ -507,7 +507,11 @@ window.axolotl.protocol = function() {
                             delete session['pendingPreKey'];
                             removeOldChains(session);
                             crypto_storage.saveSession(encodedNumber, session, registrationId);
-                            return [plaintext, session];
+                            return [plaintext, function() {
+                                closeSession(session, true);
+                                removeOldChains(session);
+                                crypto_storage.saveSession(encodedNumber, session, registrationId);
+                            }];
                         });
                     });
                 });
diff --git a/libtextsecure/axolotl_wrapper.js b/libtextsecure/axolotl_wrapper.js
index ad3c6de9..7b9834cf 100644
--- a/libtextsecure/axolotl_wrapper.js
+++ b/libtextsecure/axolotl_wrapper.js
@@ -39,10 +39,20 @@
     var decodeMessageContents = function(res) {
         var finalMessage = textsecure.protobuf.PushMessageContent.decode(res[0]);
 
-        //TODO
+        if ((finalMessage.flags & textsecure.protobuf.PushMessageContent.Flags.END_SESSION)
+                == textsecure.protobuf.PushMessageContent.Flags.END_SESSION)
+            res[1]();
+
+        return finalMessage;
+    }
+
+    var decodeDeviceContents = function(res) {
+        var finalMessage = textsecure.protobuf.DeviceControl.decode(res[0]);
+
+        //TODO: Add END_SESSION flag for device control messages
         /*if ((finalMessage.flags & textsecure.protobuf.PushMessageContent.Flags.END_SESSION)
                 == textsecure.protobuf.PushMessageContent.Flags.END_SESSION)
-            axolotl.protocol.closeSession(res[1], true);*/
+            res[1]();*/
 
         return finalMessage;
     }
@@ -67,14 +77,10 @@
                 if (proto.message.readUint8() != ((3 << 4) | 3))
                     throw new Error("Bad version byte");
                 var from = proto.source + "." + (proto.sourceDevice == null ? 0 : proto.sourceDevice);
-                return axolotl.protocol.handlePreKeyWhisperMessage(from, getString(proto.message)).then(function(res) {
-                    return textsecure.protobuf.DeviceControl.decode(res[0]);
-                });
+                return axolotl.protocol.handlePreKeyWhisperMessage(from, getString(proto.message)).then(decodeDeviceContents);
             case textsecure.protobuf.IncomingPushMessageSignal.Type.DEVICE_CONTROL:
                 var from = proto.source + "." + (proto.sourceDevice == null ? 0 : proto.sourceDevice);
-                return axolotl.protocol.decryptWhisperMessage(from, getString(proto.message)).then(function(res) {
-                    return textsecure.protobuf.DeviceControl.decode(res[0]);
-                });
+                return axolotl.protocol.decryptWhisperMessage(from, getString(proto.message)).then(decodeDeviceContents);
             default:
                 return new Promise(function(resolve, reject) { reject(new Error("Unknown message type")); });
             }