[thirdparty/pdns.git] / regression-tests.dnsdist / dnsdisttests.py

#!/usr/bin/env python2

import copy
import os
import socket
import ssl
import struct
import subprocess
import sys
import threading
import time
import unittest
import clientsubnetoption
import dns
import dns.message
import libnacl
import libnacl.utils

from eqdnsmessage import AssertEqualDNSMessageMixin

# Python2/3 compatibility hacks
try:
  from queue import Queue
except ImportError:
  from Queue import Queue

try:
  range = xrange
except NameError:
  pass


class DNSDistTest(AssertEqualDNSMessageMixin, unittest.TestCase):
    """
    Set up a dnsdist instance and responder threads.
    Queries sent to dnsdist are relayed to the responder threads,
    who reply with the response provided by the tests themselves
    on a queue. Responder threads also queue the queries received
    from dnsdist on a separate queue, allowing the tests to check
    that the queries sent from dnsdist were as expected.
    """
    _dnsDistPort = 5340
    _dnsDistListeningAddr = "127.0.0.1"
    _testServerPort = 5350
    _toResponderQueue = Queue()
    _fromResponderQueue = Queue()
    _queueTimeout = 1
    _dnsdistStartupDelay = 2.0
    _dnsdist = None
    _responsesCounter = {}
    _config_template = """
    """
    _config_params = ['_testServerPort']
    _acl = ['127.0.0.1/32']
    _consolePort = 5199
    _consoleKey = None
    _healthCheckName = 'a.root-servers.net.'
    _healthCheckCounter = 0
    _answerUnexpected = True
    _checkConfigExpectedOutput = None

    @classmethod
    def startResponders(cls):
        print("Launching responders..")

        cls._UDPResponder = threading.Thread(name='UDP Responder', target=cls.UDPResponder, args=[cls._testServerPort, cls._toResponderQueue, cls._fromResponderQueue])
        cls._UDPResponder.setDaemon(True)
        cls._UDPResponder.start()
        cls._TCPResponder = threading.Thread(name='TCP Responder', target=cls.TCPResponder, args=[cls._testServerPort, cls._toResponderQueue, cls._fromResponderQueue])
        cls._TCPResponder.setDaemon(True)
        cls._TCPResponder.start()

    @classmethod
    def startDNSDist(cls):
        print("Launching dnsdist..")
        confFile = os.path.join('configs', 'dnsdist_%s.conf' % (cls.__name__))
        params = tuple([getattr(cls, param) for param in cls._config_params])
        print(params)
        with open(confFile, 'w') as conf:
            conf.write("-- Autogenerated by dnsdisttests.py\n")
            conf.write(cls._config_template % params)

        dnsdistcmd = [os.environ['DNSDISTBIN'], '--supervised', '-C', confFile,
                      '-l', '%s:%d' % (cls._dnsDistListeningAddr, cls._dnsDistPort) ]
        for acl in cls._acl:
            dnsdistcmd.extend(['--acl', acl])
        print(' '.join(dnsdistcmd))

        # validate config with --check-config, which sets client=true, possibly exposing bugs.
        testcmd = dnsdistcmd + ['--check-config']
        try:
            output = subprocess.check_output(testcmd, stderr=subprocess.STDOUT, close_fds=True)
        except subprocess.CalledProcessError as exc:
            raise AssertionError('dnsdist --check-config failed (%d): %s' % (exc.returncode, exc.output))
        if cls._checkConfigExpectedOutput is not None:
          expectedOutput = cls._checkConfigExpectedOutput
        else:
          expectedOutput = ('Configuration \'%s\' OK!\n' % (confFile)).encode()
        if output != expectedOutput:
            raise AssertionError('dnsdist --check-config failed: %s' % output)

        logFile = os.path.join('configs', 'dnsdist_%s.log' % (cls.__name__))
        with open(logFile, 'w') as fdLog:
          cls._dnsdist = subprocess.Popen(dnsdistcmd, close_fds=True, stdout=fdLog, stderr=fdLog)

        if 'DNSDIST_FAST_TESTS' in os.environ:
            delay = 0.5
        else:
            delay = cls._dnsdistStartupDelay

        time.sleep(delay)

        if cls._dnsdist.poll() is not None:
            cls._dnsdist.kill()
            sys.exit(cls._dnsdist.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._dnsDistPort))

    @classmethod
    def setUpClass(cls):

        cls.startResponders()
        cls.startDNSDist()
        cls.setUpSockets()

        print("Launching tests..")

    @classmethod
    def tearDownClass(cls):
        if 'DNSDIST_FAST_TESTS' in os.environ:
            delay = 0.1
        else:
            delay = 1.0
        if cls._dnsdist:
            cls._dnsdist.terminate()
            if cls._dnsdist.poll() is None:
                time.sleep(delay)
                if cls._dnsdist.poll() is None:
                    cls._dnsdist.kill()
                cls._dnsdist.wait()

    @classmethod
    def _ResponderIncrementCounter(cls):
        if threading.currentThread().name in cls._responsesCounter:
            cls._responsesCounter[threading.currentThread().name] += 1
        else:
            cls._responsesCounter[threading.currentThread().name] = 1

    @classmethod
    def _getResponse(cls, request, fromQueue, toQueue, synthesize=None):
        response = None
        if len(request.question) != 1:
            print("Skipping query with question count %d" % (len(request.question)))
            return None
        healthCheck = str(request.question[0].name).endswith(cls._healthCheckName)
        if healthCheck:
            cls._healthCheckCounter += 1
            response = dns.message.make_response(request)
        else:
            cls._ResponderIncrementCounter()
            if not fromQueue.empty():
                toQueue.put(request, True, cls._queueTimeout)
                if synthesize is None:
                    response = fromQueue.get(True, cls._queueTimeout)
                    if response:
                        response = copy.copy(response)
                        response.id = request.id

        if not response:
            if synthesize is not None:
                response = dns.message.make_response(request)
                response.set_rcode(synthesize)
            elif cls._answerUnexpected:
                response = dns.message.make_response(request)
                response.set_rcode(dns.rcode.SERVFAIL)

        return response

    @classmethod
    def UDPResponder(cls, port, fromQueue, toQueue, trailingDataResponse=False, callback=None):
        # trailingDataResponse=True means "ignore trailing data".
        # Other values are either False (meaning "raise an exception")
        # or are interpreted as a response RCODE for queries with trailing data.
        # callback is invoked for every -even healthcheck ones- query and should return a raw response
        ignoreTrailing = trailingDataResponse is True

        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
        sock.bind(("127.0.0.1", port))
        while True:
            data, addr = sock.recvfrom(4096)
            forceRcode = None
            try:
                request = dns.message.from_wire(data, ignore_trailing=ignoreTrailing)
            except dns.message.TrailingJunk as e:
                if trailingDataResponse is False or forceRcode is True:
                    raise
                print("UDP query with trailing data, synthesizing response")
                request = dns.message.from_wire(data, ignore_trailing=True)
                forceRcode = trailingDataResponse

            wire = None
            if callback:
              wire = callback(request)
            else:
              response = cls._getResponse(request, fromQueue, toQueue, synthesize=forceRcode)
              if response:
                wire = response.to_wire()

            if not wire:
              continue

            sock.settimeout(2.0)
            sock.sendto(wire, addr)
            sock.settimeout(None)
        sock.close()

    @classmethod
    def TCPResponder(cls, port, fromQueue, toQueue, trailingDataResponse=False, multipleResponses=False, callback=None):
        # trailingDataResponse=True means "ignore trailing data".
        # Other values are either False (meaning "raise an exception")
        # or are interpreted as a response RCODE for queries with trailing data.
        # callback is invoked for every -even healthcheck ones- query and should return a raw response
        ignoreTrailing = trailingDataResponse is True

        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
        try:
            sock.bind(("127.0.0.1", port))
        except socket.error as e:
            print("Error binding in the TCP responder: %s" % str(e))
            sys.exit(1)

        sock.listen(100)
        while True:
            (conn, _) = sock.accept()
            conn.settimeout(5.0)
            data = conn.recv(2)
            if not data:
                conn.close()
                continue

            (datalen,) = struct.unpack("!H", data)
            data = conn.recv(datalen)
            forceRcode = None
            try:
                request = dns.message.from_wire(data, ignore_trailing=ignoreTrailing)
            except dns.message.TrailingJunk as e:
                if trailingDataResponse is False or forceRcode is True:
                    raise
                print("TCP query with trailing data, synthesizing response")
                request = dns.message.from_wire(data, ignore_trailing=True)
                forceRcode = trailingDataResponse

            if callback:
              wire = callback(request)
            else:
              response = cls._getResponse(request, fromQueue, toQueue, synthesize=forceRcode)
              if response:
                wire = response.to_wire(max_size=65535)

            if not wire:
                conn.close()
                continue

            conn.send(struct.pack("!H", len(wire)))
            conn.send(wire)

            while multipleResponses:
                if fromQueue.empty():
                    break

                response = fromQueue.get(True, cls._queueTimeout)
                if not response:
                    break

                response = copy.copy(response)
                response.id = request.id
                wire = response.to_wire(max_size=65535)
                try:
                    conn.send(struct.pack("!H", len(wire)))
                    conn.send(wire)
                except socket.error as e:
                    # some of the tests are going to close
                    # the connection on us, just deal with it
                    break

            conn.close()

        sock.close()

    @classmethod
    def sendUDPQuery(cls, query, response, useQueue=True, timeout=2.0, rawQuery=False):
        if useQueue:
            cls._toResponderQueue.put(response, True, timeout)

        if timeout:
            cls._sock.settimeout(timeout)

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

        receivedQuery = None
        message = None
        if useQueue and not cls._fromResponderQueue.empty():
            receivedQuery = cls._fromResponderQueue.get(True, timeout)
        if data:
            message = dns.message.from_wire(data)
        return (receivedQuery, message)

    @classmethod
    def openTCPConnection(cls, timeout=None):
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        if timeout:
            sock.settimeout(timeout)

        sock.connect(("127.0.0.1", cls._dnsDistPort))
        return sock

    @classmethod
    def openTLSConnection(cls, port, serverName, caCert=None, timeout=None):
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        if timeout:
            sock.settimeout(timeout)

        # 2.7.9+
        if hasattr(ssl, 'create_default_context'):
            sslctx = ssl.create_default_context(cafile=caCert)
            sslsock = sslctx.wrap_socket(sock, server_hostname=serverName)
        else:
            sslsock = ssl.wrap_socket(sock, ca_certs=caCert, cert_reqs=ssl.CERT_REQUIRED)

        sslsock.connect(("127.0.0.1", port))
        return sslsock

    @classmethod
    def sendTCPQueryOverConnection(cls, sock, query, rawQuery=False, response=None, timeout=2.0):
        if not rawQuery:
            wire = query.to_wire()
        else:
            wire = query

        if response:
            cls._toResponderQueue.put(response, True, timeout)

        sock.send(struct.pack("!H", len(wire)))
        sock.send(wire)

    @classmethod
    def recvTCPResponseOverConnection(cls, sock, useQueue=False, timeout=2.0):
        message = None
        data = sock.recv(2)
        if data:
            (datalen,) = struct.unpack("!H", data)
            data = sock.recv(datalen)
            if data:
                message = dns.message.from_wire(data)

        if useQueue and not cls._fromResponderQueue.empty():
            receivedQuery = cls._fromResponderQueue.get(True, timeout)
            return (receivedQuery, message)
        else:
            return message

    @classmethod
    def sendTCPQuery(cls, query, response, useQueue=True, timeout=2.0, rawQuery=False):
        message = None
        if useQueue:
            cls._toResponderQueue.put(response, True, timeout)

        sock = cls.openTCPConnection(timeout)

        try:
            cls.sendTCPQueryOverConnection(sock, query, rawQuery)
            message = cls.recvTCPResponseOverConnection(sock)
        except socket.timeout as e:
            print("Timeout: %s" % (str(e)))
        except socket.error as e:
            print("Network error: %s" % (str(e)))
        finally:
            sock.close()

        receivedQuery = None
        if useQueue and not cls._fromResponderQueue.empty():
            receivedQuery = cls._fromResponderQueue.get(True, timeout)

        return (receivedQuery, message)

    @classmethod
    def sendTCPQueryWithMultipleResponses(cls, query, responses, useQueue=True, timeout=2.0, rawQuery=False):
        if useQueue:
            for response in responses:
                cls._toResponderQueue.put(response, True, timeout)
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        if timeout:
            sock.settimeout(timeout)

        sock.connect(("127.0.0.1", cls._dnsDistPort))
        messages = []

        try:
            if not rawQuery:
                wire = query.to_wire()
            else:
                wire = query

            sock.send(struct.pack("!H", len(wire)))
            sock.send(wire)
            while True:
                data = sock.recv(2)
                if not data:
                    break
                (datalen,) = struct.unpack("!H", data)
                data = sock.recv(datalen)
                messages.append(dns.message.from_wire(data))

        except socket.timeout as e:
            print("Timeout: %s" % (str(e)))
        except socket.error as e:
            print("Network error: %s" % (str(e)))
        finally:
            sock.close()

        receivedQuery = None
        if useQueue and not cls._fromResponderQueue.empty():
            receivedQuery = cls._fromResponderQueue.get(True, timeout)
        return (receivedQuery, messages)

    def setUp(self):
        # This function is called before every tests

        # Clear the responses counters
        for key in self._responsesCounter:
            self._responsesCounter[key] = 0

        self._healthCheckCounter = 0

        # Make sure the queues are empty, in case
        # a previous test failed
        while not self._toResponderQueue.empty():
            self._toResponderQueue.get(False)

        while not self._fromResponderQueue.empty():
            self._fromResponderQueue.get(False)

        super(DNSDistTest, self).setUp()

    @classmethod
    def clearToResponderQueue(cls):
        while not cls._toResponderQueue.empty():
            cls._toResponderQueue.get(False)

    @classmethod
    def clearFromResponderQueue(cls):
        while not cls._fromResponderQueue.empty():
            cls._fromResponderQueue.get(False)

    @classmethod
    def clearResponderQueues(cls):
        cls.clearToResponderQueue()
        cls.clearFromResponderQueue()

    @staticmethod
    def generateConsoleKey():
        return libnacl.utils.salsa_key()

    @classmethod
    def _encryptConsole(cls, command, nonce):
        command = command.encode('UTF-8')
        if cls._consoleKey is None:
            return command
        return libnacl.crypto_secretbox(command, nonce, cls._consoleKey)

    @classmethod
    def _decryptConsole(cls, command, nonce):
        if cls._consoleKey is None:
            result = command
        else:
            result = libnacl.crypto_secretbox_open(command, nonce, cls._consoleKey)
        return result.decode('UTF-8')

    @classmethod
    def sendConsoleCommand(cls, command, timeout=1.0):
        ourNonce = libnacl.utils.rand_nonce()
        theirNonce = None
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        if timeout:
            sock.settimeout(timeout)

        sock.connect(("127.0.0.1", cls._consolePort))
        sock.send(ourNonce)
        theirNonce = sock.recv(len(ourNonce))
        if len(theirNonce) != len(ourNonce):
            print("Received a nonce of size %d, expecting %d, console command will not be sent!" % (len(theirNonce), len(ourNonce)))
            if len(theirNonce) == 0:
                raise socket.error("Got EOF while reading a nonce of size %d, console command will not be sent!" % (len(ourNonce)))
            return None

        halfNonceSize = int(len(ourNonce) / 2)
        readingNonce = ourNonce[0:halfNonceSize] + theirNonce[halfNonceSize:]
        writingNonce = theirNonce[0:halfNonceSize] + ourNonce[halfNonceSize:]
        msg = cls._encryptConsole(command, writingNonce)
        sock.send(struct.pack("!I", len(msg)))
        sock.send(msg)
        data = sock.recv(4)
        if not data:
            raise socket.error("Got EOF while reading the response size")

        (responseLen,) = struct.unpack("!I", data)
        data = sock.recv(responseLen)
        response = cls._decryptConsole(data, readingNonce)
        return response

    def compareOptions(self, a, b):
        self.assertEquals(len(a), len(b))
        for idx in range(len(a)):
            self.assertEquals(a[idx], b[idx])

    def checkMessageNoEDNS(self, expected, received):
        self.assertEquals(expected, received)
        self.assertEquals(received.edns, -1)
        self.assertEquals(len(received.options), 0)

    def checkMessageEDNSWithoutOptions(self, expected, received):
        self.assertEquals(expected, received)
        self.assertEquals(received.edns, 0)
        self.assertEquals(expected.payload, received.payload)

    def checkMessageEDNSWithoutECS(self, expected, received, withCookies=0):
        self.assertEquals(expected, received)
        self.assertEquals(received.edns, 0)
        self.assertEquals(expected.payload, received.payload)
        self.assertEquals(len(received.options), withCookies)
        if withCookies:
            for option in received.options:
                self.assertEquals(option.otype, 10)

    def checkMessageEDNSWithECS(self, expected, received, additionalOptions=0):
        self.assertEquals(expected, received)
        self.assertEquals(received.edns, 0)
        self.assertEquals(expected.payload, received.payload)
        self.assertEquals(len(received.options), 1 + additionalOptions)
        hasECS = False
        for option in received.options:
            if option.otype == clientsubnetoption.ASSIGNED_OPTION_CODE:
                hasECS = True
            else:
                self.assertNotEquals(additionalOptions, 0)

        self.compareOptions(expected.options, received.options)
        self.assertTrue(hasECS)

    def checkQueryEDNSWithECS(self, expected, received, additionalOptions=0):
        self.checkMessageEDNSWithECS(expected, received, additionalOptions)

    def checkResponseEDNSWithECS(self, expected, received, additionalOptions=0):
        self.checkMessageEDNSWithECS(expected, received, additionalOptions)

    def checkQueryEDNSWithoutECS(self, expected, received):
        self.checkMessageEDNSWithoutECS(expected, received)

    def checkResponseEDNSWithoutECS(self, expected, received, withCookies=0):
        self.checkMessageEDNSWithoutECS(expected, received, withCookies)

    def checkQueryNoEDNS(self, expected, received):
        self.checkMessageNoEDNS(expected, received)

    def checkResponseNoEDNS(self, expected, received):
        self.checkMessageNoEDNS(expected, received)