Limit the number of queries sent out to get NS addresses per query.

[thirdparty/pdns.git] / regression-tests.ixfrdist / ixfrdisttests.py

#!/usr/bin/env python2

import errno
import shutil
import os
import socket
import struct
import subprocess
import sys
import time
import unittest
import dns
import dns.message

from eqdnsmessage import AssertEqualDNSMessageMixin

class IXFRDistTest(AssertEqualDNSMessageMixin, unittest.TestCase):

    _ixfrDistStartupDelay = 2.0
    _ixfrDistPort = 5342

    _config_template = """
listen:
  - '127.0.0.1:%d'
acl:
  - '127.0.0.0/8'
axfr-timeout: 20
keep: 20
tcp-in-threads: 10
work-dir: 'ixfrdist.dir'
failed-soa-retry: 3
"""
    _config_domains = None
    _config_params = ['_ixfrDistPort']

    @classmethod
    def startIXFRDist(cls):
        print("Launching ixfrdist..")
        conffile = 'ixfrdist.yml'
        params = tuple([getattr(cls, param) for param in cls._config_params])
        print(params)
        with open(conffile, 'w') as conf:
            conf.write("# Autogenerated by ixfrdisttests.py\n")
            conf.write(cls._config_template % params)

            if cls._config_domains is not None:
                conf.write("domains:\n")

                for domain, master in cls._config_domains.items():
                    conf.write("  - domain: %s\n" % (domain))
                    conf.write("    master: %s\n" % (master))

        ixfrdistcmd = [os.environ['IXFRDISTBIN'], '--config', conffile, '--debug']

        logFile = 'ixfrdist.log'
        with open(logFile, 'w') as fdLog:
            cls._ixfrdist = subprocess.Popen(ixfrdistcmd, close_fds=True,
                                             stdout=fdLog, stderr=fdLog)

        if 'IXFRDIST_FAST_TESTS' in os.environ:
            delay = 0.5
        else:
            delay = cls._ixfrDistStartupDelay

        time.sleep(delay)

        if cls._ixfrdist.poll() is not None:
            cls._ixfrdist.kill()
            sys.exit(cls._ixfrdist.returncode)

    @classmethod
    def setUpSockets(cls):
        print("Setting up UDP socket..")
        cls._sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        cls._sock.settimeout(2.0)
        cls._sock.connect(("127.0.0.1", cls._ixfrDistPort))

    @classmethod
    def setUpClass(cls):
        cls.startIXFRDist()
        cls.setUpSockets()

        print("Launching tests..")

    @classmethod
    def tearDownClass(cls):
        cls.tearDownIXFRDist()

    @classmethod
    def tearDownIXFRDist(cls):
        if 'IXFRDIST_FAST_TESTS' in os.environ:
            delay = 0.1
        else:
            delay = 1.0

        try:
            if cls._ixfrdist:
                cls._ixfrdist.terminate()
                if cls._ixfrdist.poll() is None:
                    time.sleep(delay)
                    if cls._ixfrdist.poll() is None:
                        cls._ixfrdist.kill()
                    cls._ixfrdist.wait()
        except OSError as e:
            # There is a race-condition with the poll() and
            # kill() statements, when the process is dead on the
            # kill(), this is fine
            if e.errno != errno.ESRCH:
                raise

    @classmethod
    def sendUDPQuery(cls, query, timeout=2.0, decode=True, fwparams=dict()):
        if timeout:
            cls._sock.settimeout(timeout)

        try:
            cls._sock.send(query.to_wire())
            data = cls._sock.recv(4096)
        except socket.timeout:
            data = None
        finally:
            if timeout:
                cls._sock.settimeout(None)

        message = None
        if data:
            if not decode:
                return data
            message = dns.message.from_wire(data, **fwparams)
        return message

    # FIXME: sendTCPQuery and sendTCPQueryMultiResponse, when they are done reading
    # should wait for a short while on the socket to see if more data is coming
    # and error if it does!
    @classmethod
    def sendTCPQuery(cls, query, timeout=2.0):
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        if timeout:
            sock.settimeout(timeout)

        sock.connect(("127.0.0.1", cls._ixfrDistPort))

        try:
            wire = query.to_wire()
            sock.send(struct.pack("!H", len(wire)))
            sock.send(wire)
            data = sock.recv(2)
            if data:
                (datalen,) = struct.unpack("!H", data)
                data = sock.recv(datalen)
        except socket.timeout as e:
            print("Timeout: %s" % (str(e)))
            data = None
        except socket.error as e:
            print("Network error: %s" % (str(e)))
            data = None
        finally:
            sock.close()

        message = None
        if data:
            message = dns.message.from_wire(data)
        return message

    @classmethod
    def sendTCPQueryMultiResponse(cls, query, timeout=2.0, count=1):
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        if timeout:
            sock.settimeout(timeout)

        sock.connect(("127.0.0.1", cls._ixfrDistPort))

        try:
            wire = query.to_wire()
            sock.send(struct.pack("!H", len(wire)))
            sock.send(wire)
        except socket.timeout as e:
            raise Exception("Timeout: %s" % (str(e)))
        except socket.error as e:
            raise Exception("Network error: %s" % (str(e)))

        messages = []
        for i in range(count):
            try:
                data = sock.recv(2)
                if data:
                    (datalen,) = struct.unpack("!H", data)
                    data = sock.recv(datalen)
                    messages.append(dns.message.from_wire(data))
                else:
                    break
            except socket.timeout as e:
                raise Exception("Timeout: %s" % (str(e)))
            except socket.error as e:
                raise Exception("Network error: %s" % (str(e)))

        return messages

    def setUp(self):
        # This function is called before every tests
        super(IXFRDistTest, self).setUp()