Merge pull request #7909 from qvr/expungebyname-stats

[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

class IXFRDistTest(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
        return