4 files changed, 147 insertions, 54 deletions
diff --git a/Swift/Controllers/Chat/UserSearchController.cpp b/Swift/Controllers/Chat/UserSearchController.cpp
index 886aa47..0539f9b 100644
--- a/Swift/Controllers/Chat/UserSearchController.cpp
+++ b/Swift/Controllers/Chat/UserSearchController.cpp
@@ -4,7 +4,7 @@
  * See Documentation/Licenses/GPLv3.txt for more information.
  */
 
-#include "Swift/Controllers/Chat/UserSearchController.h"
+#include <Swift/Controllers/Chat/UserSearchController.h>
 
 #include <boost/bind.hpp>
 #include <boost/shared_ptr.hpp>
@@ -20,28 +20,33 @@
 #include <Swift/Controllers/UIInterfaces/UserSearchWindowFactory.h>
 
 namespace Swift {
-UserSearchController::UserSearchController(Type type, const JID& jid, UIEventStream* uiEventStream, UserSearchWindowFactory* factory, IQRouter* iqRouter) : type_(type), jid_(jid) {
-	iqRouter_ = iqRouter;
-	uiEventStream_ = uiEventStream;
-	uiEventConnection_ = uiEventStream_->onUIEvent.connect(boost::bind(&UserSearchController::handleUIEvent, this, _1));
+UserSearchController::UserSearchController(Type type, const JID& jid, UIEventStream* uiEventStream, UserSearchWindowFactory* factory, IQRouter* iqRouter) : type_(type), jid_(jid), uiEventStream_(uiEventStream), factory_(factory), iqRouter_(iqRouter) {
+	uiEventStream_->onUIEvent.connect(boost::bind(&UserSearchController::handleUIEvent, this, _1));
 	window_ = NULL;
-	factory_ = factory;
 	discoWalker_ = NULL;
 }
 
 UserSearchController::~UserSearchController() {
-	delete window_;
+	endDiscoWalker();
 	delete discoWalker_;
+	if (window_) {
+		window_->onFormRequested.disconnect(boost::bind(&UserSearchController::handleFormRequested, this, _1));
+		window_->onSearchRequested.disconnect(boost::bind(&UserSearchController::handleSearch, this, _1, _2));
+		delete window_;
+	}
+	uiEventStream_->onUIEvent.disconnect(boost::bind(&UserSearchController::handleUIEvent, this, _1));
 }
 
 void UserSearchController::handleUIEvent(boost::shared_ptr<UIEvent> event) {
 	bool handle = false;
 	if (type_ == AddContact) {
-		boost::shared_ptr<RequestAddUserDialogUIEvent> searchEvent = boost::dynamic_pointer_cast<RequestAddUserDialogUIEvent>(event);
-		if (searchEvent) handle = true;
+		if (boost::dynamic_pointer_cast<RequestAddUserDialogUIEvent>(event)) {
+			handle = true;
+		}
 	} else {
-		boost::shared_ptr<RequestChatWithUserDialogUIEvent> searchEvent = boost::dynamic_pointer_cast<RequestChatWithUserDialogUIEvent>(event);
-		if (searchEvent) handle = true;
+		if (boost::dynamic_pointer_cast<RequestChatWithUserDialogUIEvent>(event)) {
+			handle = true;
+		}
 	}
 	if (handle) {
 		if (!window_) {
@@ -59,15 +64,28 @@ void UserSearchController::handleUIEvent(boost::shared_ptr<UIEvent> event) {
 void UserSearchController::handleFormRequested(const JID& service) {
 	window_->setSearchError(false);
 	window_->setServerSupportsSearch(true);
+
 	//Abort a previous search if is active
+	endDiscoWalker();
 	delete discoWalker_;
 	discoWalker_ = new DiscoServiceWalker(service, iqRouter_);
-	discoWalker_->onServiceFound.connect(boost::bind(&UserSearchController::handleDiscoServiceFound, this, _1, _2, discoWalker_));
-	discoWalker_->onWalkComplete.connect(boost::bind(&UserSearchController::handleDiscoWalkFinished, this, discoWalker_));
+	discoWalker_->onServiceFound.connect(boost::bind(&UserSearchController::handleDiscoServiceFound, this, _1, _2));
+	discoWalker_->onWalkComplete.connect(boost::bind(&UserSearchController::handleDiscoWalkFinished, this));
 	discoWalker_->beginWalk();
 }
 
-void UserSearchController::handleDiscoServiceFound(const JID& jid, boost::shared_ptr<DiscoInfo> info, DiscoServiceWalker* walker) {
+void UserSearchController::endDiscoWalker() {
+	if (discoWalker_) {
+		discoWalker_->endWalk();
+		discoWalker_->onServiceFound.disconnect(boost::bind(&UserSearchController::handleDiscoServiceFound, this, _1, _2));
+		discoWalker_->onWalkComplete.disconnect(boost::bind(&UserSearchController::handleDiscoWalkFinished, this));
+		delete discoWalker_;
+		discoWalker_ = NULL;
+	}
+}
+
+
+void UserSearchController::handleDiscoServiceFound(const JID& jid, boost::shared_ptr<DiscoInfo> info) {
 	bool isUserDirectory = false;
 	bool supports55 = false;
 	foreach (DiscoInfo::Identity identity, info->getIdentities()) {
@@ -80,15 +98,14 @@ void UserSearchController::handleDiscoServiceFound(const JID& jid, boost::shared
 	supports55 = std::find(features.begin(), features.end(), DiscoInfo::JabberSearchFeature) != features.end();
 	if (/*isUserDirectory && */supports55) { //FIXME: once M-Link correctly advertises directoryness.
 		/* Abort further searches.*/
-		delete discoWalker_;
-		discoWalker_ = NULL;
+		endDiscoWalker();
 		boost::shared_ptr<GenericRequest<SearchPayload> > searchRequest(new GenericRequest<SearchPayload>(IQ::Get, jid, boost::shared_ptr<SearchPayload>(new SearchPayload()), iqRouter_));
-		searchRequest->onResponse.connect(boost::bind(&UserSearchController::handleFormResponse, this, _1, _2, jid));
+		searchRequest->onResponse.connect(boost::bind(&UserSearchController::handleFormResponse, this, _1, _2));
 		searchRequest->send();
 	}
 }
 
-void UserSearchController::handleFormResponse(boost::shared_ptr<SearchPayload> fields, ErrorPayload::ref error, const JID& jid) {
+void UserSearchController::handleFormResponse(boost::shared_ptr<SearchPayload> fields, ErrorPayload::ref error) {
 	if (error || !fields) {
 		window_->setServerSupportsSearch(false);
 		return;
@@ -98,11 +115,11 @@ void UserSearchController::handleFormResponse(boost::shared_ptr<SearchPayload> f
 
 void UserSearchController::handleSearch(boost::shared_ptr<SearchPayload> fields, const JID& jid) {
 	boost::shared_ptr<GenericRequest<SearchPayload> > searchRequest(new GenericRequest<SearchPayload>(IQ::Set, jid, fields, iqRouter_));
-	searchRequest->onResponse.connect(boost::bind(&UserSearchController::handleSearchResponse, this, _1, _2, jid));
+	searchRequest->onResponse.connect(boost::bind(&UserSearchController::handleSearchResponse, this, _1, _2));
 	searchRequest->send();
 }
 
-void UserSearchController::handleSearchResponse(boost::shared_ptr<SearchPayload> resultsPayload, ErrorPayload::ref error, const JID& jid) {
+void UserSearchController::handleSearchResponse(boost::shared_ptr<SearchPayload> resultsPayload, ErrorPayload::ref error) {
 	if (error || !resultsPayload) {
 		window_->setSearchError(true);
 		return;
@@ -121,10 +138,9 @@ void UserSearchController::handleSearchResponse(boost::shared_ptr<SearchPayload>
 	window_->setResults(results);
 }
 
-void UserSearchController::handleDiscoWalkFinished(DiscoServiceWalker* walker) {
+void UserSearchController::handleDiscoWalkFinished() {
 	window_->setServerSupportsSearch(false);
-	delete discoWalker_;
-	discoWalker_ = NULL;
+	endDiscoWalker();
 }
 
 }
diff --git a/Swift/Controllers/Chat/UserSearchController.h b/Swift/Controllers/Chat/UserSearchController.h
index c54b8d5..9b81020 100644
--- a/Swift/Controllers/Chat/UserSearchController.h
+++ b/Swift/Controllers/Chat/UserSearchController.h
@@ -41,21 +41,24 @@ namespace Swift {
 			enum Type {AddContact, StartChat};
 			UserSearchController(Type type, const JID& jid, UIEventStream* uiEventStream, UserSearchWindowFactory* userSearchWindowFactory, IQRouter* iqRouter);
 			~UserSearchController();
+
 		private:
 			void handleUIEvent(boost::shared_ptr<UIEvent> event);
 			void handleFormRequested(const JID& service);
-			void handleDiscoServiceFound(const JID& jid, boost::shared_ptr<DiscoInfo> info, DiscoServiceWalker* walker);
-			void handleDiscoWalkFinished(DiscoServiceWalker* walker);
-			void handleFormResponse(boost::shared_ptr<SearchPayload> items, ErrorPayload::ref error, const JID& jid);
+			void handleDiscoServiceFound(const JID& jid, boost::shared_ptr<DiscoInfo> info);
+			void handleDiscoWalkFinished();
+			void handleFormResponse(boost::shared_ptr<SearchPayload> items, ErrorPayload::ref error);
 			void handleSearch(boost::shared_ptr<SearchPayload> fields, const JID& jid);
-			void handleSearchResponse(boost::shared_ptr<SearchPayload> results, ErrorPayload::ref error, const JID& jid);
+			void handleSearchResponse(boost::shared_ptr<SearchPayload> results, ErrorPayload::ref error);
+			void endDiscoWalker();
+
+		private:
 			Type type_;
+			JID jid_;
 			UIEventStream* uiEventStream_;
-			UserSearchWindow* window_;
 			UserSearchWindowFactory* factory_;
-			boost::bsignals::scoped_connection uiEventConnection_;
 			IQRouter* iqRouter_;
-			JID jid_;
+			UserSearchWindow* window_;
 			DiscoServiceWalker* discoWalker_;
 	};
 }
diff --git a/Swift/Controllers/DiscoServiceWalker.cpp b/Swift/Controllers/DiscoServiceWalker.cpp
index 95ce23b..505acb4 100644
--- a/Swift/Controllers/DiscoServiceWalker.cpp
+++ b/Swift/Controllers/DiscoServiceWalker.cpp
@@ -5,32 +5,55 @@
  */
 
 #include <Swift/Controllers/DiscoServiceWalker.h>
+#include <Swiften/Base/Log.h>
 
 #include <boost/bind.hpp>
 
-#include "Swiften/Disco/GetDiscoInfoRequest.h"
-#include "Swiften/Disco/GetDiscoItemsRequest.h"
-
 namespace Swift {
 
-DiscoServiceWalker::DiscoServiceWalker(const JID& service, IQRouter* iqRouter, size_t maxSteps) : service_(service), iqRouter_(iqRouter), maxSteps_(maxSteps) {
+DiscoServiceWalker::DiscoServiceWalker(const JID& service, IQRouter* iqRouter, size_t maxSteps) : service_(service), iqRouter_(iqRouter), maxSteps_(maxSteps), active_(false) {
 
 }
 
 void DiscoServiceWalker::beginWalk() {
-	assert(servicesBeingSearched_.size() == 0);
+	SWIFT_LOG(debug) << "Starting walk to " << service_ << std::endl;
+	assert(!active_);
+	assert(servicesBeingSearched_.empty());
+	active_ = true;
 	walkNode(service_);
 }
 
+void DiscoServiceWalker::endWalk() {
+	if (active_) {
+		SWIFT_LOG(debug) << "Ending walk to " << service_ << std::endl;
+		foreach (GetDiscoInfoRequest::ref request, pendingDiscoInfoRequests_) {
+			request->onResponse.disconnect(boost::bind(&DiscoServiceWalker::handleDiscoInfoResponse, this, _1, _2, request));
+		}
+		foreach (GetDiscoItemsRequest::ref request, pendingDiscoItemsRequests_) {
+			request->onResponse.disconnect(boost::bind(&DiscoServiceWalker::handleDiscoItemsResponse, this, _1, _2, request));
+		}
+		active_ = false;
+	}
+}
+
 void DiscoServiceWalker::walkNode(const JID& jid) {
-	servicesBeingSearched_.push_back(jid);
-	searchedServices_.push_back(jid);
+	SWIFT_LOG(debug) << "Walking node " << jid << std::endl;
+	servicesBeingSearched_.insert(jid);
+	searchedServices_.insert(jid);
 	GetDiscoInfoRequest::ref discoInfoRequest = GetDiscoInfoRequest::create(jid, iqRouter_);
-	discoInfoRequest->onResponse.connect(boost::bind(&DiscoServiceWalker::handleDiscoInfoResponse, this, _1, _2, jid));
+	discoInfoRequest->onResponse.connect(boost::bind(&DiscoServiceWalker::handleDiscoInfoResponse, this, _1, _2, discoInfoRequest));
+	pendingDiscoInfoRequests_.insert(discoInfoRequest);
 	discoInfoRequest->send();
 }
 
 void DiscoServiceWalker::handleReceivedDiscoItem(const JID& item) {
+	SWIFT_LOG(debug) << "Received disco item " << item << std::endl;
+
+	/* If we got canceled, don't do anything */
+	if (!active_) {
+		return;
+	}
+
 	if (std::find(searchedServices_.begin(), searchedServices_.end(), item) != searchedServices_.end()) {
 		/* Don't recurse infinitely */
 		return;
@@ -38,9 +61,17 @@ void DiscoServiceWalker::handleReceivedDiscoItem(const JID& item) {
 	walkNode(item);
 }
 
-void DiscoServiceWalker::handleDiscoInfoResponse(boost::shared_ptr<DiscoInfo> info, ErrorPayload::ref error, const JID& jid) {
+void DiscoServiceWalker::handleDiscoInfoResponse(boost::shared_ptr<DiscoInfo> info, ErrorPayload::ref error, GetDiscoInfoRequest::ref request) {
+	/* If we got canceled, don't do anything */
+	if (!active_) {
+		return;
+	}
+
+	SWIFT_LOG(debug) << "Disco info response from " << request->getReceiver() << std::endl;
+
+	pendingDiscoInfoRequests_.erase(request);
 	if (error) {
-		handleDiscoError(jid, error);
+		handleDiscoError(request->getReceiver(), error);
 		return;
 	}
 
@@ -52,21 +83,30 @@ void DiscoServiceWalker::handleDiscoInfoResponse(boost::shared_ptr<DiscoInfo> in
 	}
 	bool completed = false;
 	if (couldContainServices) {
-		GetDiscoItemsRequest::ref discoItemsRequest = GetDiscoItemsRequest::create(jid, iqRouter_);
-		discoItemsRequest->onResponse.connect(boost::bind(&DiscoServiceWalker::handleDiscoItemsResponse, this, _1, _2, jid));
+		GetDiscoItemsRequest::ref discoItemsRequest = GetDiscoItemsRequest::create(request->getReceiver(), iqRouter_);
+		discoItemsRequest->onResponse.connect(boost::bind(&DiscoServiceWalker::handleDiscoItemsResponse, this, _1, _2, discoItemsRequest));
+		pendingDiscoItemsRequests_.insert(discoItemsRequest);
 		discoItemsRequest->send();
 	} else {
 		completed = true;
 	}
-	onServiceFound(jid, info);
+	onServiceFound(request->getReceiver(), info);
 	if (completed) {
-		markNodeCompleted(jid);
+		markNodeCompleted(request->getReceiver());
 	}
 }
 
-void DiscoServiceWalker::handleDiscoItemsResponse(boost::shared_ptr<DiscoItems> items, ErrorPayload::ref error, const JID& jid) {
+void DiscoServiceWalker::handleDiscoItemsResponse(boost::shared_ptr<DiscoItems> items, ErrorPayload::ref error, GetDiscoItemsRequest::ref request) {
+	/* If we got canceled, don't do anything */
+	if (!active_) {
+		return;
+	}
+
+	SWIFT_LOG(debug) << "Received disco item from " << request->getReceiver() << std::endl;
+
+	pendingDiscoItemsRequests_.erase(request);
 	if (error) {
-		handleDiscoError(jid, error);
+		handleDiscoError(request->getReceiver(), error);
 		return;
 	}
 	foreach (DiscoItems::Item item, items->getItems()) {
@@ -77,15 +117,28 @@ void DiscoServiceWalker::handleDiscoItemsResponse(boost::shared_ptr<DiscoItems>
 			handleReceivedDiscoItem(item.getJID());
 		}
 	}
-	markNodeCompleted(jid);
+	markNodeCompleted(request->getReceiver());
 }
 
 void DiscoServiceWalker::handleDiscoError(const JID& jid, ErrorPayload::ref /*error*/) {
+	/* If we got canceled, don't do anything */
+	if (!active_) {
+		return;
+	}
+
+	SWIFT_LOG(debug) << "Disco error from " << jid << std::endl;
+
 	markNodeCompleted(jid);
 }
 
 void DiscoServiceWalker::markNodeCompleted(const JID& jid) {
-	servicesBeingSearched_.erase(std::remove(servicesBeingSearched_.begin(), servicesBeingSearched_.end(), jid), servicesBeingSearched_.end());
+	// Check whether we weren't canceled in between a 'emit result' and this call
+	if (!active_) {
+		return;
+	}
+	SWIFT_LOG(debug) << "Node completed " << jid << std::endl;
+
+	servicesBeingSearched_.erase(jid);
 	/* All results are in */
 	if (servicesBeingSearched_.size() == 0) {
 		onWalkComplete();
diff --git a/Swift/Controllers/DiscoServiceWalker.h b/Swift/Controllers/DiscoServiceWalker.h
index 5b2a47e..167174a 100644
--- a/Swift/Controllers/DiscoServiceWalker.h
+++ b/Swift/Controllers/DiscoServiceWalker.h
@@ -7,6 +7,7 @@
 #pragma once
 
 #include <vector>
+#include <set>
 
 #include <boost/shared_ptr.hpp>
 #include <Swiften/Base/boost_bsignals.h>
@@ -15,6 +16,8 @@
 #include <Swiften/Elements/DiscoInfo.h>
 #include <Swiften/Elements/DiscoItems.h>
 #include <Swiften/Elements/ErrorPayload.h>
+#include <Swiften/Disco/GetDiscoInfoRequest.h>
+#include <Swiften/Disco/GetDiscoItemsRequest.h>
 
 namespace Swift {
 	class IQRouter;
@@ -22,26 +25,44 @@ namespace Swift {
 	 * Recursively walk service discovery trees to find all services offered.
 	 * This stops on any disco item that's not reporting itself as a server.
 	 */
-	class DiscoServiceWalker : public boost::signals::trackable {
+	class DiscoServiceWalker {
 		public:
 			DiscoServiceWalker(const JID& service, IQRouter* iqRouter, size_t maxSteps = 200);
-			/** Start the walk. Call this exactly once.*/
+
+			/**
+			 * Start the walk.
+			 *
+			 * Call this exactly once.
+			 */
 			void beginWalk();
+
+			/**
+			 * End the walk.
+			 */
+			void endWalk();
+
 			/** Emitted for each service found. */
 			boost::signal<void(const JID&, boost::shared_ptr<DiscoInfo>)> onServiceFound;
+
 			/** Emitted when walking is complete.*/
 			boost::signal<void()> onWalkComplete;
+
 		private:
 			void handleReceivedDiscoItem(const JID& item);
 			void walkNode(const JID& jid);
 			void markNodeCompleted(const JID& jid);
-			void handleDiscoInfoResponse(boost::shared_ptr<DiscoInfo> info, ErrorPayload::ref error, const JID& jid);
-			void handleDiscoItemsResponse(boost::shared_ptr<DiscoItems> items, ErrorPayload::ref error, const JID& jid);
+			void handleDiscoInfoResponse(boost::shared_ptr<DiscoInfo> info, ErrorPayload::ref error, GetDiscoInfoRequest::ref request);
+			void handleDiscoItemsResponse(boost::shared_ptr<DiscoItems> items, ErrorPayload::ref error, GetDiscoItemsRequest::ref request);
 			void handleDiscoError(const JID& jid, ErrorPayload::ref error);
+
+		private:
 			JID service_;
 			IQRouter* iqRouter_;
 			size_t maxSteps_;
-			std::vector<JID> servicesBeingSearched_;
-			std::vector<JID> searchedServices_;
+			bool active_;
+			std::set<JID> servicesBeingSearched_;
+			std::set<JID> searchedServices_;
+			std::set<GetDiscoInfoRequest::ref> pendingDiscoInfoRequests_;
+			std::set<GetDiscoItemsRequest::ref> pendingDiscoItemsRequests_;
 	};
 }