Use aiodns instead of dnspython to query DNS records

2 files changed, 116 insertions, 118 deletions

diff --git a/slixmpp/xmlstream/resolver.py b/slixmpp/xmlstream/resolver.py
index d7fe1ea6..b4144cc3 100644
--- a/slixmpp/xmlstream/resolver.py
+++ b/slixmpp/xmlstream/resolver.py
@@ -8,6 +8,7 @@
     :license: MIT, see LICENSE for more details
 """
 
+import asyncio
 import socket
 import logging
 import random
@@ -16,51 +17,42 @@ import random
 log = logging.getLogger(__name__)
 
 
-#: Global flag indicating the availability of the ``dnspython`` package.
-#: Installing ``dnspython`` can be done via:
+#: Global flag indicating the availability of the ``aiodns`` package.
+#: Installing ``aiodns`` can be done via:
 #:
 #: .. code-block:: sh
 #:
-#:     pip install dnspython
-#:
-#: For Python3, installation may require installing from source using
-#: the ``python3`` branch:
-#:
-#: .. code-block:: sh
-#:
-#:     git clone http://github.com/rthalley/dnspython
-#:     cd dnspython
-#:     git checkout python3
-#:     python3 setup.py install
-DNSPYTHON_AVAILABLE = False
+#:     pip install aiodns
+AIODNS_AVAILABLE = False
 try:
-    import dns.resolver
-    DNSPYTHON_AVAILABLE = True
+    import aiodns
+    AIODNS_AVAILABLE = True
 except ImportError as e:
-    log.debug("Could not find dnspython package. " + \
+    log.debug("Could not find aiodns package. " + \
               "Not all features will be available")
 
 
 def default_resolver():
     """Return a basic DNS resolver object.
 
-    :returns: A :class:`dns.resolver.Resolver` object if dnspython
+    :returns: A :class:`aiodns.DNSResolver` object if aiodns
               is available. Otherwise, ``None``.
     """
-    if DNSPYTHON_AVAILABLE:
-        return dns.resolver.get_default_resolver()
+    if AIODNS_AVAILABLE:
+        return aiodns.DNSResolver(loop=asyncio.get_event_loop())
     return None
 
 
+@asyncio.coroutine
 def resolve(host, port=None, service=None, proto='tcp',
-            resolver=None, use_ipv6=True, use_dnspython=True):
+            resolver=None, use_ipv6=True, use_aiodns=True):
     """Peform DNS resolution for a given hostname.
 
     Resolution may perform SRV record lookups if a service and protocol
     are specified. The returned addresses will be sorted according to
     the SRV priorities and weights.
 
-    If no resolver is provided, the dnspython resolver will be used if
+    If no resolver is provided, the aiodns resolver will be used if
     available. Otherwise the built-in socket facilities will be used,
     but those do not provide SRV support.
 
