From 34f945b7f43989fd8723071bf078670b8f7985e6 Mon Sep 17 00:00:00 2001
From: hippoz <10706925-hippoz@users.noreply.gitlab.com>
Date: Sat, 20 Aug 2022 00:12:27 +0300
Subject: [PATCH] add optimistic_id field for messages, allowing us to show our
 own messages from other clients

---
 frontend/src/components/MessageInput.svelte |  3 ++-
 frontend/src/main.js                        |  4 ++--
 frontend/src/stores.js                      | 13 ++++++++----
 src/gateway/index.ts                        |  8 +++++--
 src/routes/api/v1/channels.ts               | 23 +++++++++++++++++----
 5 files changed, 38 insertions(+), 13 deletions(-)

diff --git a/frontend/src/components/MessageInput.svelte b/frontend/src/components/MessageInput.svelte
index a97e3a0..313eb85 100644
--- a/frontend/src/components/MessageInput.svelte
+++ b/frontend/src/components/MessageInput.svelte
@@ -60,7 +60,8 @@
         messageInput = "";
 
         const res = await request("POST", apiRoute(`channels/${channel.id}/messages`), true, {
-            content: optimisticMessage.content
+            content: optimisticMessage.content,
+            optimistic_id: optimisticMessageId
         });
         
         if (res.success && res.ok) {
diff --git a/frontend/src/main.js b/frontend/src/main.js
index e8d6bb4..52ff208 100644
--- a/frontend/src/main.js
+++ b/frontend/src/main.js
@@ -12,11 +12,11 @@ authWithToken(getItem("auth:token"));
 // Remove loading screen
 const loadingElement = document.getElementById("pre--loading-screen");
 if (loadingElement) {
-	loadingElement.parentElement.removeChild(loadingElement);
+    loadingElement.parentElement.removeChild(loadingElement);
 }
 
 const app = new Main({
-	target: document.body
+    target: document.body
 });
 window.__waffle.app = app;
 
diff --git a/frontend/src/stores.js b/frontend/src/stores.js
index 5ee0fc7..7e841b3 100644
--- a/frontend/src/stores.js
+++ b/frontend/src/stores.js
@@ -162,6 +162,15 @@ class MessageStore extends Store {
     }
 
     addMessage(message) {
+        if (message.optimistic_id) {
+            const index = this.value.findIndex(e => e.id === message.optimistic_id);
+            if (index !== -1) {
+                this.value[index] = message;
+                this.updated();
+                return;
+            }
+        }
+
         this.value.push(message);
         // only dispatch update if collectOldMessages didn't
         if (!this.collectOldMessages()) {
@@ -255,10 +264,6 @@ class MessagesStoreProvider {
         this.storeByChannel = new Map();
 
         gateway.subscribe(GatewayEventType.MessageCreate, (message) => {
-            // we currently don't care about our own messages
-            if (gateway.user && message.author_id === gateway.user.id)
-                return;
-
             const store = this.getStoreOrNull(message.channel_id);
             if (store)
                 store.addMessage(message);
diff --git a/src/gateway/index.ts b/src/gateway/index.ts
index a54d602..09ef7ad 100644
--- a/src/gateway/index.ts
+++ b/src/gateway/index.ts
@@ -67,13 +67,17 @@ function clientUnsubscribeAll(ws: WebSocket) {
     ws.state.dispatchChannels = new Set();
 }
 
-export function dispatch(channel: string, message: GatewayPayload) {
+export function dispatch(channel: string, message: GatewayPayload | ((ws: WebSocket) => GatewayPayload)) {
     const members = dispatchChannels.get(channel);
     if (!members) return;
         
     members.forEach(e => {
         if (e.state.ready) {
-            e.send(JSON.stringify(message));
+            let data = message;
+            if (typeof message === "function") {
+                data = message(e);
+            }
+            e.send(JSON.stringify(data));
         }
     });
 }
diff --git a/src/routes/api/v1/channels.ts b/src/routes/api/v1/channels.ts
index 5414ef4..f935495 100644
--- a/src/routes/api/v1/channels.ts
+++ b/src/routes/api/v1/channels.ts
@@ -173,12 +173,14 @@ router.post(
     authenticateRoute(),
     param("id").isNumeric(),
     body("content").isLength({ min: 1, max: 4000 }),
+    body("optimistic_id").optional().isNumeric(),
     async (req, res) => {
         const validationErrors = validationResult(req);
         if (!validationErrors.isEmpty()) {
             return res.status(400).json({ ...errors.INVALID_DATA, errors: validationErrors.array() });
         }
 
+        const optimisticId = parseInt(req.body.optimistic_id);
         const channelId = parseInt(req.params.id);
         const { content } = req.body;
         const authorId = req.user.id;
@@ -191,7 +193,7 @@ router.post(
             });
         }
 
-        const returnObject = {
+        let returnObject: any = {
             id: result.rows[0].id,
             content,
             channel_id: channelId,
@@ -200,9 +202,22 @@ router.post(
             created_at: createdAt
         };
 
-        dispatch(`channel:${channelId}`, {
-            t: GatewayPayloadType.MessageCreate,
-            d: returnObject
+        dispatch(`channel:${channelId}`, (ws) => {
+            let payload: any = returnObject;
+            if (ws.state && ws.state.user && ws.state.user.id === req.user.id && optimisticId) {
+                payload = {
+                    ...payload,
+                    optimistic_id: optimisticId
+                }
+                returnObject = {
+                    ...returnObject,
+                    optimistic_id: optimisticId
+                };
+            }
+            return {
+                t: GatewayPayloadType.MessageCreate,
+                d: payload
+            };
         });
 
         return res.status(201).send(returnObject);