From a5d857945a9300c74e4397ee0fdbcc5770cd3b4b Mon Sep 17 00:00:00 2001
From: Richard van der Hoff <richard@matrix.org>
Date: Mon, 14 Nov 2016 15:05:49 +0000
Subject: [PATCH] Retry decryption after receiving keys

m.room_keys may arrive after the messages themselves, so allow events to be
decrypted after the event (haha).
---
 lib/crypto/algorithms/megolm.js | 42 +++++++++++++++++++++++++++++++++
 lib/models/event.js             | 23 ++++++++++++++++--
 2 files changed, 63 insertions(+), 2 deletions(-)

diff --git a/lib/crypto/algorithms/megolm.js b/lib/crypto/algorithms/megolm.js
index d750845ff..7e87ad37c 100644
--- a/lib/crypto/algorithms/megolm.js
+++ b/lib/crypto/algorithms/megolm.js
@@ -445,6 +445,10 @@ MegolmEncryption.prototype.onNewDevice = function(userId, deviceId) {
  */
 function MegolmDecryption(params) {
     base.DecryptionAlgorithm.call(this, params);
+
+    // events which we couldn't decrypt due to unknown sessions / indexes: map from
+    // senderKey|sessionId to list of MatrixEvents
+    this._pendingEvents = {};
 }
 utils.inherits(MegolmDecryption, base.DecryptionAlgorithm);
 
@@ -474,11 +478,15 @@ MegolmDecryption.prototype.decryptEvent = function(event) {
             event.getRoomId(), content.sender_key, content.session_id, content.ciphertext
         );
     } catch (e) {
+        if (e.message === 'OLM.UNKNOWN_MESSAGE_INDEX') {
+            this._addEventToPendingList(event);
+        }
         throw new base.DecryptionError(e);
     }
 
     if (res === null) {
         // We've got a message for a session we don't have.
+        this._addEventToPendingList(event);
         throw new base.DecryptionError("Unknown inbound session id");
     }
 
@@ -496,6 +504,24 @@ MegolmDecryption.prototype.decryptEvent = function(event) {
     event.setClearData(payload, res.keysClaimed, res.keysProved);
 };
 
+
+/**
+ * Add an event to the list of those we couldn't decrypt the first time we
+ * saw them.
+ *
+ * @private
+ *
+ * @param {module:models/event.MatrixEvent} event
+ */
+MegolmDecryption.prototype._addEventToPendingList = function(event) {
+    var content = event.getWireContent();
+    var k = content.sender_key + "|" + content.session_id;
+    if (!this._pendingEvents[k]) {
+        this._pendingEvents[k] = [];
+    }
+    this._pendingEvents[k].push(event);
+};
+
 /**
  * @inheritdoc
  *
@@ -517,6 +543,22 @@ MegolmDecryption.prototype.onRoomKeyEvent = function(event) {
         content.room_id, event.getSenderKey(), content.session_id,
         content.session_key, event.getKeysClaimed()
     );
+
+    var k = event.getSenderKey() + "|" + content.session_id;
+    var pending = this._pendingEvents[k];
+    if (pending) {
+        // have another go at decrypting events sent with this session.
+        delete this._pendingEvents[k];
+
+        for (var i = 0; i < pending.length; i++) {
+            try {
+                this.decryptEvent(pending[i]);
+                console.log("successful re-decryption of", pending[i]);
+            } catch (e) {
+                console.log("Still can't decrypt", pending[i], e.stack || e);
+            }
+        }
+    }
 };
 
 base.registerAlgorithm(
diff --git a/lib/models/event.js b/lib/models/event.js
index 79eca1ad5..42d1b1299 100644
--- a/lib/models/event.js
+++ b/lib/models/event.js
@@ -21,6 +21,10 @@ limitations under the License.
  * @module models/event
  */
 
+var EventEmitter = require("events").EventEmitter;
+
+var utils = require('../utils.js');
+
 /**
  * Enum for event statuses.
  * @readonly
@@ -79,8 +83,10 @@ module.exports.MatrixEvent = function MatrixEvent(
     this._keysProved = {};
     this._keysClaimed = {};
 };
+utils.inherits(module.exports.MatrixEvent, EventEmitter);
 
-module.exports.MatrixEvent.prototype = {
+
+utils.extend(module.exports.MatrixEvent.prototype, {
 
     /**
      * Get the event_id for this event.
@@ -251,6 +257,7 @@ module.exports.MatrixEvent.prototype = {
         this._clearEvent = clearEvent;
         this._keysProved = keysProved || {};
         this._keysClaimed = keysClaimed || {};
+        this.emit("Event.decrypted", this);
     },
 
     /**
@@ -367,7 +374,7 @@ module.exports.MatrixEvent.prototype = {
      setPushActions: function(pushActions) {
         this._pushActions = pushActions;
      },
-};
+});
 
 
 /* http://matrix.org/docs/spec/r0.0.1/client_server.html#redactions says:
@@ -408,3 +415,15 @@ var _REDACT_KEEP_CONTENT_MAP = {
                            },
     'm.room.aliases': {'aliases': 1},
 };
+
+
+
+
+/**
+ * Fires when an event is decrypted
+ *
+ * @event module:models/event.MatrixEvent#"Event.decrypted"
+ *
+ * @param {module:models/event.MatrixEvent} event
+ *    The matrix event which has been decrypted
+ */