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

#!/usr/bin/env python2

import clientsubnetoption
import dns
import dns.message
import Queue
import os
import socket
import struct
import subprocess
import sys
import threading
import time
import unittest
import random


class DNSDistTest(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
    _testServerPort = 5350
    _toResponderQueue = Queue.Queue()
    _fromResponderQueue = Queue.Queue()
    _dnsdist = None
    _responsesCounter = {}
    _config_template = """
    newServer{address="127.0.0.1:%s"}
    truncateTC(true)
    addAnyTCRule()
    addAction(RegexRule("evil[0-9]{4,}\\\\.regex\\\\.tests\\\\.powerdns\\\\.com$"), RCodeAction(5))
    mySMN = newSuffixMatchNode()
    mySMN:add(newDNSName("nameAndQtype.tests.powerdns.com."))
    addAction(AndRule{SuffixMatchNodeRule(mySMN), QTypeRule("TXT")}, RCodeAction(4))
    block=newDNSName("powerdns.org.")
    function blockFilter(remote, qname, qtype, dh)
        if(qname:isPartOf(block))
        then
            print("Blocking *.powerdns.org")
            return true
        end
        return false
    end
    """
    _config_params = ['_testServerPort']
    _acl = ['127.0.0.1/32']

    @classmethod
    def startResponders(cls):
        print("Launching responders..")
        # clear counters
        for key in cls._responsesCounter:
            cls._responsesCounter[key] = 0

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

    @classmethod
    def startDNSDist(cls, shutUp=True):
        print("Launching dnsdist..")
        conffile = 'dnsdist_test.conf'
        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'], '-C', conffile,
                      '-l', '127.0.0.1:%d' % cls._dnsDistPort]
        for acl in cls._acl:
            dnsdistcmd.extend(['--acl', acl])
        print(' '.join(dnsdistcmd))

        if shutUp:
            with open(os.devnull, 'w') as fdDevNull:
                cls._dnsdist = subprocess.Popen(dnsdistcmd, close_fds=True, stdout=fdDevNull)
        else:
            cls._dnsdist = subprocess.Popen(dnsdistcmd, close_fds=True)

        if 'DNSDIST_FAST_TESTS' in os.environ:
            delay = 0.5
        else:
            delay = 2
        time.sleep(delay)

        if cls._dnsdist.poll() is not None:
            cls._dnsdist.kill()
            cls._dnsdist.wait()
            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.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
        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 UDPResponder(cls, port):
        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)
            request = dns.message.from_wire(data)
            if len(request.question) != 1:
                print("Skipping query with question count %d" % (len(request.question)))
                continue
            if str(request.question[0].name).endswith('tests.powerdns.com.') and not cls._toResponderQueue.empty():
                response = cls._toResponderQueue.get()
                response.id = request.id
                cls._fromResponderQueue.put(request)
                cls.ResponderIncrementCounter()
            else:
                # unexpected query, or health check
                response = dns.message.make_response(request)
                if request.question[0].rdclass == dns.rdataclass.IN:
                    if request.question[0].rdtype == dns.rdatatype.A:
                        rrset = dns.rrset.from_text(request.question[0].name,
                                                    3600,
                                                    request.question[0].rdclass,
                                                    request.question[0].rdtype,
                                                    '127.0.0.1')
                    elif request.question[0].rdtype == dns.rdatatype.AAAA:
                        rrset = dns.rrset.from_text(request.question[0].name,
                                                    3600,
                                                    request.question[0].rdclass,
                                                    request.question[0].rdtype,
                                                    '::1')
                if rrset:
                    response.answer.append(rrset)

            sock.sendto(response.to_wire(), addr)
        sock.close()

    @classmethod
    def TCPResponder(cls, port):
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        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, address) = sock.accept()
            data = conn.recv(2)
            (datalen,) = struct.unpack("!H", data)
            data = conn.recv(datalen)
            request = dns.message.from_wire(data)
            if len(request.question) != 1:
                print("Skipping query with question count %d" % (len(request.question)))
                continue
            if str(request.question[0].name).endswith('tests.powerdns.com.') and not cls._toResponderQueue.empty():
                response = cls._toResponderQueue.get()
                response.id = request.id
                cls._fromResponderQueue.put(request)
                cls.ResponderIncrementCounter()
            else:
                # unexpected query, or health check
                response = dns.message.make_response(request)
                if request.question[0].rdclass == dns.rdataclass.IN:
                    if request.question[0].rdtype == dns.rdatatype.A:
                        rrset = dns.rrset.from_text(request.question[0].name,
                                                    3600,
                                                    request.question[0].rdclass,
                                                    request.question[0].rdtype,
                                                    '127.0.0.1')
                    elif request.question[0].rdtype == dns.rdatatype.AAAA:
                        rrset = dns.rrset.from_text(request.question[0].name,
                                                    3600,
                                                    request.question[0].rdclass,
                                                    request.question[0].rdtype,
                                                    '::1')
                if rrset:
                    response.answer.append(rrset)

            wire = response.to_wire()
            conn.send(struct.pack("!H", len(wire)))
            conn.send(wire)
            conn.close()
        sock.close()

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

        if timeout:
            cls._sock.settimeout(timeout)

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

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

    @classmethod
    def sendTCPQuery(cls, query, response, useQueue=True, timeout=2.0):
        if useQueue:
            cls._toResponderQueue.put(response)
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        if timeout:
            sock.settimeout(timeout)

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

        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()

        receivedQuery = None
        message = None
        if useQueue and not cls._fromResponderQueue.empty():
            receivedQuery = cls._fromResponderQueue.get(query)
        if data:
            message = dns.message.from_wire(data)
        return (receivedQuery, message)
