From d91380e1111caa38c00744005aaab3e5e412a374 Mon Sep 17 00:00:00 2001
From: varjolintu <sami.vanttinen@ahmala.org>
Date: Sun, 1 Oct 2023 15:57:32 +0300
Subject: [PATCH] Fix handling message timeouts

---
 keepassxc-browser/background/client.js  | 62 +++++++++++++++----------
 keepassxc-browser/background/keepass.js |  3 +-
 2 files changed, 40 insertions(+), 25 deletions(-)

diff --git a/keepassxc-browser/background/client.js b/keepassxc-browser/background/client.js
index 0e926df..0f6a5df 100644
--- a/keepassxc-browser/background/client.js
+++ b/keepassxc-browser/background/client.js
@@ -55,26 +55,38 @@ const messageBuffer = {
         this.buffer.push(message);
     },
 
-    removeMessageFromIndex(index) {
+    // Returns corresponding message from the response. If the response is an error,
+    // return the first matching action from the buffer.
+    getMessage(response) {
+        const isError = Boolean(!response.nonce && response.error && response.errorCode);
+        return this.buffer.find(b => keepassClient.incrementedNonce(b.request.nonce) === response.nonce
+                                || (isError && b.request?.action === response?.action));
+    },
+
+    removeMessage(message) {
+        const index = this.buffer.indexOf(message);
         if (index >= 0 && index < this.buffer.length) {
             this.buffer.splice(index, 1);
         }
-    }
+    },
 };
 
 // Basic class for a message to be sent. The Promise inside the class will be resolved when
 // the response to the message is received.
 class Message {
     constructor(request, enableTimeout, timeoutValue) {
+        this.enableTimeout = enableTimeout;
+        this.request = request;
+        this.timeout = undefined;
+
         this.promise = new Promise((resolve, reject) => {
             this.reject = reject;
             this.resolve = resolve;
-            this.enableTimeout = enableTimeout;
 
             const messageTimeout = timeoutValue || keepassClient.messageTimeout;
 
             // Handle timeout
-            if (enableTimeout) {
+            if (this.enableTimeout) {
                 this.timeout = setTimeout(() => {
                     const errorMessage = {
                         action: request.action,
@@ -88,6 +100,11 @@ class Message {
             }
         });
     }
+
+    cancelTimeout() {
+        this.enableTimeout = false;
+        clearTimeout(this.timeout);
+    }
 }
 
 //--------------------------------------------------------------------------
@@ -102,36 +119,33 @@ keepassClient.sendNativeMessage = async function(request, enableTimeout = false,
 
     const message = new Message(request, enableTimeout, timeoutValue);
     await navigator.locks.request('messageBuffer', async (lock) => {
-        messageBuffer.addMessage({ request: request, message: message });
+        messageBuffer.addMessage(message);
     });
 
     keepassClient.nativePort.postMessage(request);
-    return await message.promise;
+
+    const response = await message.promise;
+
+    // Remove a timeouted message
+    if (response.error && response?.errorCode === kpErrors.TIMEOUT_OR_NOT_CONNECTED) {
+        messageBuffer.removeMessage(message);
+    }
+
+    return response;
 };
 
 keepassClient.handleNativeMessage = async function(response) {
-    const isError = Boolean(!response.nonce && response.error && response.errorCode);
-
     // Parse through the message buffer to find the corresponding Promise.
     await navigator.locks.request('messageBuffer', async (lock) => {
-        for (let i = 0; i < messageBuffer.buffer.length; ++i) {
-            if (!messageBuffer.buffer[i]) {
-                continue;
+        const message = messageBuffer.getMessage(response);
+        if (message) {
+            if (message.enableTimeout) {
+                message.cancelTimeout();
             }
 
-            const request = messageBuffer.buffer[i]?.request;
-            const message = messageBuffer.buffer[i]?.message;
-            const errorFound = isError && request?.action === response?.action;
-
-            if ((response.nonce && response.nonce === keepassClient.incrementedNonce(request.nonce)) || errorFound) {
-                if (message.enableTimeout) {
-                    clearTimeout(message.timeout);
-                }
-
-                message.resolve(response);
-                messageBuffer.removeMessageFromIndex(i);
-                return;
-            }
+            message.resolve(response);
+            messageBuffer.removeMessage(message);
+            return;
         }
 
         debugLogMessage('Corresponding request not found in the message buffer for response: ', response);
diff --git a/keepassxc-browser/background/keepass.js b/keepassxc-browser/background/keepass.js
index 3a749f0..7e58389 100755
--- a/keepassxc-browser/background/keepass.js
+++ b/keepassxc-browser/background/keepass.js
@@ -387,8 +387,9 @@ keepass.getDatabaseHash = async function(tab, args = []) {
 
         keepass.databaseHash = '';
         keepass.isDatabaseClosed = true;
-        if (response.message && response.message === '') {
+        if ((response.message && response.message === '') || response.errorCode === kpErrors.TIMEOUT_OR_NOT_CONNECTED) {
             keepass.isKeePassXCAvailable = false;
+            keepass.isConnected = false;
             keepass.handleError(tab, kpErrors.TIMEOUT_OR_NOT_CONNECTED);
         } else {
             keepass.handleError(tab, response.errorCode, response.error);