Fix incoming message corrections in MUCs

2 files changed, 62 insertions, 45 deletions

diff --git a/libdino/src/service/content_item_store.vala b/libdino/src/service/content_item_store.vala
index b3ec5d67..11512a53 100644
--- a/libdino/src/service/content_item_store.vala
+++ b/libdino/src/service/content_item_store.vala
@@ -148,28 +148,41 @@ public class ContentItemStore : StreamInteractionModule, Object {
         return null;
     }
 
-    public ContentItem? get_content_item_for_referencing_id(Conversation conversation, string message_id) {
-        Row? row = get_content_item_row_for_referencing_id(conversation, message_id);
+    public ContentItem? get_content_item_for_message(Conversation conversation, Message message) {
+        Row? row = get_content_item_row_for_message(conversation, message);
+
         if (row != null) {
             return get_item_from_row(row, conversation);
         }
         return null;
     }
 
-    public int get_content_item_id_for_referencing_id(Conversation conversation, string message_id) {
-        Row? row = get_content_item_row_for_referencing_id(conversation, message_id);
+    public ContentItem? get_content_item_for_referencing_id(Conversation conversation, string referencing_id) {
+        Message? message = stream_interactor.get_module(MessageStorage.IDENTITY).get_message_by_referencing_id(referencing_id, conversation);
+        if (message == null) return null;
+
+        Row? row = get_content_item_row_for_message(conversation, message);
+
+        if (row != null) {
+            return get_item_from_row(row, conversation);
+        }
+        return null;
+    }
+
+    public int get_content_item_id_for_referencing_id(Conversation conversation, string referencing_id) {
+        Message? message = stream_interactor.get_module(MessageStorage.IDENTITY).get_message_by_referencing_id(referencing_id, conversation);
+        if (message == null) return -1;
+
+        Row? row = get_content_item_row_for_message(conversation, message);
         if (row != null) {
             return row[db.content_item.id];
         }
         return -1;
     }
 
-    private Row? get_content_item_row_for_referencing_id(Conversation conversation, string message_id) {
+    private Row? get_content_item_row_for_message(Conversation conversation, Message message) {
         var content_item_row = db.content_item.select();
 
-        Message? message = stream_interactor.get_module(MessageStorage.IDENTITY).get_message_by_referencing_id(message_id, conversation);
-        if (message == null) return null;
-
         RowOption file_transfer_row = db.file_transfer.select()
                 .with(db.file_transfer.account_id, "=", conversation.account.id)
                 .with(db.file_transfer.counterpart_id, "=", db.get_jid_id(conversation.counterpart))
@@ -182,7 +195,7 @@ public class ContentItemStore : StreamInteractionModule, Object {
                     .with(db.content_item.content_type, "=", 2);
         } else {
             // Check if this message has been corrected. In that case, the foreign id is the latest correction.
-            int correction_message_db_id = stream_interactor.get_module(MessageCorrection.IDENTITY).get_latest_correction_message_id(conversation, message_id);
+            int correction_message_db_id = stream_interactor.get_module(MessageCorrection.IDENTITY).get_latest_correction_message_id(conversation, message);
             int message_db_id = correction_message_db_id != -1 ? correction_message_db_id : message.id;
 
             content_item_row.with(db.content_item.foreign_id, "=", message_db_id)
diff --git a/libdino/src/service/message_correction.vala b/libdino/src/service/message_correction.vala
index 252830b4..6c4383ba 100644
--- a/libdino/src/service/message_correction.vala
+++ b/libdino/src/service/message_correction.vala
@@ -92,6 +92,7 @@ public class MessageCorrection : StreamInteractionModule, MessageListener {
     public override string[] after_actions { get { return after_actions_const; } }
 
     public override async bool run(Entities.Message message, Xmpp.MessageStanza stanza, Conversation conversation) {
+        // Check if we already know a newer correction for this message
         if (unmatched_corrections.has_key(conversation) && unmatched_corrections[conversation].size > 0) {
             ContentItem? remove_from_list = null;
             bool? ret = null;
@@ -128,62 +129,62 @@ public class MessageCorrection : StreamInteractionModule, MessageListener {
         }
 
         if (replace_id != null) {
-            var correction_message = message;
-            correction_message.edit_to = replace_id; // This isn't persisted TODO
-
-            // Check if it's maybe accepted by the XEP rules (for MUCs without occupant id)
-            if (conversation.type_.is_muc_semantic()) {
-                if (last_messages.has_key(conversation) && last_messages[conversation].has_key(correction_message.from)) {
-                    var last_message = last_messages[conversation][correction_message.from];
-                    // TODO this should be the referencing id (-> in MUCs the server id), but this implementation is temporarily for backwards-compatibility.
-                    bool acceptable = last_message.stanza_id == replace_id;
-                    if (acceptable) {
-                        return process_in_order_correction(conversation, last_messages[conversation][correction_message.from], correction_message);
-                    }
-                }
-            }
-
-            Message? original_message = stream_interactor.get_module(MessageStorage.IDENTITY).get_message_by_stanza_id(replace_id, conversation);
-            if (original_message != null && is_correction_acceptable(original_message, correction_message)) {
-                return process_in_order_correction(conversation, original_message, correction_message);
-            }
-            return false;
+            message.edit_to = replace_id; // This isn't persisted here, but later after verifying that it's an allowed edit TODO?
+            return process_unverified_in_order_correction(conversation, message, replace_id);
         }
         return false;
     }
 
-    private bool process_wrong_order_correction(Conversation conversation, Message earlier_message, MessageItem correction_message) {
-        if (!is_correction_acceptable(earlier_message, correction_message.message)) {
+    private bool process_wrong_order_correction(Conversation conversation, Message earlier_message, MessageItem correction_item) {
+        if (!is_correction_acceptable(earlier_message, correction_item.message)) {
             return false;
         }
 
+        db.message_correction.insert()
+                .value(db.message_correction.message_id, correction_item.message.id)
+                .value(db.message_correction.to_stanza_id, correction_item.message.edit_to)
+                .perform();
+
         db.content_item.update()
-                .with(db.content_item.id, "=", correction_message.id)
+                .with(db.content_item.id, "=", correction_item.id)
                 .set(db.content_item.time, (long) earlier_message.time.to_unix())
                 .set(db.content_item.local_time, (long) earlier_message.local_time.to_unix())
                 .perform();
 
-        correction_message.time = earlier_message.time;
-
-        on_received_correction(conversation, correction_message.message.id);
+        on_received_correction(conversation, correction_item.message.id);
 
         return true;
     }
 
-    private bool process_in_order_correction(Conversation conversation, Message original_message, Message correction_message) {
-        int message_id_to_be_updated = get_latest_correction_message_id(conversation, correction_message.edit_to);
-        if (message_id_to_be_updated == -1) {
-            message_id_to_be_updated = original_message.id;
+    private bool process_unverified_in_order_correction(Conversation conversation, Message correction_message, string replace_id) {
+        // Legacy logic for MUCs without occupant id support
+        if (conversation.type_.is_muc_semantic()) {
+            if (last_messages.has_key(conversation) && last_messages[conversation].has_key(correction_message.from)) {
+                var last_message = last_messages[conversation][correction_message.from];
+                bool acceptable = last_message.stanza_id == replace_id;
+                if (acceptable) {
+                    return process_in_order_correction(conversation, last_messages[conversation][correction_message.from], correction_message);
+                }
+            }
         }
 
-        ContentItem? content_item = stream_interactor.get_module(ContentItemStore.IDENTITY).get_content_item_for_referencing_id(conversation, correction_message.edit_to);
+        Message? original_message = stream_interactor.get_module(MessageStorage.IDENTITY).get_message_by_stanza_id(replace_id, conversation);
+        if (original_message != null && is_correction_acceptable(original_message, correction_message)) {
+            correction_message.edit_to = replace_id;
+            return process_in_order_correction(conversation, original_message, correction_message);
+        }
+        return false;
+    }
+
+    private bool process_in_order_correction(Conversation conversation, Message original_message, Message correction_message) {
+        ContentItem? content_item = stream_interactor.get_module(ContentItemStore.IDENTITY).get_content_item_for_message(conversation, original_message);
 
         db.message_correction.insert()
                 .value(db.message_correction.message_id, correction_message.id)
                 .value(db.message_correction.to_stanza_id, correction_message.edit_to)
                 .perform();
 
-        int current_correction_message_id = get_latest_correction_message_id(conversation, correction_message.edit_to);
+        int current_correction_message_id = get_latest_correction_message_id(conversation, original_message);
         if (content_item != null) {
             db.content_item.update()
                     .with(db.content_item.id, "=", content_item.id)
@@ -208,16 +209,19 @@ public class MessageCorrection : StreamInteractionModule, MessageListener {
         }
     }
 
-    public int get_latest_correction_message_id(Conversation conversation, string stanza_ref_id) {
+    public int get_latest_correction_message_id(Conversation conversation, Message message) {
         var qry = db.message_correction.select({db.message.id})
                 .join_with(db.message, db.message.id, db.message_correction.message_id)
                 .with(db.message.account_id, "=", conversation.account.id)
                 .with(db.message.counterpart_id, "=", db.get_jid_id(conversation.counterpart))
-                .with(db.message_correction.to_stanza_id, "=", stanza_ref_id)
+                .with(db.message_correction.to_stanza_id, "=", message.edit_to ?? message.stanza_id)
                 .order_by(db.message.time, "DESC");
 
-        if (conversation.counterpart.resourcepart != null) {
-            qry.with(db.message.counterpart_resource, "=", conversation.counterpart.resourcepart);
+        if (message.occupant_db_id != -1) {
+            qry.outer_join_with(db.message_occupant_id, db.message_occupant_id.message_id, db.message.id)
+                .with(db.message_occupant_id.occupant_id, "=", message.occupant_db_id);
+        } else if (message.counterpart.resourcepart != null) {
+            qry.with(db.message.counterpart_resource, "=", message.counterpart.resourcepart);
         }
         RowOption row = qry.single().row();
         if (row.is_present()) {