From ff46d382acc476d06a8522c981686cee38f869ef Mon Sep 17 00:00:00 2001
From: Bruno Windels <bruno@windels.cloud>
Date: Thu, 10 Feb 2022 19:54:15 +0100
Subject: [PATCH] adjust m.direct when creating a DM

---
 src/matrix/Session.js               |  2 +-
 src/matrix/net/HomeServerApi.ts     |  3 +++
 src/matrix/room/RoomBeingCreated.ts | 35 +++++++++++++++++++++++++++++
 3 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/src/matrix/Session.js b/src/matrix/Session.js
index c5b5ad0d..be634754 100644
--- a/src/matrix/Session.js
+++ b/src/matrix/Session.js
@@ -624,8 +624,8 @@ export class Session {
             // the room id. Replace the room being created with the synced room.
             if (roomBeingCreated.roomId && !!this.rooms.get(roomBeingCreated.roomId)) {
                 this._tryReplaceRoomBeingCreated(roomBeingCreated.roomId, log);
+                await roomBeingCreated.adjustDirectMessageMapIfNeeded(this._user, this._storage, this._hsApi, log);
             }
-            // TODO: if type is DM, then adjust the m.direct account data
         });
         return roomBeingCreated;
     }
diff --git a/src/matrix/net/HomeServerApi.ts b/src/matrix/net/HomeServerApi.ts
index bf2c2c21..a6749cc4 100644
--- a/src/matrix/net/HomeServerApi.ts
+++ b/src/matrix/net/HomeServerApi.ts
@@ -286,6 +286,9 @@ export class HomeServerApi {
         return this._post(`/createRoom`, {}, payload, options);
     }
     
+    setAccountData(ownUserId: string, type: string, content: Record<string, any>, options?: IRequestOptions): IHomeServerRequest {
+        return this._put(`/user/${encodeURIComponent(ownUserId)}/account_data/${encodeURIComponent(type)}`, {}, content, options);
+    }
 }
 
 import {Request as MockRequest} from "../../mocks/Request.js";
diff --git a/src/matrix/room/RoomBeingCreated.ts b/src/matrix/room/RoomBeingCreated.ts
index 1ad84a47..9e564a87 100644
--- a/src/matrix/room/RoomBeingCreated.ts
+++ b/src/matrix/room/RoomBeingCreated.ts
@@ -26,6 +26,8 @@ import type {HomeServerApi} from "../net/HomeServerApi";
 import type {ILogItem} from "../../logging/types";
 import type {Platform} from "../../platform/web/Platform";
 import type {IBlobHandle} from "../../platform/types/types";
+import type {User} from "../User";
+import type {Storage} from "../storage/idb/Storage";
 
 type CreateRoomPayload = {
     is_direct?: boolean;
@@ -217,4 +219,37 @@ export class RoomBeingCreated extends EventEmitter<{change: never}> {
             this.options.avatar.blob.dispose();
         }
     }
+
+    async adjustDirectMessageMapIfNeeded(user: User, storage: Storage, hsApi: HomeServerApi, log: ILogItem): Promise<void> {
+        if (!this.options.invites || this.options.type !== RoomType.DirectMessage) {
+            return;
+        }
+        const userId = this.options.invites[0];
+        const DM_MAP_TYPE = "m.direct";
+        await log.wrap("set " + DM_MAP_TYPE, async log => {
+            try {
+                const txn = await storage.readWriteTxn([storage.storeNames.accountData]);
+                let mapEntry;
+                try {
+                    mapEntry = await txn.accountData.get(DM_MAP_TYPE);
+                    if (!mapEntry) {
+                        mapEntry = {type: DM_MAP_TYPE, content: {}};
+                    }
+                    const map = mapEntry.content;
+                    const userRooms = map[userId];
+                    // this is a new room id so no need to check if it's already there
+                    userRooms.push(this._roomId);
+                    txn.accountData.set(mapEntry);
+                    await txn.complete();
+                } catch (err) {
+                    txn.abort();
+                    throw err;
+                }
+                await hsApi.setAccountData(user.id, DM_MAP_TYPE, mapEntry.content, {log}).response();
+            } catch (err) {
+                // we can't really do anything else but logging here
+                log.catch(err);
+            }
+        });
+    }
 }