/builddir/build/BUILD/libutil++-1.4.0/src/sella/net/DNS.cpp

/*
** libutil++
** $Id: DNS.cpp 1197 2014-10-14 22:26:11Z sella $
** Copyright (c) 2011-2014 Digital Genesis, LLC. All Rights Reserved.
** Released under the LGPL Version 2.1 License.
** http://www.digitalgenesis.com
*/

static const char rcsid[] __attribute__((used)) = "$Id: DNS.cpp 1197 2014-10-14 22:26:11Z sella $";

#include "DNS.h"
#include "DNSException.h"
#include "../util/ThreadPool.h"

#include <string.h>
#include <fcntl.h>
#include <net/if.h>
#include <arpa/inet.h>
#include <sys/ioctl.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netdb.h>

union __sockaddr_union {
        struct sockaddr_storage         *ss;
        struct sockaddr                         *sa;
        struct sockaddr_in                      *sa_in;
        struct sockaddr_in6                     *sa_in6;
};

using namespace sella::net;
using namespace sella::util;

const std::vector<std::string> DNS::empty;
bool walk(std::vector<IPAddress> &ips, struct addrinfo *result, union __sockaddr_union &u, int protocol);

DNS::DNS(const std::vector<std::string> &nameservers, const std::vector<std::string> &searches) throw (DNSException) :
        pool(NULL),
        context(NULL), 
        flags(0)
{
        init(nameservers, searches);
}

DNS::DNS(sella::util::ThreadPool *pool, const std::vector<std::string> &nameservers, const std::vector<std::string> &searches) throw (DNSException) :
        pool(pool),
        context(NULL), 
        flags(0)
{
        init(nameservers, searches);

        Task::callback function = std::bind(&DNS::receive, this, 0, 0);
        task = Task::shared_ptr();
        task->setCallback(function);
}

DNS::~DNS() throw () {
        if (pool && task) {
                pool->unschedule(task, true);
        }

        ::dns_free(this->context);
}

std::vector<IPAddress> DNS::resolve(const std::string &name, int protocol) throw (DNSException) {
        int s;
        union __sockaddr_union u;
        struct sockaddr_storage addr;
        struct addrinfo hints, *result;
        std::vector<IPAddress> ips;

        memset(&hints, 0, sizeof(struct addrinfo));
        hints.ai_family = AF_UNSPEC;
        hints.ai_socktype = SOCK_DGRAM; /* Datagram socket */
        hints.ai_flags = AI_PASSIVE;            /* For wildcard IP address */
        hints.ai_protocol = 0;                          /* Any protocol */
        hints.ai_canonname = NULL;
        hints.ai_addr = NULL;
        hints.ai_next = NULL;

        if (protocol == IPv4Required) {
                hints.ai_family = AF_INET;
        } else if (protocol == IPv6Required) {
                hints.ai_family = AF_INET6;
        }

        if ((s = getaddrinfo(name.c_str(), NULL, &hints, &result)) != 0) {
                THROW_CODE(DNSException, s, "getaddrinfo() failed: %s", gai_strerror(s));
        }

        u.ss = (struct sockaddr_storage*) &addr;
        u.sa_in = (struct sockaddr_in*) &addr;
        memset(&addr, 0, sizeof(addr));

        try {
                if (protocol == AnyProtocol) {
                        walk(ips, result, u, AF_INET);
                } else if (protocol == IPv4Preferred) {
                        if (!walk(ips, result, u, AF_INET)) {
                                if (!walk(ips, result, u, AF_INET6)) {
                                        THROW_CODE(DNSException, EAI_NODATA, "Failed to resolve %s to an IPv4 or IPv6 address", name.c_str());
                                }
                        }
                } else if (protocol == IPv4Required) {
                        if (!walk(ips, result, u, AF_INET)) {
                                THROW_CODE(DNSException, EAI_NODATA, "Failed to resolve %s to an IPv4 address", name.c_str());
                        }
                } else if (protocol == IPv6Preferred) {
                        if (!walk(ips, result, u, AF_INET6)) {
                                if (!walk(ips, result, u, AF_INET)) {
                                        THROW_CODE(DNSException, EAI_NODATA, "Failed to resolve %s to an IPv4 or IPv6 address", name.c_str());
                                }
                        }
                } else if (protocol == IPv6Required) {
                        if (!walk(ips, result, u, AF_INET6)) {
                                THROW_CODE(DNSException, EAI_NODATA, "Failed to resolve %s to an IPv6 address", name.c_str());
                        }
                }

                freeaddrinfo(result);
        } catch (DNSException &e) {
                freeaddrinfo(result);

                throw e;
        }

        return ips;
}

