summaryrefslogtreecommitdiff
path: root/sleekxmpp/xmlstream
diff options
context:
space:
mode:
authorLance Stout <lancestout@gmail.com>2012-03-29 15:11:24 -0700
committerLance Stout <lancestout@gmail.com>2012-03-30 10:12:43 -0700
commitc1d36cad4679419679896c81e4d4520ec1627132 (patch)
tree04719dd43b2308a5ab99dbf690dfe94c84008af8 /sleekxmpp/xmlstream
parentaad2eb31fcde345431a27994bdc771f58e480e38 (diff)
downloadslixmpp-c1d36cad4679419679896c81e4d4520ec1627132.tar.gz
slixmpp-c1d36cad4679419679896c81e4d4520ec1627132.tar.bz2
slixmpp-c1d36cad4679419679896c81e4d4520ec1627132.tar.xz
slixmpp-c1d36cad4679419679896c81e4d4520ec1627132.zip
Add better DNS resolver wrapper.
Diffstat (limited to 'sleekxmpp/xmlstream')
-rw-r--r--sleekxmpp/xmlstream/resolver.py287
-rw-r--r--sleekxmpp/xmlstream/xmlstream.py124
2 files changed, 314 insertions, 97 deletions
diff --git a/sleekxmpp/xmlstream/resolver.py b/sleekxmpp/xmlstream/resolver.py
new file mode 100644
index 00000000..ecb76519
--- /dev/null
+++ b/sleekxmpp/xmlstream/resolver.py
@@ -0,0 +1,287 @@
+# -*- encoding: utf-8 -*-
+
+"""
+ sleekxmpp.xmlstream.dns
+ ~~~~~~~~~~~~~~~~~~~~~~~
+
+ :copyright: (c) 2012 Nathanael C. Fritz
+ :license: MIT, see LICENSE for more details
+"""
+
+import socket
+import logging
+import random
+
+
+log = logging.getLogger(__name__)
+
+
+#: Global flag indicating the availability of the ``dnspython`` package.
+#: Installing ``dnspython`` 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
+USE_DNSPYTHON = False
+try:
+ import dns.resolver
+ USE_DNSPYTHON = True
+except ImportError as e:
+ log.debug("Could not find dnspython 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
+ is available. Otherwise, ``None``.
+ """
+ if USE_DNSPYTHON:
+ return dns.resolver.get_default_resolver()
+ return None
+
+
+def resolve(host, port=None, service=None, proto='tcp', resolver=None):
+ """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
+ available. Otherwise the built-in socket facilities will be used,
+ but those do not provide SRV support.
+
+ If SRV records were used, queries to resolve alternative hosts will
+ be made as needed instead of all at once.
+
+ :param host: The hostname to resolve.
+ :param port: A default port to connect with. SRV records may
+ dictate use of a different port.
+ :param service: Optional SRV service name without leading underscore.
+ :param proto: Optional SRV protocol name without leading underscore.
+ :param resolver: Optionally provide a DNS resolver object that has
+ been custom configured.
+
+ :type host: string
+ :type port: int
+ :type service: string
+ :type proto: string
+ :type resolver: :class:`dns.resolver.Resolver`
+
+ :return: An iterable of IP address, port pairs in the order
+ dictated by SRV priorities and weights, if applicable.
+ """
+ if resolver is None and USE_DNSPYTHON:
+ resolver = dns.resolver.get_default_resolver()
+
+ # An IPv6 literal is allowed to be enclosed in square brackets, but
+ # the brackets must be stripped in order to process the literal;
+ # otherwise, things break.
+ host = host.strip('[]')
+
+ try:
+ # If `host` is an IPv4 literal, we can return it immediately.
+ ipv4 = socket.inet_pton(socket.AF_INET, host)
+ yield [(host, port)]
+ except socket.error:
+ pass
+
+ try:
+ # Likewise, If `host` is an IPv6 literal, we can return it immediately.
+ ipv6 = socket.inet_pton(socket.AF_INET6, host)
+ yield [(host, port)]
+ except socket.error:
+ pass
+
+ # If no service was provided, then we can just do A/AAAA lookups on the
+ # provided host. Otherwise we need to get an ordered list of hosts to
+ # resolve based on SRV records.
+ if not service:
+ hosts = [(host, port)]
+ else:
+ hosts = get_SRV(host, port, service, proto, resolver=resolver)
+
+ for host, port in hosts:
+ results = []
+ for address in get_AAAA(host, resolver=resolver):
+ results.append((address, port))
+ for address in get_A(host, resolver=resolver):
+ results.append((address, port))
+
+ for address, port in results:
+ yield address, port
+
+
+def get_A(host, resolver=None):
+ """Lookup DNS A records for a given host.
+
+ If ``resolver`` is not provided, or is ``None``, then resolution will
+ be performed using the built-in :mod:`socket` module.
+
+ :param host: The hostname to resolve for A record IPv4 addresses.
+ :param resolver: Optional DNS resolver object to use for the query.
+
+ :type host: string
+ :type resolver: :class:`dns.resolver.Resolver` or ``None``
+
+ :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
+ # getaddrinfo() method.
+ if resolver is None:
+ 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)
+ return []
+
+ # Using dnspython:
+ 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 []
+
+
+def get_AAAA(host, resolver=None):
+ """Lookup DNS AAAA records for a given host.
+
+ If ``resolver`` is not provided, or is ``None``, then resolution will
+ be performed using the built-in :mod:`socket` module.
+
+ :param host: The hostname to resolve for AAAA record IPv6 addresses.
+ :param resolver: Optional DNS resolver object to use for the query.
+
+ :type host: string
+ :type resolver: :class:`dns.resolver.Resolver` or ``None``
+
+ :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
+ # getaddrinfo() method.
+ if resolver is None:
+ try:
+ recs = socket.getaddrinfo(host, None, socket.AF_INET6,
+ socket.SOCK_STREAM)
+ return [rec[4][0] for rec in recs]
+ except socket.gaierror:
+ log.debug("DNS: Error retreiving AAAA address " + \
+ "info for %s." % host)
+ return []
+
+ # Using dnspython:
+ 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):
+ """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
+ and port without performing any DNS queries.
+
+ :param host: The hostname to resolve.
+ :param port: A default port to connect with. SRV records may
+ dictate use of a different port.
+ :param service: Optional SRV service name without leading underscore.
+ :param proto: Optional SRV protocol name without leading underscore.
+ :param resolver: Optionally provide a DNS resolver object that has
+ been custom configured.
+
+ :type host: string
+ :type port: int
+ :type service: string
+ :type proto: string
+ :type resolver: :class:`dns.resolver.Resolver`
+
+ :return: A list of hostname, port pairs in the order dictacted
+ by SRV priorities and weights.
+ """
+ if resolver is None:
+ return [(host, port)]
+
+ log.debug("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)]
+
+ answers = {}
+ for rec in recs:
+ if rec.priority not in answers:
+ answers[rec.priority] = []
+ if rec.weight == 0:
+ answers[rec.priority].insert(0, rec)
+ else:
+ answers[rec.priority].append(rec)
+
+ sorted_recs = []
+ for priority in sorted(answers.keys()):
+ while answers[priority]:
+ running_sum = 0
+ sums = {}
+ for rec in answers[priority]:
+ running_sum += rec.weight
+ sums[running_sum] = rec
+
+ selected = random.randint(0, running_sum + 1)
+ for running_sum in sums:
+ if running_sum >= selected:
+ rec = sums[running_sum]
+ sorted_recs.append((rec.target.to_text(), rec.port))
+ answers[priority].remove(rec)
+ break
+
+ return sorted_recs
diff --git a/sleekxmpp/xmlstream/xmlstream.py b/sleekxmpp/xmlstream/xmlstream.py
index 45f27929..31ca9cfe 100644
--- a/sleekxmpp/xmlstream/xmlstream.py
+++ b/sleekxmpp/xmlstream/xmlstream.py
@@ -1,4 +1,3 @@
-# -*- coding: utf-8 -*-
"""
sleekxmpp.xmlstream.xmlstream
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -39,19 +38,13 @@ from sleekxmpp.xmlstream import Scheduler, tostring
from sleekxmpp.xmlstream.stanzabase import StanzaBase, ET, ElementBase
from sleekxmpp.xmlstream.handler import Waiter, XMLCallback
from sleekxmpp.xmlstream.matcher import MatchXMLMask
+from sleekxmpp.xmlstream.resolver import resolve, default_resolver
# In Python 2.x, file socket objects are broken. A patched socket
# wrapper is provided for this case in filesocket.py.
if sys.version_info < (3, 0):
from sleekxmpp.xmlstream.filesocket import FileSocket, Socket26
-try:
- import dns.resolver
-except ImportError:
- DNSPYTHON = False
-else:
- DNSPYTHON = True
-
#: The time in seconds to wait before timing out waiting for response stanzas.
RESPONSE_TIMEOUT = 30
@@ -306,6 +299,11 @@ class XMLStream(object):
#: A list of DNS results that have not yet been tried.
self.dns_answers = []
+ #: The service name to check with DNS SRV records. For
+ #: example, setting this to ``'xmpp-client'`` would query the
+ #: ``_xmpp-client._tcp`` service.
+ self.dns_service = None
+
self.add_event_handler('connected', self._handle_connected)
self.add_event_handler('session_start', self._start_keepalive)
self.add_event_handler('disconnected', self._end_keepalive)
@@ -445,25 +443,10 @@ class XMLStream(object):
self.stop.set()
return False
- try:
- # Look for IPv6 addresses, in addition to IPv4
- for res in Socket.getaddrinfo(self.address[0],
- int(self.address[1]),
- 0,
- Socket.SOCK_STREAM):
- log.debug("Trying: %s", res[-1])
- af, sock_type, proto, canonical, sock_addr = res
- try:
- self.socket = self.socket_class(af, sock_type, proto)
- break
- except Socket.error:
- log.debug("Could not open IPv%s socket." % proto)
- except Socket.gaierror:
- log.warning("Socket could not be opened: no connectivity" + \
- " or wrong IP versions.")
- if reattempt:
- self.reconnect_delay = delay
- return False
+ af = Socket.AF_INET
+ if ':' in self.address[0]:
+ af = Socket.AF_INET6
+ self.socket = self.socket_class(af, Socket.SOCK_STREAM)
self.configure_socket()
@@ -511,7 +494,10 @@ class XMLStream(object):
except Socket.error as serr:
error_msg = "Could not connect to %s:%s. Socket Error #%s: %s"
self.event('socket_error', serr, direct=True)
- log.error(error_msg, self.address[0], self.address[1],
+ domain = self.address[0]
+ if ':' in domain:
+ domain = '[%s]' % domain
+ log.error(error_msg, domain, self.address[1],
serr.errno, serr.strerror)
if reattempt:
self.reconnect_delay = delay
@@ -915,50 +901,11 @@ class XMLStream(object):
"""
if port is None:
port = self.default_port
- if DNSPYTHON:
- resolver = dns.resolver.get_default_resolver()
- self.configure_dns(resolver, domain=domain, port=port)
-
- v4_answers = []
- v6_answers = []
- answers = []
-
- try:
- log.debug("Querying A records for %s" % domain)
- v4_answers = resolver.query(domain, dns.rdatatype.A)
- except (dns.resolver.NXDOMAIN, dns.resolver.NoAnswer):
- log.warning("No A records for %s", domain)
- v4_answers = [((domain, port), 0, 0)]
- except dns.exception.Timeout:
- log.warning("DNS resolution timed out " + \
- "for A record of %s", domain)
- v4_answers = [((domain, port), 0, 0)]
- else:
- for ans in v4_answers:
- log.debug("Found A record: %s", ans.address)
- answers.append(((ans.address, port), 0, 0))
-
- try:
- log.debug("Querying AAAA records for %s" % domain)
- v6_answers = resolver.query(domain, dns.rdatatype.AAAA)
- except (dns.resolver.NXDOMAIN, dns.resolver.NoAnswer):
- log.warning("No AAAA records for %s", domain)
- v6_answers = [((domain, port), 0, 0)]
- except dns.exception.Timeout:
- log.warning("DNS resolution timed out " + \
- "for AAAA record of %s", domain)
- v6_answers = [((domain, port), 0, 0)]
- else:
- for ans in v6_answers:
- log.debug("Found AAAA record: %s", ans.address)
- answers.append(((ans.address, port), 0, 0))
+
+ resolver = default_resolver()
+ self.configure_dns(resolver, domain=domain, port=port)
- return answers
- else:
- log.warning("dnspython is not installed -- " + \
- "relying on OS A/AAAA record resolution")
- self.configure_dns(None, domain=domain, port=port)
- return [((domain, port), 0, 0)]
+ return resolve(domain, port, service=self.dns_service, resolver=resolver)
def pick_dns_answer(self, domain, port=None):
"""Pick a server and port from DNS answers.
@@ -971,33 +918,16 @@ class XMLStream(object):
"""
if not self.dns_answers:
self.dns_answers = self.get_dns_records(domain, port)
- addresses = {}
- intmax = 0
- topprio = 65535
- for answer in self.dns_answers:
- topprio = min(topprio, answer[1])
- for answer in self.dns_answers:
- if answer[1] == topprio:
- intmax += answer[2]
- addresses[intmax] = answer[0]
-
- #python3 returns a generator for dictionary keys
- items = [x for x in addresses.keys()]
- items.sort()
-
- address = (domain, port)
- picked = random.randint(0, intmax)
- for item in items:
- if picked <= item:
- address = addresses[item]
- break
- for idx, answer in enumerate(self.dns_answers):
- if self.dns_answers[0] == address:
- self.dns_answers.pop(idx)
- break
- log.debug("Trying to connect to %s:%s", *address)
- return address
+ try:
+ if sys.version_info < (3, 0):
+ return self.dns_answers.next()
+ else:
+ return next(self.dns_answers)
+ except StopIteration:
+ self.dns_answers = None
+ return (domain, port)
+
def add_event_handler(self, name, pointer,
threaded=False, disposable=False):
"""Add a custom event handler that will be executed whenever