DNS.cpp

/*
** libutil++
** $Id: DNS.cpp 1884 2017-04-15 21:42:02Z sella $
** Copyright (c) 2011-2016 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 1884 2017-04-15 21:42:02Z sella $";

#include "DNS.h"
#include "DNSException.h"
#include "DNSCallback.h"
#include "../util/ThreadPool.h"
#include "../util/Time.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>
#include <udns.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<std::string> DNS::resolve(const std::string &name, int record, int protocol) throw (DNSException) {
        StoreDNSCallback c;

        resolve(&c, name, record, protocol);

        while (getActive() > 0) {
                receive(250000);
        }

        return c.getResponses();
}

sella::variant::Variant DNS::resolve(const std::vector<std::string> &names, int record, int protocol) throw (DNSException) {
        sella::variant::Variant responses;
        StoreDNSCallback c[names.size()];

        for (size_t i = 0; i < names.size(); i++) {
                resolve(&(c[i]), names[i], record, protocol);
        }

        while (getActive() > 0) {
                receive(250000);
        }

        for (size_t i = 0; i < names.size(); i++) {
                const auto &r = c[i].getResponses();
                auto &response = responses[c[i].qname]["response"];

                for (auto i = r.begin(); i != r.end(); ++i) {
                        response.push_back(*i);
                }

                responses[c[i].qname]["qname"] = c[i].getQname();
                responses[c[i].qname]["ttl"] = c[i].getTTL();
                responses[c[i].qname]["error"] = c[i].getErrorName();
        }

        return responses;
}

std::vector<std::string> DNS::bulkResolve(const std::vector<std::string> &names, int record, int protocol) throw (DNSException) {
        std::vector<std::string> responses;
        StoreDNSCallback c[names.size()];

        for (size_t i = 0; i < names.size(); i++) {
                resolve(&(c[i]), names[i], record, protocol);
        }

        while (getActive() > 0) {
                receive(250000);
        }

        for (size_t i = 0; i < names.size(); i++) {
                const auto &r = c[i].getResponses();

                responses.insert(responses.end(), r.begin(), r.end());
        }

        return responses;
}

std::vector<std::string> DNS::sresolve(const std::string &name, int protocol) throw (DNSException) {
        const std::vector<IPAddress> &addresses = DNS::sresolve_ip(name, protocol);
        std::vector<std::string> response;

        for (auto a = addresses.begin(); a != addresses.end(); ++a) {
                response.push_back(a->str());
        }

        return response;
}

std::vector<IPAddress> DNS::sresolve_ip(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;
        } else {
                THROW(DNSException, "unknown protocol (%d)", protocol);
        }

        if ((s = getaddrinfo(name.c_str(), NULL, &hints, &result)) != 0) {
                if (s == EAI_NODATA) {
                        return ips;
                }

                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_UNSPEC);
                } else if (protocol == IPv4Preferred) {
                        if (!walk(ips, result, u, AF_INET)) {
                                walk(ips, result, u, AF_INET6);
                        }
                } else if (protocol == IPv4Required) {
                        walk(ips, result, u, AF_INET);
                } else if (protocol == IPv6Preferred) {
                        if (!walk(ips, result, u, AF_INET6)) {
                                walk(ips, result, u, AF_INET);
                        }
                } else if (protocol == IPv6Required) {
                        walk(ips, result, u, AF_INET6);
                }

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

                throw e;
        }

        return ips;
}

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

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

        switch (record) {
                case A:
                        switch (protocol) {
                                case AnyProtocol:
                                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 AnyProtocol, IPv4Preferred, IPv4Required, IPv6Preferred, IPv6Required", protocol);
                        }
                        break;

                case PTR:
                        try {
                                IPAddress ip(name);

                                if (ip.getFamily() == AF_INET) {
                                        ::dns_submit_a4ptr(this->context, &(ip.getIPv4AddressStructure()), DNS::callback_ptr, callback);
                                } else {
                                        ::dns_submit_a6ptr(this->context, &(ip.getIPv6AddressStructure()), DNS::callback_ptr, callback);
                                }

                                break;
                        } catch (AddressException &e) {
                                RETHROW(DNSException, e);
                        }

                default:
                        THROW(DNSException, "record type %d has not been implemented yet", record);
        }


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

        return callback;
}

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

        if (usec > 0) {
                tv = sella::util::Time::usectotv(usec);
        }

        if (!now) {
                now = time(NULL);
        }

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

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

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

        int s = dns_sock(context);

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

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

                return true;
        }

        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::getFD() throw () {
        return dns_sock(context);
}

const DNS::shared DNS::shared_ptr(const std::vector<std::string> &nameservers, const std::vector<std::string> &searches) throw (DNSException) {
        return std::make_shared<DNS>(nameservers, searches);
}

const DNS::shared DNS::shared_ptr(sella::util::ThreadPool *pool, const std::vector<std::string> &nameservers, const std::vector<std::string> &searches) throw (DNSException) {
        return std::make_shared<DNS>(pool, nameservers, searches);
}

const IPAddress DNS::convertArpa(const std::string &qname) throw (DNSException) {
        if (qname.size() < 5 || qname.substr(qname.size() - 5).compare(".arpa") != 0) {
                THROW(DNSException, "Input is not a valid .in-addr.arpa or .ip6.arapa record");
        }

        try {
                std::string buf;
                const auto &tokens = String::split(qname, '.', false);

                if (tokens.size() == 6) {
                        buf.append(tokens[3]).append(".").append(tokens[2]).append(".").append(tokens[1]).append(".").append(tokens[0]);
                } else if (tokens.size() == 34) {
                        for (ssize_t i = 31; i >= 0; i--) {
                                buf.append(tokens[i]);

                                if (i % 4 == 0 && i > 0) {
                                        buf.append(":");
                                }       
                        }
                } else {
                        THROW(DNSException, "Input is not a valid .in-addr.arpa or .ip6.arapa record");
                }

                return IPAddress(buf);
        } catch (AddressException &e) {
                THROW(DNSException, "Input is not a valid .in-addr.arpa or .ip6.arapa record");
        }
}

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

        //printf("callback_a - %s\n", callback->qname.c_str());
        //usleep(99000);

        int status = ::dns_status(context);

        if (status != DNS_E_NOERROR && status != DNS_E_NODATA) {
//              printf("callback_a - failure\n");
                callback->failure(callback->qname, static_cast<DNSCallback::Error>(::dns_status(context)));
        } else if (!response || response->dnsa4_nrr == 0) {
                //printf("callback_a - !response\n");
                if (callback->protocol != IPv4Required) {
                        if (callback->protocol == IPv4Preferred) {
                                callback->protocol = IPv6Required;
                        }

                        //printf("callback_a - go v6\n");
                        ::dns_submit_p(context, callback->qname.c_str(), DNS_C_IN, DNS_T_AAAA, callback->flags, ::dns_parse_a6, DNS::callback_aaaa, callback);
                } else {
                        //printf("callback_a - success (final)\n");
                        callback->success(callback->qname, std::string(), -1, resolved);
                }
        } else {
                //printf("callback_a - got: %d\n", response->dnsa4_nrr);
                for (int i = 0; i < response->dnsa4_nrr; i++) {
                        try {
                                resolved.push_back(IPAddress(&response->dnsa4_addr[i]).str());
                        } catch (AddressException &e) { }
                }

                //printf("callback_a - success\n");
                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);
                }

                if (callback->protocol != IPv4Required && callback->protocol != IPv4Preferred) {
                        //printf("callback_a - go v6: %d\n", callback->protocol);
                        ::dns_submit_p(context, response->dnsa4_qname, DNS_C_IN, DNS_T_AAAA, callback->flags, ::dns_parse_a6, DNS::callback_aaaa, callback);
                }
        }
}

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

        //printf("callback_aaaa - %s\n", callback->qname.c_str());
        //usleep(99000);

        int status = ::dns_status(context);

        if (status != DNS_E_NOERROR && status != DNS_E_NODATA) {
                //printf("callback_aaaa - failure\n");
                callback->failure(callback->qname, static_cast<DNSCallback::Error>(::dns_status(context)));
        } else if (!response || response->dnsa6_nrr == 0) {
                //printf("callback_aaaa - !response\n");
                if (callback->protocol != IPv6Required && callback->protocol != AnyProtocol) {
                        if (callback->protocol == IPv6Preferred) {
                                callback->protocol = IPv4Required;
                        }

                        //printf("callback_aaaa - go v4\n");
                        ::dns_submit_p(context, callback->qname.c_str(), DNS_C_IN, DNS_T_A, callback->flags, ::dns_parse_a4, DNS::callback_a, callback);
                } else {
                        //printf("callback_aaaa - success (final)\n");
                        callback->success(callback->qname, std::string(), -1, resolved);
                }
        } else {
                //printf("callback_aaaa - got: %d\n", response->dnsa6_nrr);
                for (int i = 0; i < response->dnsa6_nrr; i++) {
                        try {
                                resolved.push_back(IPAddress(&response->dnsa6_addr[i]).str());
                        } catch (AddressException &e) { }
                }

                //printf("callback_aaaa - success\n");
                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);
                }

                if (callback->protocol != IPv6Required && callback->protocol != IPv6Preferred && callback->protocol != AnyProtocol) {
                        //printf("callback_aaaa - go v4\n");
                        ::dns_submit_p(context, response->dnsa6_qname, DNS_C_IN, DNS_T_A, callback->flags, ::dns_parse_a4, DNS::callback_a, callback);
                }
        }
}

void DNS::callback_ptr(struct dns_ctx *context, dns_rr_ptr *result, void *argument) {
        std::vector<std::string> reverses;
        DNSCallback *callback = reinterpret_cast<DNSCallback*>(argument);
        struct dns_rr_ptr *response = reinterpret_cast<struct dns_rr_ptr*>(result);

//      printf("callback_ptr - %s\n", callback->qname.c_str());
//      usleep(99000);

        int status = ::dns_status(context);

        if (status != DNS_E_NOERROR && status != DNS_E_NODATA) {
//              printf("callback_ptr - failure\n");
                callback->failure(callback->qname, static_cast<DNSCallback::Error>(::dns_status(context)));
        } else if (!response || response->dnsptr_nrr == 0) {
                callback->success(callback->qname, std::string(), -1, reverses);
        } else {
//              printf("callback_ptr - got: %d\n", response->dnsptr_nrr);
                for (int i = 0; i < response->dnsptr_nrr; i++) {
                        reverses.push_back(response->dnsptr_ptr[i]);
                }

//              printf("callback_ptr - success\n");
                if (strcmp(response->dnsptr_qname, response->dnsptr_cname) != 0) {
                        callback->success(response->dnsptr_qname, response->dnsptr_cname, response->dnsptr_ttl, reverses);
                } else {
                        callback->success(response->dnsptr_qname, std::string(), response->dnsptr_ttl, reverses);
                }
        }
}

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
*/