DNSCallback* DNS::resolve(DNSCallback* callback, const std::string &name, int protocol) throw (DNSException) {
        if (callback == NULL) {
                THROW(DNSException, "invalid NULL callback supplied");
        }

        callback->qname = name;
        callback->protocol = protocol;
        callback->flags = flags;

        switch (protocol) {
                case IPv4Preferred:
                case IPv4Required:
                        ::dns_submit_p(this->context, name.c_str(), DNS_C_IN, DNS_T_A, flags, ::dns_parse_a4, DNS::callback_a, callback);
                        break;

                case IPv6Preferred:
                case IPv6Required:
                        ::dns_submit_p(this->context, name.c_str(), DNS_C_IN, DNS_T_AAAA, flags, ::dns_parse_a6, DNS::callback_aaaa, callback);
                        break;

                default:
                        THROW(DNSException, "invalid protocol supplied (%d); must be IPv4Preferred, IPv4Required, IPv6Preferred, IPv6Required", protocol);
        }

        if (pool) {
                pool->reschedule(task, 0);
        }

        return callback;
}

bool DNS::receive(time_t delay, time_t now) throw () {
        fd_set rfds;
        struct timeval tv = { (delay < 0) ? 0 : delay, 0 };

        std::lock_guard<std::mutex> lg(mutex);

        if (pool && getActive() > 0) {
                pool->reschedule(task, 0, interval);
        }

#if 0 /* Only delay when active requests exist */
        if (::dns_timeouts(this->context, -1, now) >= 0) { /* Returns -1 with no active requests */
                int s = getSocket();

                FD_ZERO(&rfds);
                FD_SET(s, &rfds);

                tv.tv_sec = delay;
                tv.tv_usec = 0;

                if (select(s + 1, &rfds, NULL, NULL, &tv) > 0) {
                        ::dns_ioevent(this->context, now);

                        return true;
                }
        }
#else
        int s = getSocket();

        FD_ZERO(&rfds);
        FD_SET(s, &rfds);

        ::dns_timeouts(this->context, -1, now);

        if (select(s + 1, &rfds, NULL, NULL, &tv) > 0) {
                ::dns_ioevent(this->context, now);

                return true;
        }
#endif

        return false;
}

void DNS::setInterval(uint64_t usec) throw () {
        this->interval = (usec > 0) ? usec : DefaultInterval;
}

void DNS::setNoSearch(bool on) throw () {
        flags = dns_set_opt(context, DNS_OPT_FLAGS, -1); /* -1 triggers get */

        if (on) {
                flags |= DNS_NOSRCH;
        } else {
                flags &= ~DNS_NOSRCH;
        }

        dns_set_opt(context, DNS_OPT_FLAGS, flags);
}

void DNS::setNoRecursion(bool on) throw () {
        flags = dns_set_opt(context, DNS_OPT_FLAGS, -1); /* -1 triggers get */

        if (on) {
                flags |= DNS_NORD;
        } else {
                flags &= ~DNS_NORD;
        }

        dns_set_opt(context, DNS_OPT_FLAGS, flags);
}

void DNS::setAuthoritativeOnly(bool on) throw () {
        flags = dns_set_opt(context, DNS_OPT_FLAGS, -1); /* -1 triggers get */

        if (on) {
                flags |= DNS_AAONLY;
        } else {
                flags &= ~DNS_AAONLY;
        }

        dns_set_opt(context, DNS_OPT_FLAGS, flags);
}

void DNS::setTimeout(int seconds) throw () {
        dns_set_opt(context, DNS_OPT_TIMEOUT, seconds);
}

void DNS::setRetries(int retries) throw () {
        dns_set_opt(context, DNS_OPT_NTRIES, retries);
}

void DNS::setDots(int ndots) throw () {
        dns_set_opt(context, DNS_OPT_NDOTS, ndots);
}

void DNS::setUDPSize(int bytes) throw () {
        dns_set_opt(context, DNS_OPT_UDPSIZE, bytes);
}

void DNS::setUDPPort(short port) throw () {
        dns_set_opt(context, DNS_OPT_PORT, (int) port);
}

int DNS::getActive() throw () {
        return dns_active(context);
}

int DNS::getSocket() throw () {
        return dns_sock(context);
}

const DNS::shared DNS::shared_ptr() throw () {
        return std::make_shared<DNS>();
}

