7 files changed, 88 insertions, 17 deletions
diff --git a/Swift/Controllers/Chat/ChatController.cpp b/Swift/Controllers/Chat/ChatController.cpp
index 3bc9d20..3e58e40 100644
--- a/Swift/Controllers/Chat/ChatController.cpp
+++ b/Swift/Controllers/Chat/ChatController.cpp
@@ -565,6 +565,10 @@ bool ChatController::shouldIgnoreMessage(std::shared_ptr<Message> message) {
     }
     return false;
 }
+    
+JID ChatController::messageCorrectionJID(const JID& fromJID) {
+    return fromJID.toBare();
+}
 
 ChatWindow* ChatController::detachChatWindow() {
     chatWindow_->onUserTyping.disconnect(boost::bind(&ChatStateNotifier::setUserIsTyping, chatStateNotifier_));
diff --git a/Swift/Controllers/Chat/ChatController.h b/Swift/Controllers/Chat/ChatController.h
index ffc6be9..1deb0bb 100644
--- a/Swift/Controllers/Chat/ChatController.h
+++ b/Swift/Controllers/Chat/ChatController.h
@@ -48,6 +48,7 @@ namespace Swift {
             virtual JID getBaseJID() 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;
             virtual bool shouldIgnoreMessage(std::shared_ptr<Message> message) SWIFTEN_OVERRIDE;
+            virtual JID messageCorrectionJID(const JID& fromJID) SWIFTEN_OVERRIDE;
 
         private:
             void handlePresenceChange(std::shared_ptr<Presence> newPresence);
diff --git a/Swift/Controllers/Chat/ChatControllerBase.cpp b/Swift/Controllers/Chat/ChatControllerBase.cpp
index ecaf186..7ae7dbd 100644
--- a/Swift/Controllers/Chat/ChatControllerBase.cpp
+++ b/Swift/Controllers/Chat/ChatControllerBase.cpp
@@ -279,10 +279,10 @@ void ChatControllerBase::handleIncomingMessage(std::shared_ptr<MessageEvent> mes
         if (replace) {
             // Should check if the user has a previous message
             std::map<JID, std::string>::iterator lastMessage;
-            lastMessage = lastMessagesUIID_.find(from);
+            lastMessage = lastMessagesUIID_.find(messageCorrectionJID(from));
             if (lastMessage != lastMessagesUIID_.end()) {
                 chatMessage = buildChatWindowChatMessage(body, senderHighlightNameFromMessage(from), senderIsSelf);
-                replaceMessage(chatMessage, lastMessagesUIID_[from], timeStamp);
+                replaceMessage(chatMessage, lastMessagesUIID_[messageCorrectionJID(from)], timeStamp);
             }
         }
         else {
@@ -298,7 +298,7 @@ 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_[from] = addMessage(message, senderDisplayNameFromMessage(from), senderIsSelf, label, avatarManager_->getAvatarPath(from), 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 c84c4d4..d10df8f 100644
--- a/Swift/Controllers/Chat/ChatControllerBase.h
+++ b/Swift/Controllers/Chat/ChatControllerBase.h
@@ -101,6 +101,11 @@ namespace Swift {
                 return false;
             }
 
+            /**
+             * What JID should be used for last message correction (XEP-0308) tracking.
+             */
+            virtual JID messageCorrectionJID(const JID& fromJID) = 0;
+
         private:
             IDGenerator idGenerator_;
             std::string lastSentMessageStanzaID_;
diff --git a/Swift/Controllers/Chat/MUCController.cpp b/Swift/Controllers/Chat/MUCController.cpp
index 73cf748..b685e04 100644
--- a/Swift/Controllers/Chat/MUCController.cpp
+++ b/Swift/Controllers/Chat/MUCController.cpp
@@ -1078,6 +1078,10 @@ void MUCController::logMessage(const std::string& message, const JID& fromJID, c
     }
 }
 
+JID MUCController::messageCorrectionJID(const JID& fromJID) {
+    return fromJID;
+}
+
 void MUCController::addRecentLogs() {
     if (!historyController_) {
         return;
diff --git a/Swift/Controllers/Chat/MUCController.h b/Swift/Controllers/Chat/MUCController.h
index 9d6428b..ba68ec6 100644
--- a/Swift/Controllers/Chat/MUCController.h
+++ b/Swift/Controllers/Chat/MUCController.h
@@ -85,6 +85,7 @@ namespace Swift {
             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;
+            virtual JID messageCorrectionJID(const JID& fromJID) SWIFTEN_OVERRIDE;
 
         private:
             void setAvailableRoomActions(const MUCOccupant::Affiliation& affiliation, const MUCOccupant::Role& role);
diff --git a/Swift/Controllers/Chat/UnitTest/ChatsManagerTest.cpp b/Swift/Controllers/Chat/UnitTest/ChatsManagerTest.cpp
index d104fbd..0356c6a 100644
--- a/Swift/Controllers/Chat/UnitTest/ChatsManagerTest.cpp
+++ b/Swift/Controllers/Chat/UnitTest/ChatsManagerTest.cpp
@@ -137,6 +137,11 @@ class ChatsManagerTest : public CppUnit::TestFixture {
     CPPUNIT_TEST(testCarbonsForwardedOutgoingMessageFromSecondResource);
     CPPUNIT_TEST(testCarbonsForwardedIncomingDuplicates);
 
+
+    // Message correction tests
+    CPPUNIT_TEST(testChatControllerMessageCorrectionReplaceBySameResource);
+    CPPUNIT_TEST(testChatControllerMessageCorrectionReplaceByOtherResource);
+
     CPPUNIT_TEST_SUITE_END();
 
 public:
@@ -217,7 +222,6 @@ public:
         delete chatListWindow_;
         delete mocks_;
         delete settings_;
-
     }
 
     void testFirstOpenWindowIncoming() {
@@ -732,7 +736,7 @@ public:
         presence->setShow(StatusShow::None);
         presence->setType(Presence::Unavailable);
         stanzaChannel_->onPresenceReceived(presence);
-        CPPUNIT_ASSERT_EQUAL(std::string("participantA has gone offline."), MockChatWindow::bodyFromMessage(window->lastReplacedMessage_));
+        CPPUNIT_ASSERT_EQUAL(std::string("participantA has gone offline."), MockChatWindow::bodyFromMessage(window->lastReplacedLastMessage_));
     }
 
     void testChatControllerMucPmUnavailableErrorHandling() {
@@ -1097,7 +1101,7 @@ public:
         CPPUNIT_ASSERT_EQUAL(JID("room@muc.service.com"), window->lastMUCInvitationJID_);
 
         stanzaChannel_->onPresenceReceived(generateIncomingPresence(Presence::Unavailable));
-        CPPUNIT_ASSERT_EQUAL(std::string(""), MockChatWindow::bodyFromMessage(window->lastReplacedMessage_));
+        CPPUNIT_ASSERT_EQUAL(std::string(""), MockChatWindow::bodyFromMessage(window->lastReplacedLastMessage_));
         CPPUNIT_ASSERT_EQUAL(std::string("testling@test.com has gone offline."), MockChatWindow::bodyFromMessage(window->lastAddedPresence_));
     }
 
@@ -1134,7 +1138,7 @@ public:
             stanzaChannel_->onPresenceReceived(presence);
         }
         CPPUNIT_ASSERT_EQUAL(std::string("someDifferentNickname has entered the room."), window->bodyFromMessage(window->lastAddedPresence_));
-        CPPUNIT_ASSERT_EQUAL(std::string(), window->bodyFromMessage(window->lastReplacedMessage_));
+        CPPUNIT_ASSERT_EQUAL(std::string(), window->bodyFromMessage(window->lastReplacedLastMessage_));
         window->resetLastMessages();
 
         {
@@ -1145,7 +1149,7 @@ public:
             stanzaChannel_->onPresenceReceived(presence);
         }
         CPPUNIT_ASSERT_EQUAL(std::string(), window->bodyFromMessage(window->lastAddedPresence_));
-        CPPUNIT_ASSERT_EQUAL(std::string("someDifferentNickname and Romeo have entered the room"), window->bodyFromMessage(window->lastReplacedMessage_));
+        CPPUNIT_ASSERT_EQUAL(std::string("someDifferentNickname and Romeo have entered the room"), window->bodyFromMessage(window->lastReplacedLastMessage_));
         window->resetLastMessages();
 
         {
@@ -1165,7 +1169,7 @@ public:
             stanzaChannel_->onPresenceReceived(presence);
         }
         CPPUNIT_ASSERT_EQUAL(std::string(), window->bodyFromMessage(window->lastAddedPresence_));
-        CPPUNIT_ASSERT_EQUAL(std::string("someDifferentNickname, Romeo and Juliet have entered the room"), window->bodyFromMessage(window->lastReplacedMessage_));
+        CPPUNIT_ASSERT_EQUAL(std::string("someDifferentNickname, Romeo and Juliet have entered the room"), window->bodyFromMessage(window->lastReplacedLastMessage_));
     }
 
     template <typename CarbonsType>
@@ -1265,21 +1269,21 @@ public:
             CPPUNIT_ASSERT_EQUAL(ChatWindow::ReceiptReceived, window->receiptChanges_[1].second);
         }
     }
-
+    
     void testCarbonsForwardedIncomingDuplicates() {
         JID messageJID("testling@test.com/resource1");
         JID jid2 = jid_.toBare().withResource("someOtherResource");
-
+        
         MockChatWindow* window = new MockChatWindow();
         mocks_->ExpectCall(chatWindowFactory_, ChatWindowFactory::createChatWindow).With(messageJID, uiEventStream_).Return(window);
-
+        
         std::shared_ptr<Message> message(new Message());
         message->setFrom(messageJID);
         std::string body("This is a legible message. >HEH@)oeueu");
         message->setBody(body);
         manager_->handleIncomingMessage(message);
         CPPUNIT_ASSERT_EQUAL(body, MockChatWindow::bodyFromMessage(window->lastAddedMessage_));
-
+        
         // incoming carbons message from another resource and duplicate of it
         {
             auto originalMessage = std::make_shared<Message>();
@@ -1289,15 +1293,15 @@ public:
             originalMessage->setType(Message::Chat);
             std::string forwardedBody = "Some further text.";
             originalMessage->setBody(forwardedBody);
-
+            
             auto messageWrapper = createCarbonsMessage(std::make_shared<CarbonsReceived>(), originalMessage);
-
+            
             manager_->handleIncomingMessage(messageWrapper);
-
+            
             CPPUNIT_ASSERT_EQUAL(forwardedBody, MockChatWindow::bodyFromMessage(window->lastAddedMessage_));
             CPPUNIT_ASSERT_EQUAL(false, window->lastAddedMessageSenderIsSelf_);
             window->resetLastMessages();
-
+            
             messageWrapper = createCarbonsMessage(std::make_shared<CarbonsReceived>(), originalMessage);
             manager_->handleIncomingMessage(messageWrapper);
             CPPUNIT_ASSERT_EQUAL(std::string(), MockChatWindow::bodyFromMessage(window->lastAddedMessage_));
@@ -1305,6 +1309,58 @@ public:
         }
     }
 
+    void testChatControllerMessageCorrectionReplaceBySameResource() {
+        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");
+        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 after edit"), MockChatWindow::bodyFromMessage(window->lastReplacedMessage_));
+    }
+
+    void testChatControllerMessageCorrectionReplaceByOtherResource() {
+        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");
+        manager_->handleIncomingMessage(message);
+
+        CPPUNIT_ASSERT_EQUAL(std::string("text before edit"), MockChatWindow::bodyFromMessage(window->lastAddedMessage_));
+
+        message = std::make_shared<Message>();
+        message->setFrom(messageJID.toBare().withResource("resource2"));
+        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 after edit"), MockChatWindow::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>();