Commit	Line	Data
ca404e94 RG	1	#!/usr/bin/env python2
	2
	3	import clientsubnetoption
	4	import dns
b8db58a2	5	import dns.message
ca404e94 RG	6	import Queue
	7	import os
	8	import socket
	9	import struct
	10	import subprocess
	11	import sys
	12	import threading
	13	import time
	14	import unittest
18a0e7c6 CH	15	import random
18a0e7c6 CH	16
ca404e94 RG	17
	18	class DNSDistTest(unittest.TestCase):
	19	"""
	20	Set up a dnsdist instance and responder threads.
	21	Queries sent to dnsdist are relayed to the responder threads,
	22	who reply with the response provided by the tests themselves
	23	on a queue. Responder threads also queue the queries received
	24	from dnsdist on a separate queue, allowing the tests to check
	25	that the queries sent from dnsdist were as expected.
	26	"""
	27	_dnsDistPort = 5340
	28	_testServerPort = 5350
ca404e94 RG	29	_toResponderQueue = Queue.Queue()
	30	_fromResponderQueue = Queue.Queue()
	31	_dnsdist = None
ec5f5c6b	32	_responsesCounter = {}
18a0e7c6 CH	33	_config_template = """
	34	newServer{address="127.0.0.1:%s"}
	35	truncateTC(true)
	36	addAnyTCRule()
	37	addAction(RegexRule("evil[0-9]{4,}\\\\.regex\\\\.tests\\\\.powerdns\\\\.com$"), RCodeAction(5))
	38	mySMN = newSuffixMatchNode()
	39	mySMN:add(newDNSName("nameAndQtype.tests.powerdns.com."))
	40	addAction(AndRule{SuffixMatchNodeRule(mySMN), QTypeRule("TXT")}, RCodeAction(4))
	41	block=newDNSName("powerdns.org.")
	42	function blockFilter(remote, qname, qtype, dh)
	43	if(qname:isPartOf(block))
	44	then
	45	print("Blocking *.powerdns.org")
	46	return true
	47	end
	48	return false
	49	end
	50	"""
	51	_config_params = ['_testServerPort']
	52	_acl = ['127.0.0.1/32']
ca404e94 RG	53
	54	@classmethod
	55	def startResponders(cls):
	56	print("Launching responders..")
ec5f5c6b RG	57	# clear counters
	58	for key in cls._responsesCounter:
	59	cls._responsesCounter[key] = 0
	60
	61	cls._UDPResponder = threading.Thread(name='UDP Responder', target=cls.UDPResponder, args=[cls._testServerPort])
ca404e94 RG	62	cls._UDPResponder.setDaemon(True)
ca404e94 RG	63	cls._UDPResponder.start()
ec5f5c6b	64	cls._TCPResponder = threading.Thread(name='TCP Responder', target=cls.TCPResponder, args=[cls._testServerPort])
ca404e94 RG	65	cls._TCPResponder.setDaemon(True)
	66	cls._TCPResponder.start()
	67
	68	@classmethod
	69	def startDNSDist(cls, shutUp=True):
	70	print("Launching dnsdist..")
