3 files changed, 44 insertions, 1 deletions
diff --git a/Swift/Controllers/Chat/ChatControllerBase.cpp b/Swift/Controllers/Chat/ChatControllerBase.cpp
index bc9c14e..2bfff4f 100644
--- a/Swift/Controllers/Chat/ChatControllerBase.cpp
+++ b/Swift/Controllers/Chat/ChatControllerBase.cpp
@@ -124,7 +124,10 @@ void ChatControllerBase::handleSendMessageRequest(const std::string &body, bool
 	message->setType(Swift::Message::Chat);
 	message->setBody(body);
 	if (labelsEnabled_) {
-		SecurityLabelsCatalog::Item labelItem = chatWindow_->getSelectedSecurityLabel();
+		if (!isCorrectionMessage) {
+			lastLabel_ = chatWindow_->getSelectedSecurityLabel();
+		}
+		SecurityLabelsCatalog::Item labelItem = lastLabel_;
 		if (labelItem.getLabel()) {
 			message->addPayload(labelItem.getLabel());
 		}
diff --git a/Swift/Controllers/Chat/ChatControllerBase.h b/Swift/Controllers/Chat/ChatControllerBase.h
index b698581..8aed069 100644
--- a/Swift/Controllers/Chat/ChatControllerBase.h
+++ b/Swift/Controllers/Chat/ChatControllerBase.h
@@ -110,5 +110,6 @@ namespace Swift {
 			boost::shared_ptr<Timer> dateChangeTimer_;
 			TimerFactory* timerFactory_;
 			EntityCapsProvider* entityCapsProvider_;
+			SecurityLabelsCatalog::Item lastLabel_; 
 	};
 }
diff --git a/Swift/Controllers/Chat/UnitTest/MUCControllerTest.cpp b/Swift/Controllers/Chat/UnitTest/MUCControllerTest.cpp
index 76e716b..2cc62bc 100644
--- a/Swift/Controllers/Chat/UnitTest/MUCControllerTest.cpp
+++ b/Swift/Controllers/Chat/UnitTest/MUCControllerTest.cpp
@@ -39,6 +39,7 @@ class MUCControllerTest : public CppUnit::TestFixture {
 	CPPUNIT_TEST(testAddressedToSelfBySelf);
 	CPPUNIT_TEST(testMessageWithEmptyLabelItem);
 	CPPUNIT_TEST(testMessageWithLabelItem);
+	CPPUNIT_TEST(testCorrectMessageWithLabelItem);
 	CPPUNIT_TEST_SUITE_END();
 
 public:
@@ -210,6 +211,44 @@ public:
 		CPPUNIT_ASSERT_EQUAL(label, message->getPayload<SecurityLabel>());
 	}
 
+	void testCorrectMessageWithLabelItem() {
+		SecurityLabel::ref label = boost::make_shared<SecurityLabel>();
+		label->setLabel("a");
+		SecurityLabelsCatalog::Item labelItem;
+		labelItem.setSelector("Bob");
+		labelItem.setLabel(label);
+		SecurityLabel::ref label2 = boost::make_shared<SecurityLabel>();
+		label->setLabel("b");
+		SecurityLabelsCatalog::Item labelItem2;
+		labelItem2.setSelector("Charlie");
+		labelItem2.setLabel(label2);
+		window_->label_ = labelItem;
+		boost::shared_ptr<DiscoInfo> features = boost::make_shared<DiscoInfo>();
+		features->addFeature(DiscoInfo::SecurityLabelsCatalogFeature);
+		controller_->setAvailableServerFeatures(features);
+		IQ::ref iq = iqChannel_->iqs_[iqChannel_->iqs_.size() - 1];
+		SecurityLabelsCatalog::ref labelPayload = boost::make_shared<SecurityLabelsCatalog>();
+		labelPayload->addItem(labelItem);
+		IQ::ref result = IQ::createResult(self_, iq->getID(), labelPayload);
+		iqChannel_->onIQReceived(result);
+		std::string messageBody("agamemnon");
+		window_->onSendMessageRequest(messageBody, false);
+		boost::shared_ptr<Stanza> rawStanza = stanzaChannel_->sentStanzas[stanzaChannel_->sentStanzas.size() - 1];
+		Message::ref message = boost::dynamic_pointer_cast<Message>(rawStanza);
+		CPPUNIT_ASSERT_EQUAL(iq->getTo(), result->getFrom());
+		CPPUNIT_ASSERT(window_->labelsEnabled_);
+		CPPUNIT_ASSERT(stanzaChannel_->isAvailable()); /* Otherwise will prevent sends. */
+		CPPUNIT_ASSERT(message);
+		CPPUNIT_ASSERT_EQUAL(messageBody, message->getBody());
+		CPPUNIT_ASSERT_EQUAL(label, message->getPayload<SecurityLabel>());
+		window_->label_ = labelItem2;
+		window_->onSendMessageRequest(messageBody, true);
+		rawStanza = stanzaChannel_->sentStanzas[stanzaChannel_->sentStanzas.size() - 1];
+		message = boost::dynamic_pointer_cast<Message>(rawStanza);
+		CPPUNIT_ASSERT_EQUAL(messageBody, message->getBody());
+		CPPUNIT_ASSERT_EQUAL(label, message->getPayload<SecurityLabel>());
+	}
+
 	void checkEqual(const std::vector<NickJoinPart>& expected, const std::vector<NickJoinPart>& actual) {
 		CPPUNIT_ASSERT_EQUAL(expected.size(), actual.size());
 		for (size_t i = 0; i < expected.size(); i++) {