1 files changed, 20 insertions, 9 deletions
diff --git a/Swiften/Network/Connector.cpp b/Swiften/Network/Connector.cpp
index 5e7f8d9..1db1eac 100644
--- a/Swiften/Network/Connector.cpp
+++ b/Swiften/Network/Connector.cpp
@@ -18,5 +18,5 @@
 namespace Swift {
 
-Connector::Connector(const std::string& hostname, DomainNameResolver* resolver, ConnectionFactory* connectionFactory, TimerFactory* timerFactory, int defaultPort) : hostname(hostname), resolver(resolver), connectionFactory(connectionFactory), timerFactory(timerFactory), defaultPort(defaultPort), timeoutMilliseconds(0), queriedAllServices(true), foundSomeDNS(false) {
+Connector::Connector(const std::string& hostname, int port, const boost::optional<std::string>& serviceLookupPrefix, DomainNameResolver* resolver, ConnectionFactory* connectionFactory, TimerFactory* timerFactory) : hostname(hostname), port(port), serviceLookupPrefix(serviceLookupPrefix), resolver(resolver), connectionFactory(connectionFactory), timerFactory(timerFactory), timeoutMilliseconds(0), queriedAllServices(true), foundSomeDNS(false) {
 }
 
@@ -32,13 +32,17 @@ void Connector::start() {
 	assert(!timer);
 	queriedAllServices = false;
-	serviceQuery = resolver->createServiceQuery("_xmpp-client._tcp." + hostname);
-	serviceQuery->onResult.connect(boost::bind(&Connector::handleServiceQueryResult, shared_from_this(), _1));
 	if (timeoutMilliseconds > 0) {
 		timer = timerFactory->createTimer(timeoutMilliseconds);
 		timer->onTick.connect(boost::bind(&Connector::handleTimeout, shared_from_this()));
-		timer->start();
 	}
+	if (serviceLookupPrefix) {
+		serviceQuery = resolver->createServiceQuery(*serviceLookupPrefix, hostname);
+		serviceQuery->onResult.connect(boost::bind(&Connector::handleServiceQueryResult, shared_from_this(), _1));
 		serviceQuery->run();
 	}
+	else {
+		queryAddress(hostname);
+	}
+}
 
 void Connector::stop() {
@@ -110,10 +114,10 @@ void Connector::tryNextAddress() {
 		addressQueryResults.pop_front();
 
-		int port = defaultPort;
+		int connectPort = (port == -1 ? 5222 : port);
 		if (!serviceQueryResults.empty()) {
-			port = serviceQueryResults.front().port;
+			connectPort = serviceQueryResults.front().port;
 		}
 
-		tryConnect(HostAddressPort(address, port));
+		tryConnect(HostAddressPort(address, connectPort));
 	}
 }
@@ -125,8 +129,15 @@ void Connector::tryConnect(const HostAddressPort& target) {
 	currentConnection->onConnectFinished.connect(boost::bind(&Connector::handleConnectionConnectFinished, shared_from_this(), _1));
 	currentConnection->connect(target);
+	if (timer) {
+		timer->start();
+	}
 }
 
 void Connector::handleConnectionConnectFinished(bool error) {
 	SWIFT_LOG(debug) << "ConnectFinished: " << (error ? "error" : "success") << std::endl;
+	if (timer) {
+			timer->stop();
+			timer.reset();
+	}
 	currentConnection->onConnectFinished.disconnect(boost::bind(&Connector::handleConnectionConnectFinished, shared_from_this(), _1));
 	if (error) {
@@ -170,6 +181,6 @@ void Connector::finish(boost::shared_ptr<Connection> connection) {
 void Connector::handleTimeout() {
 	SWIFT_LOG(debug) << "Timeout" << std::endl;
-	finish(boost::shared_ptr<Connection>());
+	handleConnectionConnectFinished(true);
 }
 
-};
+}