18a0e7c6 CH	71	conffile = 'dnsdist_test.conf'
	72	params = tuple([getattr(cls, param) for param in cls._config_params])
	73	print(params)
	74	with open(conffile, 'w') as conf:
	75	conf.write("-- Autogenerated by dnsdisttests.py\n")
	76	conf.write(cls._config_template % params)
	77
	78	dnsdistcmd = [os.environ['DNSDISTBIN'], '-C', conffile,
	79	'-l', '127.0.0.1:%d' % cls._dnsDistPort]
	80	for acl in cls._acl:
	81	dnsdistcmd.extend(['--acl', acl])
	82	print(' '.join(dnsdistcmd))
	83
ca404e94 RG	84	if shutUp:
ca404e94 RG	85	with open(os.devnull, 'w') as fdDevNull:
bd64cc44	86	cls._dnsdist = subprocess.Popen(dnsdistcmd, close_fds=True, stdout=fdDevNull)
ca404e94	87	else:
18a0e7c6	88	cls._dnsdist = subprocess.Popen(dnsdistcmd, close_fds=True)
ca404e94	89
0a2087eb RG	90	if 'DNSDIST_FAST_TESTS' in os.environ:
	91	delay = 0.5
	92	else:
	93	delay = 2
	94	time.sleep(delay)
ca404e94 RG	95
ca404e94 RG	96	if cls._dnsdist.poll() is not None:
0a2087eb	97	cls._dnsdist.kill()
ca404e94 RG	98	cls._dnsdist.wait()
	99	sys.exit(cls._dnsdist.returncode)
	100
	101	@classmethod
	102	def setUpSockets(cls):
	103	print("Setting up UDP socket..")
	104	cls._sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
	105	cls._sock.connect(("127.0.0.1", cls._dnsDistPort))
	106
	107	@classmethod
	108	def setUpClass(cls):
	109
	110	cls.startResponders()
	111	cls.startDNSDist()
	112	cls.setUpSockets()
	113
	114	print("Launching tests..")
	115
	116	@classmethod
	117	def tearDownClass(cls):
0a2087eb RG	118	if 'DNSDIST_FAST_TESTS' in os.environ:
	119	delay = 0.1
	120	else:
	121	delay = 1
ca404e94 RG	122	if cls._dnsdist:
ca404e94 RG	123	cls._dnsdist.terminate()
0a2087eb RG	124	if cls._dnsdist.poll() is None:
	125	time.sleep(delay)
	126	if cls._dnsdist.poll() is None:
	127	cls._dnsdist.kill()
ca404e94 RG	128	cls._dnsdist.wait()
	129
	130	@classmethod
ec5f5c6b RG	131	def ResponderIncrementCounter(cls):
	132	if threading.currentThread().name in cls._responsesCounter:
	133	cls._responsesCounter[threading.currentThread().name] += 1
	134	else:
	135	cls._responsesCounter[threading.currentThread().name] = 1
	136
	137	@classmethod
	138	def UDPResponder(cls, port):
ca404e94 RG	139	sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
ca404e94 RG	140	sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
ec5f5c6b	141	sock.bind(("127.0.0.1", port))
ca404e94 RG	142	while True:
	143	data, addr = sock.recvfrom(4096)
	144	request = dns.message.from_wire(data)
	145	if len(request.question) != 1:
	146	print("Skipping query with question count %d" % (len(request.question)))
	147	continue
	148	if str(request.question[0].name).endswith('tests.powerdns.com.') and not cls._toResponderQueue.empty():
	149	response = cls._toResponderQueue.get()
	150	response.id = request.id
	151	cls._fromResponderQueue.put(request)
ec5f5c6b	152	cls.ResponderIncrementCounter()
ca404e94 RG	153	else:
	154	# unexpected query, or health check
	155	response = dns.message.make_response(request)
87c605c4 RG	156	if request.question[0].rdclass == dns.rdataclass.IN:
	157	if request.question[0].rdtype == dns.rdatatype.A:
	158	rrset = dns.rrset.from_text(request.question[0].name,
	159	3600,
	160	request.question[0].rdclass,
	161	request.question[0].rdtype,
	162	'127.0.0.1')
	163	elif request.question[0].rdtype == dns.rdatatype.AAAA:
	164	rrset = dns.rrset.from_text(request.question[0].name,
	165	3600,
	166	request.question[0].rdclass,
	167	request.question[0].rdtype,
	168	'::1')
	169	if rrset:
	170	response.answer.append(rrset)
	171
