From c5c9505ce2474ca9b1cb897438aaf801a69a9d14 Mon Sep 17 00:00:00 2001
From: Bruno Windels <brunow@matrix.org>
Date: Thu, 3 Sep 2020 17:50:28 +0200
Subject: [PATCH] hookup megolm encryption in session

---
 src/matrix/Session.js             | 14 ++++++++++++--
 src/matrix/e2ee/RoomEncryption.js | 28 +++++++++++++++++-----------
 2 files changed, 29 insertions(+), 13 deletions(-)

diff --git a/src/matrix/Session.js b/src/matrix/Session.js
index c1a26ded..e3f82bda 100644
--- a/src/matrix/Session.js
+++ b/src/matrix/Session.js
@@ -23,6 +23,7 @@ import {DeviceMessageHandler} from "./DeviceMessageHandler.js";
 import {Decryption as OlmDecryption} from "./e2ee/olm/Decryption.js";
 import {Encryption as OlmEncryption} from "./e2ee/olm/Encryption.js";
 import {Decryption as MegOlmDecryption} from "./e2ee/megolm/Decryption.js";
+import {Encryption as MegOlmEncryption} from "./e2ee/megolm/Encryption.js";
 import {RoomEncryption} from "./e2ee/RoomEncryption.js";
 import {DeviceTracker} from "./e2ee/DeviceTracker.js";
 import {LockMap} from "../utils/LockMap.js";
@@ -83,11 +84,19 @@ export class Session {
             olmUtil: this._olmUtil,
             senderKeyLock
         });
+        this._megolmEncryption = new MegOlmEncryption({
+            account: this._e2eeAccount,
+            pickleKey: PICKLE_KEY,
+            olm: this._olm,
+            storage: this._storage,
+            now: this._clock.now,
+            ownDeviceId: this._sessionInfo.deviceId,
+        })
         const megolmDecryption = new MegOlmDecryption({pickleKey: PICKLE_KEY, olm: this._olm});
         this._deviceMessageHandler.enableEncryption({olmDecryption, megolmDecryption});
     }
 
-    _createRoomEncryption(room, encryptionEventContent) {
+    _createRoomEncryption(room, encryptionParams) {
         // TODO: this will actually happen when users start using the e2ee version for the first time
 
         // this should never happen because either a session was already synced once
@@ -103,7 +112,8 @@ export class Session {
             room,
             deviceTracker: this._deviceTracker,
             olmEncryption: this._olmEncryption,
-            encryptionEventContent
+            megolmEncryption: this._megolmEncryption,
+            encryptionParams
         });
     }
 
diff --git a/src/matrix/e2ee/RoomEncryption.js b/src/matrix/e2ee/RoomEncryption.js
index 32bd061c..5f0c4cc1 100644
--- a/src/matrix/e2ee/RoomEncryption.js
+++ b/src/matrix/e2ee/RoomEncryption.js
@@ -17,14 +17,16 @@ limitations under the License.
 import {groupBy} from "../../utils/groupBy.js";
 import {makeTxnId} from "../common.js";
 
+const ENCRYPTED_TYPE = "m.room.encrypted";
 
 export class RoomEncryption {
-    constructor({room, deviceTracker, olmEncryption, encryptionEventContent}) {
+    constructor({room, deviceTracker, olmEncryption, megolmEncryption, encryptionParams}) {
         this._room = room;
         this._deviceTracker = deviceTracker;
         this._olmEncryption = olmEncryption;
+        this._megolmEncryption = megolmEncryption;
         // content of the m.room.encryption event
-        this._encryptionEventContent = encryptionEventContent;
+        this._encryptionParams = encryptionParams;
     }
 
     async writeMemberChanges(memberChanges, txn) {
@@ -32,15 +34,19 @@ export class RoomEncryption {
     }
 
     async encrypt(type, content, hsApi) {
-        await this._deviceTracker.trackRoom(this._room);
-        const devices = await this._deviceTracker.deviceIdentitiesForTrackedRoom(this._room.id, hsApi);
-        const messages = await this._olmEncryption.encrypt("m.foo", {body: "hello at " + new Date()}, devices, hsApi);
-        await this._sendMessagesToDevices("m.room.encrypted", messages, hsApi);
-        return {type, content};
-        // return {
-        //     type: "m.room.encrypted",
-        //     content: encryptedContent,
-        // }
+        const megolmResult = await this._megolmEncryption.encrypt(this._room.id, type, content, this._encryptionParams);
+        // share the new megolm session if needed
+        if (megolmResult.roomKeyMessage) {
+            await this._deviceTracker.trackRoom(this._room);
+            const devices = await this._deviceTracker.deviceIdentitiesForTrackedRoom(this._room.id, hsApi);
+            const messages = await this._olmEncryption.encrypt(
+                "m.room_key", megolmResult.roomKeyMessage, devices, hsApi);
+            await this._sendMessagesToDevices(ENCRYPTED_TYPE, messages, hsApi);
+        }
+        return {
+            type: ENCRYPTED_TYPE,
+            content: megolmResult.content
+        };
     }
 
     async _sendMessagesToDevices(type, messages, hsApi) {