#include "Swiften/Network/Connector.h"

#include <boost/bind.hpp>

#include <iostream>

#include "Swiften/Network/ConnectionFactory.h"

#include "Swiften/Network/DomainNameResolver.h"

#include "Swiften/Network/DomainNameAddressQuery.h"

namespace Swift {

Connector::Connector(const String& hostname, DomainNameResolver* resolver, ConnectionFactory* connectionFactory) : hostname(hostname), resolver(resolver), connectionFactory(connectionFactory), queriedAllHosts(true) {

}

void Connector::start() {

assert(!currentConnection);

assert(!serviceQuery);

queriedAllHosts = false;

serviceQuery = resolver->createServiceQuery("_xmpp-client._tcp." + hostname);

serviceQuery->onResult.connect(boost::bind(&Connector::handleServiceQueryResult, this, _1));

serviceQuery->run();

}

void Connector::queryAddress(const String& hostname) {

assert(!addressQuery);

addressQuery = resolver->createAddressQuery(hostname);

addressQuery->onResult.connect(boost::bind(&Connector::handleAddressQueryResult, this, _1, _2));

addressQuery->run();

}

void Connector::handleServiceQueryResult(const std::vector<DomainNameServiceQuery::Result>& result) {

serviceQueryResults = std::deque<DomainNameServiceQuery::Result>(result.begin(), result.end());

serviceQuery.reset();

tryNextHostname();

}

void Connector::tryNextHostname() {

if (queriedAllHosts) {

onConnectFinished(boost::shared_ptr<Connection>());

}

else if (serviceQueryResults.empty()) {

// Fall back on simple address resolving

queriedAllHosts = true;

queryAddress(hostname);

}

else {

queryAddress(serviceQueryResults.front().hostname);

}

}

void Connector::handleAddressQueryResult(const HostAddress& address, boost::optional<DomainNameResolveError> error) {

addressQuery.reset();

if (!serviceQueryResults.empty()) {

DomainNameServiceQuery::Result serviceQueryResult = serviceQueryResults.front();

serviceQueryResults.pop_front();

if (error) {

tryNextHostname();

}

else {

tryConnect(HostAddressPort(address, serviceQueryResult.port));

}

}

else if (error) {

// The fallback address query failed

assert(queriedAllHosts);

onConnectFinished(boost::shared_ptr<Connection>());

}

else {

// The fallback query succeeded

tryConnect(HostAddressPort(address, 5222));

}

}

void Connector::tryConnect(const HostAddressPort& target) {

assert(!currentConnection);

currentConnection = connectionFactory->createConnection();

currentConnection->onConnectFinished.connect(boost::bind(&Connector::handleConnectionConnectFinished, this, _1));

currentConnection->connect(target);

}

void Connector::handleConnectionConnectFinished(bool error) {

if (error) {

currentConnection.reset();

tryNextHostname();

}

else {

onConnectFinished(currentConnection);

}

}

};

#include "Swiften/Network/PlatformDomainNameResolver.h"

// Putting this early on, because some system types conflict with thread

#include "Swiften/Network/PlatformDomainNameServiceQuery.h"

#include <string>

#include <vector>

#include <boost/asio.hpp>

#include <boost/bind.hpp>

#include <boost/thread.hpp>

#include <boost/enable_shared_from_this.hpp>

#include <algorithm>

#include "Swiften/Base/String.h"

#include "Swiften/Network/HostAddress.h"

#include "Swiften/EventLoop/MainEventLoop.h"

#include "Swiften/Network/HostAddressPort.h"

#include "Swiften/Network/DomainNameAddressQuery.h"

using namespace Swift;

namespace {

struct AddressQuery : public DomainNameAddressQuery, public boost::enable_shared_from_this<AddressQuery>, public EventOwner {

AddressQuery(const String& host) : hostname(host), thread(NULL), safeToJoin(false) {}

~AddressQuery() {

if (safeToJoin) {

thread->join();

}

else {

// FIXME: UGLYYYYY

}

delete thread;

}

void run() {

safeToJoin = false;

thread = new boost::thread(boost::bind(&AddressQuery::doRun, shared_from_this()));

}

void doRun() {

boost::asio::ip::tcp::resolver resolver(ioService);

boost::asio::ip::tcp::resolver::query query(hostname.getUTF8String(), "5222");

try {

boost::asio::ip::tcp::resolver::iterator endpointIterator = resolver.resolve(query);

if (endpointIterator == boost::asio::ip::tcp::resolver::iterator()) {

emitError();

}

else {

boost::asio::ip::address address = (*endpointIterator).endpoint().address();

HostAddress result = (address.is_v4() ? HostAddress(&address.to_v4().to_bytes()[0], 4) : HostAddress(&address.to_v6().to_bytes()[0], 16));

MainEventLoop::postEvent(

boost::bind(boost::ref(onResult), result, boost::optional<DomainNameResolveError>()),

shared_from_this());

}

}

catch (...) {

emitError();

}

safeToJoin = true;

}

void emitError() {

MainEventLoop::postEvent(boost::bind(boost::ref(onResult), HostAddress(), boost::optional<DomainNameResolveError>(DomainNameResolveError())), shared_from_this());

}

boost::asio::io_service ioService;

String hostname;

boost::thread* thread;

bool safeToJoin;

};

}

