3 files changed, 62 insertions, 17 deletions
diff --git a/Swiften/Queries/IQRouter.cpp b/Swiften/Queries/IQRouter.cpp
index 671566a..b20d344 100644
--- a/Swiften/Queries/IQRouter.cpp
+++ b/Swiften/Queries/IQRouter.cpp
@@ -10,20 +10,18 @@
 
 namespace Swift {
 
-namespace {
-	void noop(IQHandler*) {}
-	struct PointerEquals {
-		PointerEquals(IQHandler* handler) : handler_(handler) {}
-		bool operator()(boost::shared_ptr<IQHandler> o) { return handler_ == o.get(); }
-		IQHandler* handler_;
-	};
-}
+static void noop(IQHandler*) {}
 
-IQRouter::IQRouter(IQChannel* channel) : channel_(channel) {
+IQRouter::IQRouter(IQChannel* channel) : channel_(channel), queueRemoves_(false) {
 	channel->onIQReceived.connect(boost::bind(&IQRouter::handleIQ, this, _1));
 }
 
+IQRouter::~IQRouter() {
+}
+
 void IQRouter::handleIQ(boost::shared_ptr<IQ> iq) {
+	queueRemoves_ = true;
+
 	bool handled = false;
 	foreach(boost::shared_ptr<IQHandler> handler, handlers_) {
 		handled |= handler->handleIQ(iq);
@@ -34,22 +32,38 @@ void IQRouter::handleIQ(boost::shared_ptr<IQ> iq) {
 	if (!handled && (iq->getType() == IQ::Get || iq->getType() == IQ::Set) ) {
 		channel_->sendIQ(IQ::createError(iq->getFrom(), iq->getID(), Error::FeatureNotImplemented, Error::Cancel));
 	}
+
+	processPendingRemoves();
+
+	queueRemoves_ = false;
 }
 
-void IQRouter::addHandler(IQHandler* handler) {
-	handlers_.push_back(boost::shared_ptr<IQHandler>(handler, noop));
+void IQRouter::processPendingRemoves() {
+	foreach(boost::shared_ptr<IQHandler> handler, queuedRemoves_) {
+		handlers_.erase(std::remove(handlers_.begin(), handlers_.end(), handler), handlers_.end());
+	}
+	queuedRemoves_.clear();
 }
 
-void IQRouter::addHandler(boost::shared_ptr<IQHandler> handler) {
-	handlers_.push_back(handler);
+void IQRouter::addHandler(IQHandler* handler) {
+	addHandler(boost::shared_ptr<IQHandler>(handler, noop));
 }
 
 void IQRouter::removeHandler(IQHandler* handler) {
-	handlers_.erase(std::remove_if(handlers_.begin(), handlers_.end(), PointerEquals(handler)), handlers_.end());
+	removeHandler(boost::shared_ptr<IQHandler>(handler, noop));
+}
+
+void IQRouter::addHandler(boost::shared_ptr<IQHandler> handler) {
+	handlers_.push_back(handler);
 }
 
 void IQRouter::removeHandler(boost::shared_ptr<IQHandler> handler) {
-	handlers_.erase(std::remove(handlers_.begin(), handlers_.end(), handler));
+	if (queueRemoves_) {
+		queuedRemoves_.push_back(handler);
+	}
+	else {
+		handlers_.erase(std::remove(handlers_.begin(), handlers_.end(), handler), handlers_.end());
+	}
 }
 
 void IQRouter::sendIQ(boost::shared_ptr<IQ> iq) {
diff --git a/Swiften/Queries/IQRouter.h b/Swiften/Queries/IQRouter.h
index ea80bf5..6de2ff9 100644
--- a/Swiften/Queries/IQRouter.h
+++ b/Swiften/Queries/IQRouter.h
@@ -14,10 +14,11 @@ namespace Swift {
 	class IQRouter {
 		public:
 			IQRouter(IQChannel* channel);
+			~IQRouter();
 
 			void addHandler(IQHandler* handler);
-			void addHandler(boost::shared_ptr<IQHandler> handler);
 			void removeHandler(IQHandler* handler);
+			void addHandler(boost::shared_ptr<IQHandler> handler);
 			void removeHandler(boost::shared_ptr<IQHandler> handler);
 
 			void sendIQ(boost::shared_ptr<IQ> iq);
@@ -25,10 +26,13 @@ namespace Swift {
 
 		private:
 			void handleIQ(boost::shared_ptr<IQ> iq);
+			void processPendingRemoves();
 
 		private:
 			IQChannel* channel_;
 			std::vector< boost::shared_ptr<IQHandler> > handlers_;
+			std::vector< boost::shared_ptr<IQHandler> > queuedRemoves_;
+			bool queueRemoves_;
 	};
 }
 
diff --git a/Swiften/Queries/UnitTest/IQRouterTest.cpp b/Swiften/Queries/UnitTest/IQRouterTest.cpp
index 8fa8d2a..94b7de8 100644
--- a/Swiften/Queries/UnitTest/IQRouterTest.cpp
+++ b/Swiften/Queries/UnitTest/IQRouterTest.cpp
@@ -12,6 +12,8 @@ using namespace Swift;
 class IQRouterTest : public CppUnit::TestFixture
 {
 		CPPUNIT_TEST_SUITE(IQRouterTest);
+		CPPUNIT_TEST(testRemoveHandler);
+		CPPUNIT_TEST(testRemoveHandler_AfterHandleIQ);
 		CPPUNIT_TEST(testHandleIQ_SuccesfulHandlerFirst);
 		CPPUNIT_TEST(testHandleIQ_SuccesfulHandlerLast);
 		CPPUNIT_TEST(testHandleIQ_NoSuccesfulHandler);
@@ -29,6 +31,31 @@ class IQRouterTest : public CppUnit::TestFixture
 			delete channel_;
 		}
 
+		void testRemoveHandler() {
+			IQRouter testling(channel_);
+			DummyIQHandler handler1(true, &testling);
+			DummyIQHandler handler2(true, &testling);
+			testling.removeHandler(&handler1);
+
+			channel_->onIQReceived(boost::shared_ptr<IQ>(new IQ()));
+
+			CPPUNIT_ASSERT_EQUAL(0, handler1.called);
+			CPPUNIT_ASSERT_EQUAL(1, handler2.called);
+		}
+
+		void testRemoveHandler_AfterHandleIQ() {
+			IQRouter testling(channel_);
+			DummyIQHandler handler1(true, &testling);
+			DummyIQHandler handler2(true, &testling);
+
+			channel_->onIQReceived(boost::shared_ptr<IQ>(new IQ()));
+			testling.removeHandler(&handler1);
+			channel_->onIQReceived(boost::shared_ptr<IQ>(new IQ()));
+
+			CPPUNIT_ASSERT_EQUAL(1, handler1.called);
+			CPPUNIT_ASSERT_EQUAL(1, handler2.called);
+		}
+
 		void testHandleIQ_SuccesfulHandlerFirst() {
 			IQRouter testling(channel_);
 			DummyIQHandler handler1(true, &testling);
@@ -75,7 +102,7 @@ class IQRouterTest : public CppUnit::TestFixture
 			CPPUNIT_ASSERT_EQUAL(1, handler1.called);
 			CPPUNIT_ASSERT_EQUAL(2, handler2.called);
 		}
-	
+
 	private:
 		struct DummyIQHandler : public IQHandler {
 			DummyIQHandler(bool handle, IQRouter* router) : handle(handle), router(router), called(0) {