4 files changed, 47 insertions, 11 deletions
diff --git a/Swift/Controllers/Chat/ChatController.cpp b/Swift/Controllers/Chat/ChatController.cpp
index 911f3e2..9767f8d 100644
--- a/Swift/Controllers/Chat/ChatController.cpp
+++ b/Swift/Controllers/Chat/ChatController.cpp
@@ -102,8 +102,23 @@ void ChatController::preHandleIncomingMessage(boost::shared_ptr<MessageEvent> me
 			setToJID(from);
 		}
 	}
-	chatStateNotifier_->receivedMessageFromContact(message->getPayload<ChatState>());
+	boost::shared_ptr<Replace> replace = message->getPayload<Replace>();
+	if (replace) {
+		// Determine the timestamp
+		boost::posix_time::ptime timeStamp = boost::posix_time::microsec_clock::universal_time();
+		boost::optional<boost::posix_time::ptime> messageTimeStamp = getMessageTimestamp(message);
+		if (messageTimeStamp) {
+			timeStamp = *messageTimeStamp;
+		}
+		std::string body = message->getBody();
+		chatWindow_->replaceMessage(body, lastMessageUIID_, timeStamp);
+		replacedMessage_ = true;
+	}
+	else {
+		replacedMessage_ = false;
+	}
 	chatStateTracker_->handleMessageReceived(message);