@@ -77,7 +69,7 @@ def resolve(host, port=None, service=None, proto='tcp',
     :param use_ipv6: Optionally control the use of IPv6 in situations
                      where it is either not available, or performance
                      is degraded. Defaults to ``True``.
-    :param use_dnspython: Optionally control if dnspython is used to make
+    :param use_aiodns: Optionally control if aiodns is used to make
                           the DNS queries instead of the built-in DNS
                           library.
 
@@ -85,25 +77,25 @@ def resolve(host, port=None, service=None, proto='tcp',
     :type     port: int
     :type  service: string
     :type    proto: string
-    :type resolver: :class:`dns.resolver.Resolver`
+    :type resolver: :class:`aiodns.DNSResolver`
     :type use_ipv6: bool
-    :type use_dnspython: bool
+    :type use_aiodns: bool
 
     :return: An iterable of IP address, port pairs in the order
              dictated by SRV priorities and weights, if applicable.
     """
 
-    if not use_dnspython:
-        if DNSPYTHON_AVAILABLE:
-            log.debug("DNS: Not using dnspython, but dnspython is installed.")
+    if not use_aiodns:
+        if AIODNS_AVAILABLE:
+            log.debug("DNS: Not using aiodns, but aiodns is installed.")
         else:
-            log.debug("DNS: Not using dnspython.")
+            log.debug("DNS: Not using aiodns.")
 
     if not use_ipv6:
         log.debug("DNS: Use of IPv6 has been disabled.")
 
-    if resolver is None and DNSPYTHON_AVAILABLE and use_dnspython:
-        resolver = dns.resolver.get_default_resolver()
+    if resolver is None and AIODNS_AVAILABLE and use_aiodns:
+        resolver = aiodns.DNSResolver(loop=asyncio.get_event_loop())
 
     # An IPv6 literal is allowed to be enclosed in square brackets, but
     # the brackets must be stripped in order to process the literal;
@@ -113,7 +105,7 @@ def resolve(host, port=None, service=None, proto='tcp',
     try:
         # If `host` is an IPv4 literal, we can return it immediately.
         ipv4 = socket.inet_aton(host)
-        yield (host, host, port)
+        return [(host, host, port)]
     except socket.error:
         pass
 
@@ -123,7 +115,7 @@ def resolve(host, port=None, service=None, proto='tcp',
             # it immediately.
             if hasattr(socket, 'inet_pton'):
                 ipv6 = socket.inet_pton(socket.AF_INET6, host)
-                yield (host, host, port)
+                return [(host, host, port)]
         except (socket.error, ValueError):
             pass
 
@@ -133,29 +125,31 @@ def resolve(host, port=None, service=None, proto='tcp',
     if not service:
         hosts = [(host, port)]
     else:
-        hosts = get_SRV(host, port, service, proto,
-                        resolver=resolver,
-                        use_dnspython=use_dnspython)
-
+        hosts = yield from get_SRV(host, port, service, proto,
+                                   resolver=resolver,
+                                   use_aiodns=use_aiodns)
+    results = []
     for host, port in hosts:
-        results = []
         if host == 'localhost':
             if use_ipv6:
                 results.append((host, '::1', port))
             results.append((host, '127.0.0.1', port))
+
         if use_ipv6:
-            for address in get_AAAA(host, resolver=resolver,
-                                          use_dnspython=use_dnspython):
+            aaaa = yield from get_AAAA(host, resolver=resolver,
+                                       use_aiodns=use_aiodns)
+            for address in aaaa:
                 results.append((host, address, port))
-        for address in get_A(host, resolver=resolver,
-                                   use_dnspython=use_dnspython):
-            results.append((host, address, port))
 
-        for host, address, port in results:
-            yield host, address, port
+        a = yield from get_A(host, resolver=resolver,
+                             use_aiodns=use_aiodns)
+        for address in a:
+            results.append((host, address, port))
 
+    return results
 
-def get_A(host, resolver=None, use_dnspython=True):
+@asyncio.coroutine
+def get_A(host, resolver=None, use_aiodns=True):
     """Lookup DNS A records for a given host.
 
     If ``resolver`` is not provided, or is ``None``, then resolution will
@@ -163,46 +157,41 @@ def get_A(host, resolver=None, use_dnspython=True):
 
     :param     host: The hostname to resolve for A record IPv4 addresses.
     :param resolver: Optional DNS resolver object to use for the query.
-    :param use_dnspython: Optionally control if dnspython is used to make
+    :param use_aiodns: Optionally control if aiodns is used to make
                           the DNS queries instead of the built-in DNS
                           library.
 
     :type     host: string
-    :type resolver: :class:`dns.resolver.Resolver` or ``None``
-    :type use_dnspython: bool
+    :type resolver: :class:`aiodns.DNSResolver` or ``None``
+    :type use_aiodns: bool
 
     :return: A list of IPv4 literals.
     """
     log.debug("DNS: Querying %s for A records." % host)
 
-    # If not using dnspython, attempt lookup using the OS level
+    # If not using aiodns, attempt lookup using the OS level
     # getaddrinfo() method.
-    if resolver is None or not use_dnspython:
+    if resolver is None or not use_aiodns:
         try:
             recs = socket.getaddrinfo(host, None, socket.AF_INET,
                                                   socket.SOCK_STREAM)
             return [rec[4][0] for rec in recs]
         except socket.gaierror:
-            log.debug("DNS: Error retreiving A address info for %s." % host)
+            log.debug("DNS: Error retrieving A address info for %s." % host)
             return []
 
-    # Using dnspython:
+    # Using aiodns:
+    future = resolver.query(host, 'A')
     try:
-        recs = resolver.query(host, dns.rdatatype.A)
-        return [rec.to_text() for rec in recs]
-    except (dns.resolver.NXDOMAIN, dns.resolver.NoAnswer):
-        log.debug("DNS: No A records for %s" % host)
-        return []
-    except dns.exception.Timeout:
-        log.debug("DNS: A record resolution timed out for %s" % host)
-        return []
-    except dns.exception.DNSException as e:
-        log.debug("DNS: Error querying A records for %s" % host)
-        log.exception(e)
-        return []
+        recs = yield from future
+    except Exception as e:
+        log.debug('DNS: Exception while querying for %s A records: %s', host, e)
+        recs = []
+    return recs
 
 
-def get_AAAA(host, resolver=None, use_dnspython=True):
+@asyncio.coroutine
+def get_AAAA(host, resolver=None, use_aiodns=True):
     """Lookup DNS AAAA records for a given host.
 
     If ``resolver`` is not provided, or is ``None``, then resolution will
@@ -210,23 +199,23 @@ def get_AAAA(host, resolver=None, use_dnspython=True):
 
     :param     host: The hostname to resolve for AAAA record IPv6 addresses.
     :param resolver: Optional DNS resolver object to use for the query.
-    :param use_dnspython: Optionally control if dnspython is used to make
+    :param use_aiodns: Optionally control if aiodns is used to make
                           the DNS queries instead of the built-in DNS
                           library.
 
     :type     host: string
-    :type resolver: :class:`dns.resolver.Resolver` or ``None``
-    :type use_dnspython: bool
+    :type resolver: :class:`aiodns.DNSResolver` or ``None``
+    :type use_aiodns: bool
 
     :return: A list of IPv6 literals.
     """
     log.debug("DNS: Querying %s for AAAA records." % host)
 
-    # If not using dnspython, attempt lookup using the OS level
+    # If not using aiodns, attempt lookup using the OS level
     # getaddrinfo() method.
-    if resolver is None or not use_dnspython:
+    if resolver is None or not use_aiodns:
         if not socket.has_ipv6:
-            log.debug("Unable to query %s for AAAA records: IPv6 is not supported", host)
+            log.debug("DNS: Unable to query %s for AAAA records: IPv6 is not supported", host)
             return []
         try:
             recs = socket.getaddrinfo(host, None, socket.AF_INET6,
@@ -237,29 +226,23 @@ def get_AAAA(host, resolver=None, use_dnspython=True):
                       "info for %s." % host)
             return []
 
-    # Using dnspython:
+    # Using aiodns:
+    future = resolver.query(host, 'AAAA')
     try:
-        recs = resolver.query(host, dns.rdatatype.AAAA)
-        return [rec.to_text() for rec in recs]
-    except (dns.resolver.NXDOMAIN, dns.resolver.NoAnswer):
-        log.debug("DNS: No AAAA records for %s" % host)
-        return []
-    except dns.exception.Timeout:
-        log.debug("DNS: AAAA record resolution timed out for %s" % host)
-        return []
-    except dns.exception.DNSException as e:
-        log.debug("DNS: Error querying AAAA records for %s" % host)
-        log.exception(e)
-        return []
-
-
-def get_SRV(host, port, service, proto='tcp', resolver=None, use_dnspython=True):
+        recs = yield from future
+    except Exception as e:
+        log.debug('DNS: Exception while querying for %s AAAA records: %s', host, e)
+        recs = []
+    return recs
+
+@asyncio.coroutine
+def get_SRV(host, port, service, proto='tcp', resolver=None, use_aiodns=True):
     """Perform SRV record resolution for a given host.
 
     .. note::
 
-        This function requires the use of the ``dnspython`` package. Calling
-        :func:`get_SRV` without ``dnspython`` will return the provided host
+        This function requires the use of the ``aiodns`` package. Calling
+        :func:`get_SRV` without ``aiodns`` will return the provided host
         and port without performing any DNS queries.
 
     :param     host: The hostname to resolve.
@@ -274,32 +257,23 @@ def get_SRV(host, port, service, proto='tcp', resolver=None, use_dnspython=True)
     :type     port: int
     :type  service: string
     :type    proto: string
-    :type resolver: :class:`dns.resolver.Resolver`
+    :type resolver: :class:`aiodns.DNSResolver`
 
     :return: A list of hostname, port pairs in the order dictacted
              by SRV priorities and weights.
     """
-    if resolver is None or not use_dnspython:
-        log.warning("DNS: dnspython not found. Can not use SRV lookup.")
+    if resolver is None or not use_aiodns:
+        log.warning("DNS: aiodns not found. Can not use SRV lookup.")
         return [(host, port)]
 
     log.debug("DNS: Querying SRV records for %s" % host)
     try:
-        recs = resolver.query('_%s._%s.%s' % (service, proto, host),
-                              dns.rdatatype.SRV)
-    except (dns.resolver.NXDOMAIN, dns.resolver.NoAnswer):
-        log.debug("DNS: No SRV records for %s." % host)
-        return [(host, port)]
-    except dns.exception.Timeout:
-        log.debug("DNS: SRV record resolution timed out for %s." % host)
-        return [(host, port)]
-    except dns.exception.DNSException as e:
-        log.debug("DNS: Error querying SRV records for %s." % host)
-        log.exception(e)
-        return [(host, port)]
-
-    if len(recs) == 1 and recs[0].target == '.':
-        return [(host, port)]
+        future = resolver.query('_%s._%s.%s' % (service, proto, host),
+                                'SRV')
+        recs = yield from future
+    except Exception as e:
+        log.debug('DNS: Exception while querying for %s SRV records: %s', host, e)
+        return []
 
     answers = {}
     for rec in recs:
@@ -323,7 +297,7 @@ def get_SRV(host, port, service, proto='tcp', resolver=None, use_dnspython=True)
             for running_sum in sums:
                 if running_sum >= selected:
                     rec = sums[running_sum]
-                    host = rec.target.to_text()
+                    host = rec.host
                     if host.endswith('.'):
                         host = host[:-1]
                     sorted_recs.append((host, rec.port))
diff --git a/slixmpp/xmlstream/xmlstream.py b/slixmpp/xmlstream/xmlstream.py
index f4f0b87e..40080dc1 100644
--- a/slixmpp/xmlstream/xmlstream.py
+++ b/slixmpp/xmlstream/xmlstream.py
@@ -162,7 +162,7 @@ class XMLStream(object):
         #: If set to ``True``, allow using the ``dnspython`` DNS library
         #: if available. If set to ``False``, the builtin DNS resolver
         #: will be used, even if ``dnspython`` is installed.
-        self.use_dnspython = True
+        self.use_aiodns = True
 
         #: Use CDATA for escaping instead of XML entities. Defaults
         #: to ``False``.
@@ -287,13 +287,32 @@ class XMLStream(object):
     def _connect_routine(self):
         loop = asyncio.get_event_loop()
         self.event_when_connected = "connected"
+
+        try:
+            record = yield from self.pick_dns_answer(self.default_domain)
+        except StopIteration:
+            # No more DNS records to try
+            self.dns_answers = None
+            return
+        else:
+            if record:
+                host, address, port = record
+                self._service_name = host
+            else:
+                self.event('connection_failed',
+                           'No DNS record available for %s' % self.default_domain)
+                self.dns_answers = None
+                return
+
         try:
             yield from loop.create_connection(lambda: self,
-                                              self.address[0],
-                                              self.address[1],
+                                              address,
+                                              port,
                                               ssl=self.use_ssl)
         except OSError as e:
+            log.debug('Connection failed: %s', e)
             self.event("connection_failed", e)
+            asyncio.async(self._connect_routine())
 
     def process(self, timeout=None):
         """Process all the available XMPP events (receiving or sending data on the
@@ -578,6 +597,7 @@ class XMLStream(object):
             idx += 1
         return False
 
+    @asyncio.coroutine
     def get_dns_records(self, domain, port=None):
         """Get the DNS records for a domain.
 
@@ -590,11 +610,14 @@ class XMLStream(object):
         resolver = default_resolver()
         self.configure_dns(resolver, domain=domain, port=port)
 
-        return resolve(domain, port, service=self.dns_service,
-                                     resolver=resolver,
-                                     use_ipv6=self.use_ipv6,
-                                     use_dnspython=self.use_dnspython)
+        result = yield from resolve(domain, port,
+                                    service=self.dns_service,
+                                    resolver=resolver,
+                                    use_ipv6=self.use_ipv6,
+                                    use_aiodns=self.use_aiodns)
+        return result
 
+    @asyncio.coroutine
     def pick_dns_answer(self, domain, port=None):
         """Pick a server and port from DNS answers.
 
@@ -604,8 +627,9 @@ class XMLStream(object):
         :param domain: The domain in question.
         :param port: If the results don't include a port, use this one.
         """
-        if not self.dns_answers:
-            self.dns_answers = self.get_dns_records(domain, port)
+        if self.dns_answers is None:
+            dns_records = yield from self.get_dns_records(domain, port)
+            self.dns_answers = iter(dns_records)
 
         return next(self.dns_answers)