6 files changed, 94 insertions, 17 deletions
diff --git a/Swiften/LinkLocal/DNSSD/Fake/FakeDNSSDQuerier.cpp b/Swiften/LinkLocal/DNSSD/Fake/FakeDNSSDQuerier.cpp
index 0ba4dbb..f7dfa11 100644
--- a/Swiften/LinkLocal/DNSSD/Fake/FakeDNSSDQuerier.cpp
+++ b/Swiften/LinkLocal/DNSSD/Fake/FakeDNSSDQuerier.cpp
@@ -48,20 +48,13 @@ void FakeDNSSDQuerier::addRunningQuery(boost::shared_ptr<FakeDNSSDQuery> query)
 		MainEventLoop::postEvent(boost::bind(boost::ref(registerQuery->onRegisterFinished), service), shared_from_this());
 	}
 	else if (boost::shared_ptr<FakeDNSSDResolveHostnameQuery> resolveHostnameQuery = boost::dynamic_pointer_cast<FakeDNSSDResolveHostnameQuery>(query)) {
-		std::map<String,HostAddress>::const_iterator i = addresses.find(resolveHostnameQuery->hostname);
+		std::map<String,boost::optional<HostAddress> >::const_iterator i = addresses.find(resolveHostnameQuery->hostname);
 		if (i != addresses.end()) {
 			MainEventLoop::postEvent(
 					boost::bind(
 						boost::ref(resolveHostnameQuery->onHostnameResolved), i->second), 
 					shared_from_this());
 		}
-		else {
-			MainEventLoop::postEvent(
-					boost::bind(
-						boost::ref(resolveHostnameQuery->onHostnameResolved), 
-						boost::optional<HostAddress>()), 
-					shared_from_this());
-		}
 	}
 }
 
@@ -118,8 +111,14 @@ void FakeDNSSDQuerier::setRegisterError() {
 	}
 }
 
-void FakeDNSSDQuerier::setAddress(const String& hostname, const HostAddress& address) {
+void FakeDNSSDQuerier::setAddress(const String& hostname, boost::optional<HostAddress> address) {
 	addresses[hostname] = address;
+	foreach(const boost::shared_ptr<FakeDNSSDResolveHostnameQuery>& query, getQueries<FakeDNSSDResolveHostnameQuery>()) {
+		if (query->hostname == hostname) {
+			MainEventLoop::postEvent(boost::bind(
+					boost::ref(query->onHostnameResolved), address), shared_from_this());
+		}
+	}
 }
 
 }
diff --git a/Swiften/LinkLocal/DNSSD/Fake/FakeDNSSDQuerier.h b/Swiften/LinkLocal/DNSSD/Fake/FakeDNSSDQuerier.h
index 94d4074..5af49dc 100644
--- a/Swiften/LinkLocal/DNSSD/Fake/FakeDNSSDQuerier.h
+++ b/Swiften/LinkLocal/DNSSD/Fake/FakeDNSSDQuerier.h
@@ -39,7 +39,7 @@ namespace Swift {
 			void removeService(const DNSSDServiceID& id);
 			void setServiceInfo(const DNSSDServiceID& id, const DNSSDResolveServiceQuery::Result& info);
 			bool isServiceRegistered(const String& name, int port, const ByteArray& info);
-			void setAddress(const String& hostname, const HostAddress& address);
+			void setAddress(const String& hostname, boost::optional<HostAddress> address);
 
 			void setBrowseError();
 			void setRegisterError();
@@ -62,6 +62,6 @@ namespace Swift {
 			std::set<DNSSDServiceID> services;
 			typedef std::map<DNSSDServiceID,DNSSDResolveServiceQuery::Result> ServiceInfoMap;
 			ServiceInfoMap serviceInfo;
-			std::map<String, HostAddress> addresses;
+			std::map<String, boost::optional<HostAddress> > addresses;
 	};
 }