+	chatStateNotifier_->receivedMessageFromContact(message->getPayload<ChatState>());
 }
 
 void ChatController::postHandleIncomingMessage(boost::shared_ptr<MessageEvent> messageEvent) {
@@ -116,10 +131,19 @@ void ChatController::preSendMessageRequest(boost::shared_ptr<Message> message) {
 }
 
 void ChatController::postSendMessage(const std::string& body, boost::shared_ptr<Stanza> sentStanza) {
-	std::string id = addMessage(body, QT_TRANSLATE_NOOP("", "me"), true, labelsEnabled_ ? chatWindow_->getSelectedSecurityLabel().getLabel() : boost::shared_ptr<SecurityLabel>(), std::string(avatarManager_->getAvatarPath(selfJID_).string()), boost::posix_time::microsec_clock::universal_time());
-	if (stanzaChannel_->getStreamManagementEnabled() && !id.empty()) {
-		chatWindow_->setAckState(id, ChatWindow::Pending);
-		unackedStanzas_[sentStanza] = id;
+	if (stanzaChannel_->getStreamManagementEnabled() && !myLastMessageUIID_.empty()) {
+		chatWindow_->setAckState(myLastMessageUIID_, ChatWindow::Pending);
+		unackedStanzas_[sentStanza] = myLastMessageUIID_;
+	}
+	boost::shared_ptr<Replace> replace = sentStanza->getPayload<Replace>();
+	if (replace) {
+		chatWindow_->replaceMessage(body, myLastMessageUIID_, boost::posix_time::microsec_clock::universal_time());
+	} else {
+		myLastMessageUIID_ = addMessage(body, QT_TRANSLATE_NOOP("", "me"), true, labelsEnabled_ ? chatWindow_->getSelectedSecurityLabel().getLabel() : boost::shared_ptr<SecurityLabel>(), std::string(avatarManager_->getAvatarPath(selfJID_).string()), boost::posix_time::microsec_clock::universal_time());
+		if (stanzaChannel_->getStreamManagementEnabled()) {
+			chatWindow_->setAckState(myLastMessageUIID_, ChatWindow::Pending);
+			unackedStanzas_[sentStanza] = myLastMessageUIID_;
+		}
 	}
 	lastWasPresence_ = false;
 	chatStateNotifier_->userSentMessage();
diff --git a/Swift/Controllers/Chat/ChatController.h b/Swift/Controllers/Chat/ChatController.h
index dd4bf90..4fafb44 100644
--- a/Swift/Controllers/Chat/ChatController.h
+++ b/Swift/Controllers/Chat/ChatController.h
@@ -40,6 +40,7 @@ namespace Swift {
 			NickResolver* nickResolver_;
 			ChatStateNotifier* chatStateNotifier_;
 			ChatStateTracker* chatStateTracker_;
+			std::string myLastMessageUIID_;
 			bool isInMUC_;
 			bool lastWasPresence_;
 			std::string lastStatusChangeString_;
diff --git a/Swift/Controllers/Chat/ChatControllerBase.cpp b/Swift/Controllers/Chat/ChatControllerBase.cpp
index 566d075..14e17cd 100644
--- a/Swift/Controllers/Chat/ChatControllerBase.cpp
+++ b/Swift/Controllers/Chat/ChatControllerBase.cpp
@@ -30,9 +30,10 @@ namespace Swift {
 ChatControllerBase::ChatControllerBase(const JID& self, StanzaChannel* stanzaChannel, IQRouter* iqRouter, ChatWindowFactory* chatWindowFactory, const JID &toJID, PresenceOracle* presenceOracle, AvatarManager* avatarManager, bool useDelayForLatency, UIEventStream* eventStream, EventController* eventController, TimerFactory* timerFactory) : selfJID_(self), stanzaChannel_(stanzaChannel), iqRouter_(iqRouter), chatWindowFactory_(chatWindowFactory), toJID_(toJID), labelsEnabled_(false), presenceOracle_(presenceOracle), avatarManager_(avatarManager), useDelayForLatency_(useDelayForLatency), eventController_(eventController), timerFactory_(timerFactory) {
 	chatWindow_ = chatWindowFactory_->createChatWindow(toJID, eventStream);
 	chatWindow_->onAllMessagesRead.connect(boost::bind(&ChatControllerBase::handleAllMessagesRead, this));
-	chatWindow_->onSendMessageRequest.connect(boost::bind(&ChatControllerBase::handleSendMessageRequest, this, _1));
+	chatWindow_->onSendMessageRequest.connect(boost::bind(&ChatControllerBase::handleSendMessageRequest, this, _1, _2));
 	setOnline(stanzaChannel->isAvailable() && iqRouter->isAvailable());
 	createDayChangeTimer();
+	replacedMessage_ = false;
 }
 
 ChatControllerBase::~ChatControllerBase() {
@@ -88,7 +89,7 @@ void ChatControllerBase::handleAllMessagesRead() {
 	chatWindow_->setUnreadMessageCount(0);
 }
 
-void ChatControllerBase::handleSendMessageRequest(const std::string &body) {
+void ChatControllerBase::handleSendMessageRequest(const std::string &body, bool isCorrectionMessage) {
 	if (!stanzaChannel_->isAvailable() || body.empty()) {
 		return;
 	}
@@ -104,6 +105,10 @@ void ChatControllerBase::handleSendMessageRequest(const std::string &body) {
 		boost::posix_time::ptime now = boost::posix_time::microsec_clock::universal_time();
 		message->addPayload(boost::shared_ptr<Delay>(new Delay(now, selfJID_)));
 	}
+	if (isCorrectionMessage) {
+		message->addPayload(boost::shared_ptr<Replace> (new Replace(lastSentMessageStanzaID_)));
+	}
+	message->setID(lastSentMessageStanzaID_ = idGenerator_.generateID());
 	stanzaChannel_->sendMessage(message);
 	postSendMessage(message->getBody(), boost::dynamic_pointer_cast<Stanza>(message));
 	onActivity(message->getBody());
@@ -164,7 +169,7 @@ void ChatControllerBase::handleIncomingMessage(boost::shared_ptr<MessageEvent> m
 		JID from = message->getFrom();
 		std::vector<boost::shared_ptr<Delay> > delayPayloads = message->getPayloads<Delay>();
 		for (size_t i = 0; useDelayForLatency_ && i < delayPayloads.size(); i++) {
-			if (!delayPayloads[i]->getFrom()) { 
+			if (!delayPayloads[i]->getFrom()) {
 				continue;
 			}
 			boost::posix_time::ptime now = boost::posix_time::microsec_clock::universal_time();
@@ -180,9 +185,10 @@ void ChatControllerBase::handleIncomingMessage(boost::shared_ptr<MessageEvent> m
 		if (messageTimeStamp) {
 			timeStamp = *messageTimeStamp;
 		}
-
-		addMessage(body, senderDisplayNameFromMessage(from), isIncomingMessageFromMe(message), label, std::string(avatarManager_->getAvatarPath(from).string()), timeStamp);
 		onActivity(body);
+		if (!replacedMessage_) {
+			lastMessageUIID_ = addMessage(body, senderDisplayNameFromMessage(from), isIncomingMessageFromMe(message), label, std::string(avatarManager_->getAvatarPath(from).string()), timeStamp);
+		}
 	}
 	chatWindow_->show();
 	chatWindow_->setUnreadMessageCount(unreadMessages_.size());
diff --git a/Swift/Controllers/Chat/ChatControllerBase.h b/Swift/Controllers/Chat/ChatControllerBase.h
index 4898320..e0f1b94 100644
--- a/Swift/Controllers/Chat/ChatControllerBase.h
+++ b/Swift/Controllers/Chat/ChatControllerBase.h
@@ -26,6 +26,7 @@
 #include "Swiften/Elements/ErrorPayload.h"
 #include "Swiften/Presence/PresenceOracle.h"
 #include "Swiften/Queries/IQRouter.h"
+#include "Swiften/Base/IDGenerator.h"
 
 namespace Swift {
 	class IQRouter;
@@ -66,8 +67,10 @@ namespace Swift {
 			virtual void dayTicked() {};
 
 		private:
+			IDGenerator idGenerator_;
+			std::string lastSentMessageStanzaID_;
 			void createDayChangeTimer();
-			void handleSendMessageRequest(const std::string &body);
+			void handleSendMessageRequest(const std::string &body, bool isCorrectionMessage);
 			void handleAllMessagesRead();
 			void handleSecurityLabelsCatalogResponse(boost::shared_ptr<SecurityLabelsCatalog>, ErrorPayload::ref error);
 			std::string getErrorMessage(boost::shared_ptr<ErrorPayload>);
@@ -82,6 +85,8 @@ namespace Swift {
 			ChatWindow* chatWindow_;
 			JID toJID_;
 			bool labelsEnabled_;
+			bool replacedMessage_;
+			std::string lastMessageUIID_;
 			PresenceOracle* presenceOracle_;
 			AvatarManager* avatarManager_;
 			bool useDelayForLatency_;