ca404e94 RG	172	sock.sendto(response.to_wire(), addr)
	173	sock.close()
	174
	175	@classmethod
ec5f5c6b	176	def TCPResponder(cls, port):
ca404e94 RG	177	sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
	178	sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
	179	try:
ec5f5c6b	180	sock.bind(("127.0.0.1", port))
ca404e94 RG	181	except socket.error as e:
	182	print("Error binding in the TCP responder: %s" % str(e))
	183	sys.exit(1)
	184
	185	sock.listen(100)
	186	while True:
	187	(conn, address) = sock.accept()
	188	data = conn.recv(2)
	189	(datalen,) = struct.unpack("!H", data)
	190	data = conn.recv(datalen)
	191	request = dns.message.from_wire(data)
	192	if len(request.question) != 1:
	193	print("Skipping query with question count %d" % (len(request.question)))
	194	continue
	195	if str(request.question[0].name).endswith('tests.powerdns.com.') and not cls._toResponderQueue.empty():
	196	response = cls._toResponderQueue.get()
	197	response.id = request.id
	198	cls._fromResponderQueue.put(request)
ec5f5c6b	199	cls.ResponderIncrementCounter()
ca404e94 RG	200	else:
	201	# unexpected query, or health check
	202	response = dns.message.make_response(request)
87c605c4 RG	203	if request.question[0].rdclass == dns.rdataclass.IN:
	204	if request.question[0].rdtype == dns.rdatatype.A:
	205	rrset = dns.rrset.from_text(request.question[0].name,
	206	3600,
	207	request.question[0].rdclass,
	208	request.question[0].rdtype,
	209	'127.0.0.1')
	210	elif request.question[0].rdtype == dns.rdatatype.AAAA:
	211	rrset = dns.rrset.from_text(request.question[0].name,
	212	3600,
	213	request.question[0].rdclass,
	214	request.question[0].rdtype,
	215	'::1')
	216	if rrset:
	217	response.answer.append(rrset)
ca404e94 RG	218
	219	wire = response.to_wire()
	220	conn.send(struct.pack("!H", len(wire)))
	221	conn.send(wire)
	222	conn.close()
	223	sock.close()
	224
	225	@classmethod
	226	def sendUDPQuery(cls, query, response, useQueue=True, timeout=2.0):
	227	if useQueue:
	228	cls._toResponderQueue.put(response)
	229
	230	if timeout:
	231	cls._sock.settimeout(timeout)
	232
	233	try:
	234	cls._sock.send(query.to_wire())
	235	data = cls._sock.recv(4096)
	236	except socket.timeout as e:
	237	data = None
	238	finally:
	239	if timeout:
	240	cls._sock.settimeout(None)
	241
	242	receivedQuery = None
	243	message = None
	244	if useQueue and not cls._fromResponderQueue.empty():
	245	receivedQuery = cls._fromResponderQueue.get(query)
	246	if data:
	247	message = dns.message.from_wire(data)
	248	return (receivedQuery, message)
	249
	250	@classmethod
	251	def sendTCPQuery(cls, query, response, useQueue=True, timeout=2.0):
	252	if useQueue:
	253	cls._toResponderQueue.put(response)
	254	sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
ca404e94 RG	255	if timeout:
	256	sock.settimeout(timeout)
	257
0a2087eb RG	258	sock.connect(("127.0.0.1", cls._dnsDistPort))
0a2087eb RG	259
ca404e94 RG	260	try:
	261	wire = query.to_wire()
	262	sock.send(struct.pack("!H", len(wire)))
	263	sock.send(wire)
	264	data = sock.recv(2)
	265	if data:
	266	(datalen,) = struct.unpack("!H", data)
	267	data = sock.recv(datalen)
	268	except socket.timeout as e:
	269	print("Timeout: %s" % (str(e)))
	270	data = None
	271	except socket.error as e:
	272	print("Network error: %s" % (str(e)))
	273	data = None
	274	finally:
	275	sock.close()
	276
	277	receivedQuery = None
	278	message = None
	279	if useQueue and not cls._fromResponderQueue.empty():
	280	receivedQuery = cls._fromResponderQueue.get(query)
	281	if data:
	282	message = dns.message.from_wire(data)
	283	return (receivedQuery, message)