1 files changed, 42 insertions, 20 deletions
diff --git a/Swift/Controllers/Chat/ChatController.cpp b/Swift/Controllers/Chat/ChatController.cpp
index 22ef68d..f4aa745 100644
--- a/Swift/Controllers/Chat/ChatController.cpp
+++ b/Swift/Controllers/Chat/ChatController.cpp
@@ -11,16 +11,17 @@
 
 #include <Swift/Controllers/Intl.h>
 #include <Swiften/Base/format.h>
-#include "Swiften/Avatars/AvatarManager.h"
-#include "Swiften/Chat/ChatStateNotifier.h"
-#include "Swiften/Chat/ChatStateTracker.h"
-#include "Swiften/Client/StanzaChannel.h"
-#include "Swiften/Disco/EntityCapsProvider.h"
-#include "Swift/Controllers/UIInterfaces/ChatWindow.h"
-#include "Swift/Controllers/UIInterfaces/ChatWindowFactory.h"
-#include "Swiften/Client/NickResolver.h"
-#include "Swift/Controllers/XMPPEvents/EventController.h"
+#include <Swiften/Base/Algorithm.h>
+#include <Swiften/Avatars/AvatarManager.h>
+#include <Swiften/Chat/ChatStateNotifier.h>
+#include <Swiften/Chat/ChatStateTracker.h>
+#include <Swiften/Client/StanzaChannel.h>
+#include <Swift/Controllers/UIInterfaces/ChatWindow.h>
+#include <Swift/Controllers/UIInterfaces/ChatWindowFactory.h>
+#include <Swiften/Client/NickResolver.h>
+#include <Swift/Controllers/XMPPEvents/EventController.h>
 #include <Swift/Controllers/StatusUtil.h>
+#include <Swiften/Disco/EntityCapsProvider.h>
 
 namespace Swift {
 	
@@ -28,7 +29,7 @@ namespace Swift {
  * The controller does not gain ownership of the stanzaChannel, nor the factory.
  */
 ChatController::ChatController(const JID& self, StanzaChannel* stanzaChannel, IQRouter* iqRouter, ChatWindowFactory* chatWindowFactory, const JID &contact, NickResolver* nickResolver, PresenceOracle* presenceOracle, AvatarManager* avatarManager, bool isInMUC, bool useDelayForLatency, UIEventStream* eventStream, EventController* eventController, TimerFactory* timerFactory, EntityCapsProvider* entityCapsProvider)
-	: ChatControllerBase(self, stanzaChannel, iqRouter, chatWindowFactory, contact, presenceOracle, avatarManager, useDelayForLatency, eventStream, eventController, timerFactory) {
+	: ChatControllerBase(self, stanzaChannel, iqRouter, chatWindowFactory, contact, presenceOracle, avatarManager, useDelayForLatency, eventStream, eventController, timerFactory, entityCapsProvider) {
 	isInMUC_ = isInMUC;
 	lastWasPresence_ = false;
 	chatStateNotifier_ = new ChatStateNotifier(stanzaChannel, contact, entityCapsProvider);
@@ -59,6 +60,7 @@ ChatController::ChatController(const JID& self, StanzaChannel* stanzaChannel, IQ
 	chatWindow_->addSystemMessage(startMessage);
 	chatWindow_->onUserTyping.connect(boost::bind(&ChatStateNotifier::setUserIsTyping, chatStateNotifier_));
 	chatWindow_->onUserCancelsTyping.connect(boost::bind(&ChatStateNotifier::userCancelledNewMessage, chatStateNotifier_));
+	handleBareJIDCapsChanged(toJID_);
 
 }
 
@@ -74,6 +76,19 @@ ChatController::~ChatController() {
 	delete chatStateTracker_;
 }
 
+void ChatController::handleBareJIDCapsChanged(const JID& /*jid*/) {
+	DiscoInfo::ref disco = entityCapsProvider_->getCaps(toJID_);
+	if (disco) {
+		if (disco->hasFeature(DiscoInfo::MessageCorrectionFeature)) {
+			chatWindow_->setCorrectionEnabled(ChatWindow::Yes);
+		} else {
+			chatWindow_->setCorrectionEnabled(ChatWindow::No);
+		}
+	} else {
+		chatWindow_->setCorrectionEnabled(ChatWindow::Maybe);
+	}
+}
+
 void ChatController::setToJID(const JID& jid) {
 	chatStateNotifier_->setContact(jid);
 	ChatControllerBase::setToJID(jid);
@@ -84,6 +99,7 @@ void ChatController::setToJID(const JID& jid) {
 		presence = jid.isBare() ? presenceOracle_->getHighestPriorityPresence(jid.toBare()) : presenceOracle_->getLastPresence(jid);
 	}
 	chatStateNotifier_->setContactIsOnline(presence && presence->getType() == Presence::Available);
+	handleBareJIDCapsChanged(toJID_);
 }
 
 bool ChatController::isIncomingMessageFromMe(boost::shared_ptr<Message>) {
@@ -102,8 +118,8 @@ void ChatController::preHandleIncomingMessage(boost::shared_ptr<MessageEvent> me
 			setToJID(from);
 		}
 	}
-	chatStateNotifier_->receivedMessageFromContact(message->getPayload<ChatState>());
 	chatStateTracker_->handleMessageReceived(message);
+	chatStateNotifier_->receivedMessageFromContact(message->getPayload<ChatState>());
 }
 
 void ChatController::postHandleIncomingMessage(boost::shared_ptr<MessageEvent> messageEvent) {
@@ -116,27 +132,33 @@ 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()) {
-		chatWindow_->setAckState(id, ChatWindow::Pending);
-		unackedStanzas_[sentStanza] = id;
+	boost::shared_ptr<Replace> replace = sentStanza->getPayload<Replace>();
+	if (replace) {
+		eraseIf(unackedStanzas_, PairSecondEquals<boost::shared_ptr<Stanza>, std::string>(myLastMessageUIID_));
+		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() && !myLastMessageUIID_.empty() ) {
+		chatWindow_->setAckState(myLastMessageUIID_, ChatWindow::Pending);
+		unackedStanzas_[sentStanza] = myLastMessageUIID_;
 	}
 	lastWasPresence_ = false;
 	chatStateNotifier_->userSentMessage();
 }
 
 void ChatController::handleStanzaAcked(boost::shared_ptr<Stanza> stanza) {
-	std::string id = unackedStanzas_[stanza];
-	if (id != "") {
-		chatWindow_->setAckState(id, ChatWindow::Received);
+	std::map<boost::shared_ptr<Stanza>, std::string>::iterator unackedStanza = unackedStanzas_.find(stanza);
+	if (unackedStanza != unackedStanzas_.end()) {
+		chatWindow_->setAckState(unackedStanza->second, ChatWindow::Received);
+		unackedStanzas_.erase(unackedStanza);
 	}
-	unackedStanzas_.erase(unackedStanzas_.find(stanza));
 }
 
 void ChatController::setOnline(bool online) {
 	if (!online) {
 		std::map<boost::shared_ptr<Stanza>, std::string>::iterator it = unackedStanzas_.begin();
-		for ( ; it != unackedStanzas_.end(); it++) {
+		for ( ; it != unackedStanzas_.end(); ++it) {
 			chatWindow_->setAckState(it->second, ChatWindow::Failed);
 		}
 		unackedStanzas_.clear();