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
@@ -542,2 +542,16 @@ boost::optional<boost::posix_time::ptime> ChatController::getMessageTimestamp(st
 
+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 */) {
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
@@ -58,2 +58,4 @@ namespace Swift {
             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;
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
@@ -278,14 +278,8 @@ void ChatControllerBase::handleIncomingMessage(std::shared_ptr<MessageEvent> mes
         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);
         }
@@ -299,6 +293,2 @@ void ChatControllerBase::handleIncomingMessage(std::shared_ptr<MessageEvent> mes
 
-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) {
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
@@ -50,2 +50,9 @@ namespace Swift {
         public:
+            class StreamWindowMessageIDPair {
+                public:
+                    std::string idInStream;
+                    std::string idInWindow;
+            };
+
+        public:
             virtual ~ChatControllerBase();
@@ -84,3 +91,4 @@ namespace Swift {
             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&) {}
@@ -130,3 +138,3 @@ namespace Swift {
             bool labelsEnabled_;
-            std::map<JID, std::string> lastMessagesUIID_;
+            std::map<JID, StreamWindowMessageIDPair> lastMessagesIDs_;
             PresenceOracle* presenceOracle_;
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
@@ -589,3 +589,3 @@ 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()) {
@@ -594,3 +594,13 @@ void MUCController::addMessageHandleIncomingMessage(const JID& from, const ChatW
     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
@@ -83,3 +83,4 @@ namespace Swift {
             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;
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
@@ -35,2 +35,3 @@
 #include <Swiften/MUC/MUCManager.h>
+#include <Swiften/Network/DummyTimerFactory.h>
 #include <Swiften/Presence/DirectedPresenceSender.h>
@@ -43,3 +44,2 @@
 #include <Swiften/Whiteboard/WhiteboardSessionManager.h>
-#include <Swiften/Network/DummyTimerFactory.h>
 
@@ -141,4 +141,7 @@ class ChatsManagerTest : public CppUnit::TestFixture {
     // Message correction tests
+    CPPUNIT_TEST(testChatControllerMessageCorrectionCorrectReplaceID);
+    CPPUNIT_TEST(testChatControllerMessageCorrectionIncorrectReplaceID);
     CPPUNIT_TEST(testChatControllerMessageCorrectionReplaceBySameResource);
     CPPUNIT_TEST(testChatControllerMessageCorrectionReplaceByOtherResource);
+    CPPUNIT_TEST(testMUCControllerMessageCorrectionNoIDMatchRequired);
 
@@ -1311,2 +1314,69 @@ 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() {
@@ -1322,2 +1392,3 @@ public:
         message->setBody("text before edit");
+        message->setID("someID");
         manager_->handleIncomingMessage(message);
@@ -1348,2 +1419,3 @@ public:
         message->setBody("text before edit");
+        message->setID("someID");
         manager_->handleIncomingMessage(message);
@@ -1363,2 +1435,73 @@ public:
 
+    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: