From 1c0bfa30aa1eaaa136035d28ffda4346f4e064aa Mon Sep 17 00:00:00 2001
From: Linus Jahn <lnj@kaidan.im>
Date: Mon, 28 Oct 2024 21:30:59 +0100
Subject: [PATCH] MamManager: Fix unencrypted messages not decrypted with e2ee

There were also some async issues that are now solved, especially when
all decryption jobs finished instantly (or no message were encrypted).
---
 src/client/QXmppMamManager.cpp | 86 +++++++++++++++++-----------------
 1 file changed, 42 insertions(+), 44 deletions(-)

diff --git a/src/client/QXmppMamManager.cpp b/src/client/QXmppMamManager.cpp
index cc8ca74df..cc9c0e298 100644
--- a/src/client/QXmppMamManager.cpp
+++ b/src/client/QXmppMamManager.cpp
@@ -313,58 +313,56 @@ QXmppTask<QXmppMamManager::RetrieveResult> QXmppMamManager::retrieveMessages(con
             // initialize processed messages (we need random access because
             // decryptMessage() may finish in random order)
             state.processedMessages.resize(state.messages.size());
-
-            // check for encrypted messages (once)
-            auto messagesEncrypted = transform<std::vector<bool>>(state.messages, [&](const auto &m) {
-                return e2eeExt->isEncrypted(m.element);
-            });
-            auto encryptedCount = sum(messagesEncrypted);
-
-            // We can't do this on the fly (with ++ and --) in the for loop
-            // because some decryptMessage() jobs could finish instantly
-            state.runningDecryptionJobs = encryptedCount;
-
-            int size = state.messages.size();
-            for (auto i = 0; i < size; i++) {
-                if (!messagesEncrypted[i]) {
-                    continue;
-                }
-
-                e2eeExt->decryptMessage(parseMamMessage(state.messages.at(i), Encrypted)).then(this, [this, i, queryId](auto result) {
-                    auto itr = d->ongoingRequests.find(queryId.toStdString());
-                    Q_ASSERT(itr != d->ongoingRequests.end());
-
-                    auto &state = itr->second;
-
-                    // store decrypted message, fallback to encrypted message
-                    if (std::holds_alternative<QXmppMessage>(result)) {
-                        state.processedMessages[i] = std::get<QXmppMessage>(std::move(result));
-                    } else {
-                        warning(u"Error decrypting message."_s);
-                        state.processedMessages[i] = parseMamMessage(state.messages[i], Unencrypted);
-                    }
-
-                    // finish promise if this was the last job
+            state.runningDecryptionJobs = state.messages.size();
+
+            for (qsizetype i = 0; i < state.messages.size(); i++) {
+                const auto &message = state.messages.at(i);
+
+                // decrypt message if needed
+                if (e2eeExt->isEncrypted(message.element)) {
+                    e2eeExt->decryptMessage(parseMamMessage(state.messages.at(i), Encrypted)).then(this, [this, i, queryId](auto result) {
+                        // find state (again)
+                        auto itr = d->ongoingRequests.find(queryId.toStdString());
+                        Q_ASSERT(itr != d->ongoingRequests.end());
+
+                        auto &state = itr->second;
+
+                        // store decrypted message, fallback to encrypted message
+                        if (std::holds_alternative<QXmppMessage>(result)) {
+                            state.processedMessages[i] = std::get<QXmppMessage>(std::move(result));
+                        } else {
+                            warning(u"Error decrypting message."_s);
+                            state.processedMessages[i] = parseMamMessage(state.messages[i], Unencrypted);
+                        }
+
+                        // finish promise on last job
+                        state.runningDecryptionJobs--;
+                        if (state.runningDecryptionJobs == 0) {
+                            state.finish();
+                            d->ongoingRequests.erase(itr);
+                        }
+                    });
+                } else {
+                    state.processedMessages[i] = parseMamMessage(state.messages.at(i), Unencrypted);
+
+                    // finish promise on last job (may be needed if no messages are encrypted or
+                    // decryption finishes instantly)
                     state.runningDecryptionJobs--;
                     if (state.runningDecryptionJobs == 0) {
                         state.finish();
                         d->ongoingRequests.erase(itr);
                     }
-                });
+                }
             }
+        } else {
+            // for the case without decryption, finish here
+            state.processedMessages = transform<QVector<QXmppMessage>>(state.messages, [](const auto &m) {
+                return parseMamMessage(m, Unencrypted);
+            });
 
-            // finishing the promise is done after decryptMessage()
-            if (encryptedCount > 0) {
-                return;
-            }
+            state.finish();
+            d->ongoingRequests.erase(itr);
         }
-
-        // for the case without decryption, finish here
-        state.processedMessages = transform<QVector<QXmppMessage>>(state.messages, [](const auto &m) {
-            return parseMamMessage(m, Unencrypted);
-        });
-        state.finish();
-        d->ongoingRequests.erase(itr);
     });
 
     return task;