void DNS::callback_a(struct dns_ctx * context, void * result, void * argument) {
        DNSCallback *callback = reinterpret_cast<DNSCallback *>(argument);
        struct dns_rr_a4 *response = reinterpret_cast<struct dns_rr_a4 *>(result);

        if (!response || response->dnsa4_nrr == 0) {
                if (callback->protocol == IPv4Required) {
                        callback->failure(callback->qname, static_cast<DNSCallback::Error>(::dns_status(context)));
                } else {
                        callback->protocol = IPv6Required;

                        ::dns_submit_p(context, response->dnsa4_qname, DNS_C_IN, DNS_T_AAAA, callback->flags, ::dns_parse_a6, DNS::callback_a, callback);
                }
        } else {
                std::vector<IPAddress> resolved;

                for (int i = 0; i < response->dnsa4_nrr; i++) {
                        try {
                                resolved.push_back(IPAddress(&response->dnsa4_addr[i]));
                        } catch (AddressException &e) { }
                }

                if (strcmp(response->dnsa4_qname, response->dnsa4_cname) != 0) {
                        callback->success(response->dnsa4_qname, response->dnsa4_cname, response->dnsa4_ttl, resolved);
                } else {
                        callback->success(response->dnsa4_qname, std::string(), response->dnsa4_ttl, resolved);
                }
        }
}

void DNS::callback_aaaa(struct dns_ctx * context, void * result, void * argument) {
        DNSCallback *callback = reinterpret_cast<DNSCallback *>(argument);
        struct dns_rr_a6 *response = reinterpret_cast<struct dns_rr_a6 *>(result);

        if (!response || response->dnsa6_nrr == 0) {
                if (callback->protocol == IPv6Required) {
                        callback->failure(callback->qname, static_cast<DNSCallback::Error>(::dns_status(context)));
                } else {
                        callback->protocol = IPv4Required;

                        ::dns_submit_p(context, response->dnsa6_qname, DNS_C_IN, DNS_T_A, callback->flags, ::dns_parse_a4, DNS::callback_a, callback);
                }
        } else {
                std::vector<IPAddress> resolved;

                for (int i = 0; i < response->dnsa6_nrr; i++) {
                        try {
                                resolved.push_back(IPAddress(&response->dnsa6_addr[i]));
                        } catch (AddressException &e) { }
                }

                if (strcmp(response->dnsa6_qname, response->dnsa6_cname) != 0) {
                        callback->success(response->dnsa6_qname, response->dnsa6_cname, response->dnsa6_ttl, resolved);
                } else {
                        callback->success(response->dnsa6_qname, std::string(), response->dnsa6_ttl, resolved);
                }
        }
}

void DNS::init(const std::vector<std::string> &nameservers, const std::vector<std::string> &searches) {
        if (dns_init(NULL, 0) < 0) { /* Sets up internal default context */
                THROW(DNSException, "failed to initialize default DNS context");
        }

        if (!(this->context = dns_new(NULL))) { /* Copies internal default context */
                THROW(DNSException, "failed to allocate DNS context");
        }

        if (::dns_init(this->context, 0) < 0) {
                THROW(DNSException, "failed to initialize DNS context");
        }

        if (!nameservers.empty()) {
                ::dns_add_serv(context, NULL);

                for (auto n = nameservers.begin(); n != nameservers.end(); ++n) {
                        if (::dns_add_serv(context, n->c_str()) < 0) {
                                THROW2(DNSException, "failed to set nameserver '%s'", n->c_str());
                        }
                }
        }

        if (!searches.empty()) {
                ::dns_add_srch(context, NULL);

                for (auto s = searches.begin(); s != searches.end(); ++s) {
                        if (::dns_add_srch(context, s->c_str()) < 0) {
                                THROW2(DNSException, "failed to set search '%s'", s->c_str());
                        }
                }
        }

        this->flags = dns_set_opt(context, DNS_OPT_FLAGS, -1); /* -1 triggers get */

        if (::dns_open(this->context) < 1) {
                THROW(DNSException, "failed to open DNS socket");
        }
}

bool walk(std::vector<IPAddress> &ips, struct addrinfo *result, union __sockaddr_union &u, int protocol) {
        bool match = false;
        struct addrinfo *rp;

        try {
                for (rp = result; rp != NULL; rp = rp->ai_next) {
                        if ((protocol == AF_INET || protocol == AF_UNSPEC) && rp->ai_family == AF_INET) {
                                match = true;
                                memcpy(u.sa_in, rp->ai_addr, rp->ai_addrlen);
                                u.ss->ss_family = protocol;

                                ips.push_back(IPAddress(u.ss));
                        } else if ((protocol == AF_INET6 || protocol == AF_UNSPEC) && rp->ai_family == AF_INET6) {
                                match = true;
                                memcpy(u.sa_in6, rp->ai_addr, rp->ai_addrlen);
                                u.ss->ss_family = protocol;

                                ips.push_back(IPAddress(u.ss));
                        }
                }
        } catch (AddressException &e) {
                RETHROW(DNSException, e);
        }

        return match;
}

/*
** vim: noet ts=3 sw=3
*/

---