7 files changed, 187 insertions, 19 deletions
diff --git a/Swift/Controllers/Chat/ChatController.cpp b/Swift/Controllers/Chat/ChatController.cpp
index 3e58e40..b6ca984 100644
--- a/Swift/Controllers/Chat/ChatController.cpp
+++ b/Swift/Controllers/Chat/ChatController.cpp
@@ -540,6 +540,20 @@ boost::optional<boost::posix_time::ptime> ChatController::getMessageTimestamp(st
     return message->getTimestamp();
 }
 
+void ChatController::addMessageHandleIncomingMessage(const JID& from, const ChatWindow::ChatMessage& message, const std::string& messageID, bool senderIsSelf, std::shared_ptr<SecurityLabel> label, const boost::posix_time::ptime& timeStamp) {
+    lastMessagesIDs_[from.toBare()] = {messageID, addMessage(message, senderDisplayNameFromMessage(from), senderIsSelf, label, avatarManager_->getAvatarPath(from), timeStamp)};
+}
+
+void ChatController::handleIncomingReplaceMessage(const JID& from, const ChatWindow::ChatMessage& message, const std::string& messageID, const std::string& idToReplace, bool senderIsSelf, std::shared_ptr<SecurityLabel> label, const boost::posix_time::ptime& timeStamp) {
+    auto lastMessage = lastMessagesIDs_.find(from.toBare());
+    if ((lastMessage != lastMessagesIDs_.end()) && (lastMessage->second.idInStream == idToReplace)) {
+        replaceMessage(message, lastMessage->second.idInWindow, timeStamp);
+    }
+    else {
+        addMessageHandleIncomingMessage(from, message, messageID, senderIsSelf, label, timeStamp);
+    }
+}
+
 void ChatController::logMessage(const std::string& message, const JID& fromJID, const JID& toJID, const boost::posix_time::ptime& timeStamp, bool /* isIncoming */) {
     HistoryMessage::Type type;
     if (mucRegistry_->isMUC(fromJID.toBare()) || mucRegistry_->isMUC(toJID.toBare())) {
diff --git a/Swift/Controllers/Chat/ChatController.h b/Swift/Controllers/Chat/ChatController.h
index 1deb0bb..d5011e4 100644
--- a/Swift/Controllers/Chat/ChatController.h
+++ b/Swift/Controllers/Chat/ChatController.h
@@ -56,6 +56,8 @@ namespace Swift {
             virtual bool isIncomingMessageFromMe(std::shared_ptr<Message> message) SWIFTEN_OVERRIDE;
             virtual void postSendMessage(const std::string &body, std::shared_ptr<Stanza> sentStanza) SWIFTEN_OVERRIDE;
             virtual void preHandleIncomingMessage(std::shared_ptr<MessageEvent> messageEvent) SWIFTEN_OVERRIDE;
+            virtual void addMessageHandleIncomingMessage(const JID& from, const ChatWindow::ChatMessage& message, const std::string& messageID, bool senderIsSelf, std::shared_ptr<SecurityLabel> label, const boost::posix_time::ptime& timeStamp) SWIFTEN_OVERRIDE;
+            virtual void handleIncomingReplaceMessage(const JID& from, const ChatWindow::ChatMessage& message, const std::string& messageID, const std::string& idToReplace, bool senderIsSelf, std::shared_ptr<SecurityLabel> label, const boost::posix_time::ptime& timeStamp) SWIFTEN_OVERRIDE;
             virtual void postHandleIncomingMessage(std::shared_ptr<MessageEvent> messageEvent, const ChatWindow::ChatMessage& chatMessage) SWIFTEN_OVERRIDE;
             virtual void preSendMessageRequest(std::shared_ptr<Message>) SWIFTEN_OVERRIDE;
             virtual std::string senderHighlightNameFromMessage(const JID& from) SWIFTEN_OVERRIDE;
diff --git a/Swift/Controllers/Chat/ChatControllerBase.cpp b/Swift/Controllers/Chat/ChatControllerBase.cpp
index 7ae7dbd..0fc735a 100644
--- a/Swift/Controllers/Chat/ChatControllerBase.cpp
+++ b/Swift/Controllers/Chat/ChatControllerBase.cpp
@@ -276,18 +276,12 @@ void ChatControllerBase::handleIncomingMessage(std::shared_ptr<MessageEvent> mes
 
         std::shared_ptr<Replace> replace = message->getPayload<Replace>();
         bool senderIsSelf = isIncomingMessageFromMe(message);
+        chatMessage = buildChatWindowChatMessage(body, senderHighlightNameFromMessage(from), senderIsSelf);
         if (replace) {
-            // Should check if the user has a previous message
-            std::map<JID, std::string>::iterator lastMessage;
-            lastMessage = lastMessagesUIID_.find(messageCorrectionJID(from));
-            if (lastMessage != lastMessagesUIID_.end()) {
-                chatMessage = buildChatWindowChatMessage(body, senderHighlightNameFromMessage(from), senderIsSelf);
-                replaceMessage(chatMessage, lastMessagesUIID_[messageCorrectionJID(from)], timeStamp);
-            }
+            handleIncomingReplaceMessage(from, chatMessage, message->getID(), replace->getID(), senderIsSelf, label, timeStamp);
         }
         else {
-            chatMessage = buildChatWindowChatMessage(body, senderHighlightNameFromMessage(from), senderIsSelf);
-            addMessageHandleIncomingMessage(from, chatMessage, senderIsSelf, label, timeStamp);
+            addMessageHandleIncomingMessage(from, chatMessage, message->getID(), senderIsSelf, label, timeStamp);
         }
 
         logMessage(body, from, selfJID_, timeStamp, true);
@@ -297,10 +291,6 @@ void ChatControllerBase::handleIncomingMessage(std::shared_ptr<MessageEvent> mes
     postHandleIncomingMessage(messageEvent, chatMessage);
 }
 
-void ChatControllerBase::addMessageHandleIncomingMessage(const JID& from, const ChatWindow::ChatMessage& message, bool senderIsSelf, std::shared_ptr<SecurityLabel> label, const boost::posix_time::ptime& timeStamp) {
-    lastMessagesUIID_[messageCorrectionJID(from)] = addMessage(message, senderDisplayNameFromMessage(from), senderIsSelf, label, avatarManager_->getAvatarPath(from), timeStamp);
-}
-
 std::string ChatControllerBase::getErrorMessage(std::shared_ptr<ErrorPayload> error) {
     std::string defaultMessage = QT_TRANSLATE_NOOP("", "Error sending message");
     if (!error->getText().empty()) {
diff --git a/Swift/Controllers/Chat/ChatControllerBase.h b/Swift/Controllers/Chat/ChatControllerBase.h
index d10df8f..88cb95c 100644
--- a/Swift/Controllers/Chat/ChatControllerBase.h
+++ b/Swift/Controllers/Chat/ChatControllerBase.h
@@ -48,6 +48,13 @@ namespace Swift {
 
     class ChatControllerBase : public boost::signals2::trackable {
         public:
+            class StreamWindowMessageIDPair {
+                public:
+                    std::string idInStream;
+                    std::string idInWindow;
+            };
+
+        public:
             virtual ~ChatControllerBase();
             void showChatWindow();
             void activateChatWindow();
@@ -82,7 +89,8 @@ namespace Swift {
             virtual std::string senderHighlightNameFromMessage(const JID& from) = 0;
             virtual bool isIncomingMessageFromMe(std::shared_ptr<Message>) = 0;
             virtual void preHandleIncomingMessage(std::shared_ptr<MessageEvent>) {}
-            virtual void addMessageHandleIncomingMessage(const JID& from, const ChatWindow::ChatMessage& message, bool senderIsSelf, std::shared_ptr<SecurityLabel> label, const boost::posix_time::ptime& time);
+            virtual void addMessageHandleIncomingMessage(const JID& from, const ChatWindow::ChatMessage& message, const std::string& messageID, bool senderIsSelf, std::shared_ptr<SecurityLabel> label, const boost::posix_time::ptime& time) = 0;
+            virtual void handleIncomingReplaceMessage(const JID& from, const ChatWindow::ChatMessage& chatMessage, const std::string& messageID, const std::string& idToReplace, bool senderIsSelf, std::shared_ptr<SecurityLabel> label, const boost::posix_time::ptime& timeStamp) = 0;
             virtual void postHandleIncomingMessage(std::shared_ptr<MessageEvent>, const ChatWindow::ChatMessage&) {}
             virtual void preSendMessageRequest(std::shared_ptr<Message>) {}
             virtual bool isFromContact(const JID& from);
@@ -128,7 +136,7 @@ namespace Swift {
             ChatWindow* chatWindow_;
             JID toJID_;
             bool labelsEnabled_;
-            std::map<JID, std::string> lastMessagesUIID_;
+            std::map<JID, StreamWindowMessageIDPair> lastMessagesIDs_;
             PresenceOracle* presenceOracle_;
             AvatarManager* avatarManager_;
             bool useDelayForLatency_;
diff --git a/Swift/Controllers/Chat/MUCController.cpp b/Swift/Controllers/Chat/MUCController.cpp
index b685e04..fd3dc37 100644
--- a/Swift/Controllers/Chat/MUCController.cpp
+++ b/Swift/Controllers/Chat/MUCController.cpp
@@ -587,12 +587,22 @@ void MUCController::preHandleIncomingMessage(std::shared_ptr<MessageEvent> messa
     }
 }
 
-void MUCController::addMessageHandleIncomingMessage(const JID& from, const ChatWindow::ChatMessage& message, bool senderIsSelf, std::shared_ptr<SecurityLabel> label, const boost::posix_time::ptime& time) {
+void MUCController::addMessageHandleIncomingMessage(const JID& from, const ChatWindow::ChatMessage& message, const std::string& messageID, bool senderIsSelf, std::shared_ptr<SecurityLabel> label, const boost::posix_time::ptime& time) {
     if (from.isBare()) {
         chatWindow_->addSystemMessage(message, ChatWindow::DefaultDirection);
     }
     else {
-        ChatControllerBase::addMessageHandleIncomingMessage(from, message, senderIsSelf, label, time);
+        lastMessagesIDs_[from] = {messageID, addMessage(message, senderDisplayNameFromMessage(from), senderIsSelf, label, avatarManager_->getAvatarPath(from), time)};
+    }
+}
+
+void MUCController::handleIncomingReplaceMessage(const JID& from, const ChatWindow::ChatMessage& message, const std::string& messageID, const std::string& /*idToReplace*/, bool senderIsSelf, std::shared_ptr<SecurityLabel> label, const boost::posix_time::ptime& timeStamp) {
+    auto lastMessage = lastMessagesIDs_.find(from);
+    if (lastMessage != lastMessagesIDs_.end()) {
+        replaceMessage(message, lastMessage->second.idInWindow, timeStamp);
+    }
+    else {
+        addMessageHandleIncomingMessage(from, message, messageID, senderIsSelf, label, timeStamp);
     }
 }
 
diff --git a/Swift/Controllers/Chat/MUCController.h b/Swift/Controllers/Chat/MUCController.h
index ba68ec6..b8f2aac 100644
--- a/Swift/Controllers/Chat/MUCController.h
+++ b/Swift/Controllers/Chat/MUCController.h
@@ -81,7 +81,8 @@ namespace Swift {
             virtual std::string senderDisplayNameFromMessage(const JID& from) SWIFTEN_OVERRIDE;
             virtual boost::optional<boost::posix_time::ptime> getMessageTimestamp(std::shared_ptr<Message> message) const SWIFTEN_OVERRIDE;
             virtual void preHandleIncomingMessage(std::shared_ptr<MessageEvent>) SWIFTEN_OVERRIDE;
-            virtual void addMessageHandleIncomingMessage(const JID& from, const ChatWindow::ChatMessage& message, bool senderIsSelf, std::shared_ptr<SecurityLabel> label, const boost::posix_time::ptime& time) SWIFTEN_OVERRIDE;
+            virtual void addMessageHandleIncomingMessage(const JID& from, const ChatWindow::ChatMessage& message, const std::string& messageID, bool senderIsSelf, std::shared_ptr<SecurityLabel> label, const boost::posix_time::ptime& time) SWIFTEN_OVERRIDE;
+            virtual void handleIncomingReplaceMessage(const JID& from, const ChatWindow::ChatMessage& message, const std::string& messageID, const std::string& idToReplace, bool senderIsSelf, std::shared_ptr<SecurityLabel> label, const boost::posix_time::ptime& timeStamp) SWIFTEN_OVERRIDE;
             virtual void postHandleIncomingMessage(std::shared_ptr<MessageEvent>, const ChatWindow::ChatMessage& chatMessage) SWIFTEN_OVERRIDE;
             virtual void cancelReplaces() SWIFTEN_OVERRIDE;
             virtual void logMessage(const std::string& message, const JID& fromJID, const JID& toJID, const boost::posix_time::ptime& timeStamp, bool isIncoming) SWIFTEN_OVERRIDE;
diff --git a/Swift/Controllers/Chat/UnitTest/ChatsManagerTest.cpp b/Swift/Controllers/Chat/UnitTest/ChatsManagerTest.cpp
index 0356c6a..bf645d0 100644
--- a/Swift/Controllers/Chat/UnitTest/ChatsManagerTest.cpp
+++ b/Swift/Controllers/Chat/UnitTest/ChatsManagerTest.cpp
@@ -33,6 +33,7 @@
 #include <Swiften/FileTransfer/UnitTest/DummyFileTransferManager.h>
 #include <Swiften/Jingle/JingleSessionManager.h>
 #include <Swiften/MUC/MUCManager.h>
+#include <Swiften/Network/DummyTimerFactory.h>
 #include <Swiften/Presence/DirectedPresenceSender.h>
 #include <Swiften/Presence/PresenceOracle.h>
 #include <Swiften/Presence/StanzaChannelPresenceSender.h>
@@ -41,7 +42,6 @@
 #include <Swiften/VCards/VCardManager.h>
 #include <Swiften/VCards/VCardMemoryStorage.h>
 #include <Swiften/Whiteboard/WhiteboardSessionManager.h>
-#include <Swiften/Network/DummyTimerFactory.h>
 
 #include <Swift/Controllers/Chat/ChatController.h>
 #include <Swift/Controllers/Chat/ChatsManager.h>
@@ -139,8 +139,11 @@ class ChatsManagerTest : public CppUnit::TestFixture {
 
 
     // Message correction tests
+    CPPUNIT_TEST(testChatControllerMessageCorrectionCorrectReplaceID);
+    CPPUNIT_TEST(testChatControllerMessageCorrectionIncorrectReplaceID);
     CPPUNIT_TEST(testChatControllerMessageCorrectionReplaceBySameResource);
     CPPUNIT_TEST(testChatControllerMessageCorrectionReplaceByOtherResource);
+    CPPUNIT_TEST(testMUCControllerMessageCorrectionNoIDMatchRequired);
 
     CPPUNIT_TEST_SUITE_END();
 
@@ -1309,6 +1312,73 @@ public:
         }
     }
 
+    /**
+     * This test case ensures correct handling of the ideal case where the replace
+     * message refers to a message with a known ID. This results in the last
+     * message being replaced.
+     */
+    void testChatControllerMessageCorrectionCorrectReplaceID() {
+        JID messageJID("testling@test.com/resource1");
+
+        MockChatWindow* window = new MockChatWindow();
+        mocks_->ExpectCall(chatWindowFactory_, ChatWindowFactory::createChatWindow).With(messageJID, uiEventStream_).Return(window);
+
+        auto message = std::make_shared<Message>();
+        message->setFrom(messageJID);
+        message->setTo(jid_);
+        message->setType(Message::Chat);
+        message->setBody("text before edit");
+        message->setID("someID");
+        manager_->handleIncomingMessage(message);
+
+        CPPUNIT_ASSERT_EQUAL(std::string("text before edit"), MockChatWindow::bodyFromMessage(window->lastAddedMessage_));
+
+        message = std::make_shared<Message>();
+        message->setFrom(messageJID);
+        message->setTo(jid_);
+        message->setType(Message::Chat);
+        message->setBody("text after edit");
+        message->addPayload(std::make_shared<Replace>("someID"));
+        manager_->handleIncomingMessage(message);
+
+        CPPUNIT_ASSERT_EQUAL(std::string("text before edit"), MockChatWindow::bodyFromMessage(window->lastAddedMessage_));
+        CPPUNIT_ASSERT_EQUAL(std::string("text after edit"), MockChatWindow::bodyFromMessage(window->lastReplacedMessage_));
+    }
+
+    /**
+     * This test case ensures correct handling of the case where the replace
+     * message refers to a message with a unknown ID. The replace message should
+     * be treated like a non-repalce message in this case, with no replacement
+     * occuring.
+     */
+    void testChatControllerMessageCorrectionIncorrectReplaceID() {
+        JID messageJID("testling@test.com/resource1");
+
+        MockChatWindow* window = new MockChatWindow();
+        mocks_->ExpectCall(chatWindowFactory_, ChatWindowFactory::createChatWindow).With(messageJID, uiEventStream_).Return(window);
+
+        auto message = std::make_shared<Message>();
+        message->setFrom(messageJID);
+        message->setTo(jid_);
+        message->setType(Message::Chat);
+        message->setBody("text before edit");
+        message->setID("someID");
+        manager_->handleIncomingMessage(message);
+
+        CPPUNIT_ASSERT_EQUAL(std::string("text before edit"), MockChatWindow::bodyFromMessage(window->lastAddedMessage_));
+
+        message = std::make_shared<Message>();
+        message->setFrom(messageJID);
+        message->setTo(jid_);
+        message->setType(Message::Chat);
+        message->setBody("text after failed edit");
+        message->addPayload(std::make_shared<Replace>("wrongID"));
+        manager_->handleIncomingMessage(message);
+
+        CPPUNIT_ASSERT_EQUAL(std::string("text after failed edit"), MockChatWindow::bodyFromMessage(window->lastAddedMessage_));
+        CPPUNIT_ASSERT_EQUAL(std::string(""), MockChatWindow::bodyFromMessage(window->lastReplacedMessage_));
+    }
+
     void testChatControllerMessageCorrectionReplaceBySameResource() {
         JID messageJID("testling@test.com/resource1");
 
@@ -1320,6 +1390,7 @@ public:
         message->setTo(jid_);
         message->setType(Message::Chat);
         message->setBody("text before edit");
+        message->setID("someID");
         manager_->handleIncomingMessage(message);
 
         CPPUNIT_ASSERT_EQUAL(std::string("text before edit"), MockChatWindow::bodyFromMessage(window->lastAddedMessage_));
@@ -1346,6 +1417,7 @@ public:
         message->setTo(jid_);
         message->setType(Message::Chat);
         message->setBody("text before edit");
+        message->setID("someID");
         manager_->handleIncomingMessage(message);
 
         CPPUNIT_ASSERT_EQUAL(std::string("text before edit"), MockChatWindow::bodyFromMessage(window->lastAddedMessage_));
@@ -1361,6 +1433,77 @@ public:
         CPPUNIT_ASSERT_EQUAL(std::string("text after edit"), MockChatWindow::bodyFromMessage(window->lastReplacedMessage_));
     }
 
+    void testMUCControllerMessageCorrectionNoIDMatchRequired() {
+        JID mucJID("SomeMUCRoom@test.com");
+        manager_->setOnline(true);
+
+        // Open chat window to a sender.
+        MockChatWindow* window = new MockChatWindow();
+
+        std::vector<JID> jids;
+        jids.emplace_back("foo@test.com");
+        jids.emplace_back("bar@test.com");
+
+        mocks_->ExpectCall(chatWindowFactory_, ChatWindowFactory::createChatWindow).With(mucJID, uiEventStream_).Return(window);
+
+        auto nickname = std::string("SomeNickName");
+        // Join room
+        {
+            auto joinRoomEvent = std::make_shared<JoinMUCUIEvent>(mucJID, boost::optional<std::string>(), nickname);
+            uiEventStream_->send(joinRoomEvent);
+        }
+
+        auto genRemoteMUCPresence = [=]() {
+            auto presence = Presence::create();
+            presence->setFrom(mucJID.withResource(nickname));
+            presence->setTo(jid_);
+            return presence;
+        };
+
+        {
+            auto presence = genRemoteMUCPresence();
+            auto userPayload = std::make_shared<MUCUserPayload>();
+            userPayload->addStatusCode(110);
+            userPayload->addItem(MUCItem(MUCOccupant::Owner, jid_, MUCOccupant::Moderator));
+            presence->addPayload(userPayload);
+            stanzaChannel_->onPresenceReceived(presence);
+        }
+
+        {
+            auto presence = genRemoteMUCPresence();
+            presence->setFrom(mucJID.withResource("someDifferentNickname"));
+            auto userPayload = std::make_shared<MUCUserPayload>();
+            userPayload->addItem(MUCItem(MUCOccupant::Member, JID("foo@bar.com"), MUCOccupant::Moderator));
+            presence->addPayload(userPayload);
+            stanzaChannel_->onPresenceReceived(presence);
+        }
+
+        {
+            Message::ref mucMirrored = std::make_shared<Message>();
+            mucMirrored->setFrom(mucJID.withResource(nickname));
+            mucMirrored->setTo(jid_);
+            mucMirrored->setType(Message::Groupchat);
+            mucMirrored->setID("fooBlaID_1");
+            mucMirrored->setBody("Some misssssspelled message.");
+            manager_->handleIncomingMessage(mucMirrored);
+        }
+        CPPUNIT_ASSERT_EQUAL(std::string("Some misssssspelled message."), window->bodyFromMessage(window->lastAddedMessage_));
+
+        // Replace message with non-matching ID
+        {
+            Message::ref mucMirrored = std::make_shared<Message>();
+            mucMirrored->setFrom(mucJID.withResource(nickname));
+            mucMirrored->setTo(jid_);
+            mucMirrored->setType(Message::Groupchat);
+            mucMirrored->setID("fooBlaID_3");
+            mucMirrored->setBody("Some correctly spelled message.");
+            mucMirrored->addPayload(std::make_shared<Replace>("fooBlaID_2"));
+            manager_->handleIncomingMessage(mucMirrored);
+        }
+        CPPUNIT_ASSERT_EQUAL(std::string("Some correctly spelled message."), window->bodyFromMessage(window->lastReplacedMessage_));
+    }
+
+
 private:
     std::shared_ptr<Message> makeDeliveryReceiptTestMessage(const JID& from, const std::string& id) {
         std::shared_ptr<Message> message = std::make_shared<Message>();