namespace Swift {

PlatformDomainNameResolver::PlatformDomainNameResolver() {

}

boost::shared_ptr<DomainNameServiceQuery> PlatformDomainNameResolver::createServiceQuery(const String& name) {

return boost::shared_ptr<DomainNameServiceQuery>(new PlatformDomainNameServiceQuery(getNormalized(name)));

}

boost::shared_ptr<DomainNameAddressQuery> PlatformDomainNameResolver::createAddressQuery(const String& name) {

return boost::shared_ptr<DomainNameAddressQuery>(new AddressQuery(getNormalized(name)));

}

}

if (safeToJoin) {

thread->join();

}

else {

// FIXME: UGLYYYYY

}

delete thread;

}

void PlatformDomainNameServiceQuery::run() {

safeToJoin = false;

thread = new boost::thread(boost::bind(&PlatformDomainNameServiceQuery::doRun, shared_from_this()));

}

void PlatformDomainNameServiceQuery::doRun() {

std::vector<DomainNameServiceQuery::Result> records;

#if defined(SWIFTEN_PLATFORM_WINDOWS)

DNS_RECORD* responses;

// FIXME: This conversion doesn't work if unicode is deffed above

if (DnsQuery(service.getUTF8Data(), DNS_TYPE_SRV, DNS_QUERY_STANDARD, NULL, &responses, NULL) != ERROR_SUCCESS) {

emitError();

return;

}

DNS_RECORD* currentEntry = responses;

while (currentEntry) {

if (currentEntry->wType == DNS_TYPE_SRV) {

DomainNameServiceQuery::Result record;

record.priority = currentEntry->Data.SRV.wPriority;

record.weight = currentEntry->Data.SRV.wWeight;

record.port = currentEntry->Data.SRV.wPort;

// The pNameTarget is actually a PCWSTR, so I would have expected this

// conversion to not work at all, but it does.

// Actually, it doesn't. Fix this and remove explicit cast

// Remove unicode undef above as well

record.hostname = String((const char*) currentEntry->Data.SRV.pNameTarget);

records.push_back(record);

}

currentEntry = currentEntry->pNext;

}

DnsRecordListFree(responses, DnsFreeRecordList);

#else

// Make sure we reinitialize the domain list every time

res_init();

ByteArray response;

response.resize(NS_PACKETSZ);

int responseLength = res_query(const_cast<char*>(service.getUTF8Data()), ns_c_in, ns_t_srv, reinterpret_cast<u_char*>(response.getData()), response.getSize());

if (responseLength == -1) {

emitError();

return;

}

// Parse header

HEADER* header = reinterpret_cast<HEADER*>(response.getData());

unsigned char* messageStart = reinterpret_cast<unsigned char*>(response.getData());

unsigned char* messageEnd = messageStart + responseLength;

unsigned char* currentEntry = messageStart + NS_HFIXEDSZ;

// Skip over the queries

int queriesCount = ntohs(header->qdcount);

while (queriesCount > 0) {

int entryLength = dn_skipname(currentEntry, messageEnd);

if (entryLength < 0) {

emitError();

return;

}

currentEntry += entryLength + NS_QFIXEDSZ;

queriesCount--;

}

// Process the SRV answers

int answersCount = ntohs(header->ancount);

while (answersCount > 0) {

DomainNameServiceQuery::Result record;

int entryLength = dn_skipname(currentEntry, messageEnd);

currentEntry += entryLength;

currentEntry += NS_RRFIXEDSZ;

// Priority

if (currentEntry + 2 >= messageEnd) {

emitError();

return;

}

record.priority = ns_get16(currentEntry);

currentEntry += 2;

// Weight

if (currentEntry + 2 >= messageEnd) {

emitError();

return;

}

record.weight = ns_get16(currentEntry);

currentEntry += 2;

// Port

if (currentEntry + 2 >= messageEnd) {

emitError();

return;

}

record.port = ns_get16(currentEntry);

currentEntry += 2;

// Hostname

if (currentEntry >= messageEnd) {

emitError();

return;

}

ByteArray entry;

entry.resize(NS_MAXDNAME);

entryLength = dn_expand(messageStart, messageEnd, currentEntry, entry.getData(), entry.getSize());

if (entryLength < 0) {

emitError();

return;

}

record.hostname = String(entry.getData());

records.push_back(record);

currentEntry += entryLength;

answersCount--;

}

#endif

safeToJoin = true;

std::sort(records.begin(), records.end(), ResultPriorityComparator());

MainEventLoop::postEvent(boost::bind(boost::ref(onResult), records));

}

void PlatformDomainNameServiceQuery::emitError() {

safeToJoin = true;

MainEventLoop::postEvent(boost::bind(boost::ref(onResult), std::vector<DomainNameServiceQuery::Result>()), shared_from_this());

}

}