diff --git a/Swiften/LinkLocal/LinkLocalConnector.cpp b/Swiften/LinkLocal/LinkLocalConnector.cpp
index 40558d4..326d053 100644
--- a/Swiften/LinkLocal/LinkLocalConnector.cpp
+++ b/Swiften/LinkLocal/LinkLocalConnector.cpp
@@ -20,6 +20,10 @@ LinkLocalConnector::LinkLocalConnector(
 			connection(connection) {
 }
 
+LinkLocalConnector::~LinkLocalConnector() {
+	assert(!resolveQuery);
+}
+
 void LinkLocalConnector::connect() {
 	resolveQuery = querier->createResolveHostnameQuery(
 			service.getHostname(), 
@@ -31,6 +35,14 @@ void LinkLocalConnector::connect() {
 	resolveQuery->run();
 }
 
+void LinkLocalConnector::cancel() {
+	if (resolveQuery) {
+		resolveQuery->finish();
+	}
+	resolveQuery.reset();
+	connection->disconnect();
+}
+
 void LinkLocalConnector::handleHostnameResolved(const boost::optional<HostAddress>& address) {
 	if (address) {
 		resolveQuery->finish();
diff --git a/Swiften/LinkLocal/LinkLocalConnector.h b/Swiften/LinkLocal/LinkLocalConnector.h
index fa293c4..0b6baef 100644
--- a/Swiften/LinkLocal/LinkLocalConnector.h
+++ b/Swiften/LinkLocal/LinkLocalConnector.h
@@ -23,12 +23,14 @@ namespace Swift {
 					const LinkLocalService& service, 
 					boost::shared_ptr<DNSSDQuerier> querier,
 					boost::shared_ptr<Connection> connection);
+			~LinkLocalConnector();
 
 			const LinkLocalService& getService() const {
 				return service;
 			}
 
 			void connect();
+			void cancel();
 			void queueElement(boost::shared_ptr<Element> element);
 
 			const std::vector<boost::shared_ptr<Element> >& getQueuedElements() const {
diff --git a/Swiften/LinkLocal/UnitTest/LinkLocalConnectorTest.cpp b/Swiften/LinkLocal/UnitTest/LinkLocalConnectorTest.cpp
index b052e6a..ee5e414 100644
--- a/Swiften/LinkLocal/UnitTest/LinkLocalConnectorTest.cpp
+++ b/Swiften/LinkLocal/UnitTest/LinkLocalConnectorTest.cpp
@@ -16,6 +16,8 @@ class LinkLocalConnectorTest : public CppUnit::TestFixture {
 		CPPUNIT_TEST(testConnect);
 		CPPUNIT_TEST(testConnect_UnableToResolve);
 		CPPUNIT_TEST(testConnect_UnableToConnect);
+		CPPUNIT_TEST(testCancel_DuringResolve);
+		CPPUNIT_TEST(testCancel_DuringConnect);
 		CPPUNIT_TEST_SUITE_END();
 
 	public:
@@ -49,6 +51,7 @@ class LinkLocalConnectorTest : public CppUnit::TestFixture {
 		void testConnect_UnableToResolve() {
 			boost::shared_ptr<LinkLocalConnector> 
 					testling(createConnector("rabbithole.local", 1234));
+			querier->setAddress("rabbithole.local", boost::optional<HostAddress>());
 
 			testling->connect();
 			eventLoop->processEvents();
@@ -71,6 +74,36 @@ class LinkLocalConnectorTest : public CppUnit::TestFixture {
 			CPPUNIT_ASSERT(connectError);
 			CPPUNIT_ASSERT(!connection->connectedTo);
 		}
+
+		void testCancel_DuringResolve() {
+			boost::shared_ptr<LinkLocalConnector> 
+					testling(createConnector("rabbithole.local", 1234));
+			testling->connect();
+			eventLoop->processEvents();
+			CPPUNIT_ASSERT(!connectFinished);
+
+			testling->cancel();
+			eventLoop->processEvents();
+			querier->setAddress("rabbithole.local", HostAddress("192.168.1.1"));
+			eventLoop->processEvents();
+
+			CPPUNIT_ASSERT(FakeConnection::Disconnected == connection->state);
+		}
+
+		void testCancel_DuringConnect() {
+			boost::shared_ptr<LinkLocalConnector> 
+					testling(createConnector("rabbithole.local", 1234));
+			querier->setAddress("rabbithole.local", HostAddress("192.168.1.1"));
+			connection->setDelayConnect();
+			testling->connect();
+			eventLoop->processEvents();
+			CPPUNIT_ASSERT(FakeConnection::Connecting == connection->state);
+
+			testling->cancel();
+			eventLoop->processEvents();
+
+			CPPUNIT_ASSERT(FakeConnection::Disconnected == connection->state);
+		}
 	
 	private:
 		boost::shared_ptr<LinkLocalConnector> createConnector(const String& hostname, int port) {
diff --git a/Swiften/Network/FakeConnection.h b/Swiften/Network/FakeConnection.h
index 076a1f1..92a03c3 100644
--- a/Swiften/Network/FakeConnection.h
+++ b/Swiften/Network/FakeConnection.h
@@ -16,7 +16,15 @@ namespace Swift {
 			public EventOwner,
 			public boost::enable_shared_from_this<FakeConnection> {
 		public:
-			FakeConnection() {}
+			enum State {
+				Initial,
+				Connecting,
+				Connected,
+				Disconnected,
+				DisconnectedWithError
+			};
+
+			FakeConnection() : state(Initial), delayConnect(false) {}
 
 			virtual void listen() {
 				assert(false);
@@ -24,6 +32,7 @@ namespace Swift {
 
 			void setError(const Error& e) {
 				error = boost::optional<Error>(e);
+				state = DisconnectedWithError;
 				if (connectedTo) {
 					MainEventLoop::postEvent(
 							boost::bind(boost::ref(onDisconnected), error),
@@ -32,15 +41,31 @@ namespace Swift {
 			}
 
 			virtual void connect(const HostAddressPort& address) {
-				if (!error) {
-					connectedTo = address;
+				if (delayConnect) {
+					state = Connecting;
+				}
+				else {
+					if (!error) {
+						connectedTo = address;
+						state = Connected;
+					}
+					else {
+						state = DisconnectedWithError;
+					}
+					MainEventLoop::postEvent(
+							boost::bind(boost::ref(onConnectFinished), error),
+							shared_from_this());
 				}
-				MainEventLoop::postEvent(
-						boost::bind(boost::ref(onConnectFinished), error),
-						shared_from_this());
 			}
 
 			virtual void disconnect() {
+				if (!error) {
+					state = Disconnected;
+				}
+				else {
+					state = DisconnectedWithError;
+				}
+				connectedTo.reset();
 				MainEventLoop::postEvent(
 						boost::bind(boost::ref(onDisconnected), error),
 						shared_from_this());
@@ -50,8 +75,14 @@ namespace Swift {
 				dataWritten.push_back(data);
 			}
 
+			void setDelayConnect() {
+				delayConnect = true;
+			}
+
 			boost::optional<HostAddressPort> connectedTo;
 			std::vector<ByteArray> dataWritten;
 			boost::optional<Error> error;
+			State state;
+			bool delayConnect;
 	};
 }