regression-tests.recursor-dnssec/test_Protobuf.py

   1 import dns
   2 import dnsmessage_pb2
   3 import os
   4 import socket
   5 import struct
   6 import sys
   7 import threading
   8 import time
   9
  10 # Python2/3 compatibility hacks
  11 try:
  12   from queue import Queue
  13 except ImportError:
  14   from Queue import Queue
  15
  16 try:
  17   range = xrange
  18 except NameError:
  19   pass
  20
  21 from recursortests import RecursorTest
  22
  23 def ProtobufConnectionHandler(queue, conn):
  24     data = None
  25     while True:
  26         data = conn.recv(2)
  27         if not data:
  28             break
  29         (datalen,) = struct.unpack("!H", data)
  30         data = conn.recv(datalen)
  31         if not data:
  32             break
  33
  34         queue.put(data, True, timeout=2.0)
  35
  36     conn.close()
  37
  38 def ProtobufListener(queue, port):
  39     sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  40     sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
  41     try:
  42         sock.bind(("127.0.0.1", port))
  43     except socket.error as e:
  44         print("Error binding in the protobuf listener: %s" % str(e))
  45         sys.exit(1)
  46
  47     sock.listen(100)
  48     while True:
  49         try:
  50             (conn, _) = sock.accept()
  51             thread = threading.Thread(name='Connection Handler',
  52                                       target=ProtobufConnectionHandler,
  53                                       args=[queue, conn])
  54             thread.setDaemon(True)
  55             thread.start()
  56
  57         except socket.error as e:
  58             print('Error in protobuf socket: %s' % str(e))
  59
  60     sock.close()
  61
  62
  63 class ProtobufServerParams:
  64   def __init__(self, port):
  65     self.queue = Queue()
  66     self.port = port
  67
  68 protobufServersParameters = [ProtobufServerParams(4243), ProtobufServerParams(4244)]
  69 protobufListeners = []
  70 for param in protobufServersParameters:
  71   listener = threading.Thread(name='Protobuf Listener', target=ProtobufListener, args=[param.queue, param.port])
  72   listener.setDaemon(True)
  73   listener.start()
  74   protobufListeners.append(listener)
  75
  76 class TestRecursorProtobuf(RecursorTest):
  77
  78     _lua_config_file = """
  79     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"})
  80     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
  81
  82     _auth_zones = {
  83         '8': {'threads': 1,
  84               'zones': ['ROOT']},
  85         '9': {'threads': 1,
  86               'zones': ['secure.example', 'islandofsecurity.example']},
  87         '10': {'threads': 1,
  88                'zones': ['example']},
  89         '18': {'threads': 1,
  90                'zones': ['example']}
  91     }
  92
  93     def getFirstProtobufMessage(self, retries=1, waitTime=1):
  94         msg = None
  95
  96         print("in getFirstProtobufMessage")
  97         for param in protobufServersParameters:
  98           print(param.port)
  99           failed = 0
 100
 101           while param.queue.empty:
 102             print(failed)
 103             print(retries)
 104             if failed >= retries:
 105               break
 106
 107             failed = failed + 1
 108             print("waiting")
 109             time.sleep(waitTime)
 110
 111           self.assertFalse(param.queue.empty())
 112           data = param.queue.get(False)
 113           self.assertTrue(data)
 114           oldmsg = msg
 115           msg = dnsmessage_pb2.PBDNSMessage()
 116           msg.ParseFromString(data)
 117           if oldmsg is not None:
 118             self.assertEqual(msg, oldmsg)
 119
 120         print(msg)
 121         return msg
 122
 123     def checkNoRemainingMessage(self):
 124         for param in protobufServersParameters:
 125           self.assertTrue(param.queue.empty())
 126
 127     def checkProtobufBase(self, msg, protocol, query, initiator, normalQueryResponse=True, expectedECS=None, receivedSize=None):
 128         self.assertTrue(msg)
 129         self.assertTrue(msg.HasField('timeSec'))
 130         self.assertTrue(msg.HasField('socketFamily'))
 131         self.assertEqual(msg.socketFamily, dnsmessage_pb2.PBDNSMessage.INET)
 132         self.assertTrue(msg.HasField('from'))
 133         fromvalue = getattr(msg, 'from')
 134         self.assertEqual(socket.inet_ntop(socket.AF_INET, fromvalue), initiator)
 135         self.assertTrue(msg.HasField('socketProtocol'))
 136         self.assertEqual(msg.socketProtocol, protocol)
 137         self.assertTrue(msg.HasField('messageId'))
 138         self.assertTrue(msg.HasField('serverIdentity'))
 139         self.assertTrue(msg.HasField('id'))
 140         self.assertEqual(msg.id, query.id)
 141         self.assertTrue(msg.HasField('inBytes'))
 142         if normalQueryResponse:
 143             # compare inBytes with length of query/response
 144             # Note that for responses, the size we received might differ
 145             # because dnspython might compress labels differently from
 146             # the recursor
 147             if receivedSize:
 148                 self.assertEqual(msg.inBytes, receivedSize)
 149             else:
 150                 self.assertEqual(msg.inBytes, len(query.to_wire()))
 151         if expectedECS is not None:
 152             self.assertTrue(msg.HasField('originalRequestorSubnet'))
 153             # v4 only for now
 154             self.assertEqual(len(msg.originalRequestorSubnet), 4)
 155             self.assertEqual(socket.inet_ntop(socket.AF_INET, msg.originalRequestorSubnet), '127.0.0.1')
 156
 157     def checkOutgoingProtobufBase(self, msg, protocol, query, initiator, length=None):
 158         self.assertTrue(msg)
 159         self.assertTrue(msg.HasField('timeSec'))
 160         self.assertTrue(msg.HasField('socketFamily'))
 161         self.assertEqual(msg.socketFamily, dnsmessage_pb2.PBDNSMessage.INET)
 162         self.assertTrue(msg.HasField('socketProtocol'))
 163         self.assertEqual(msg.socketProtocol, protocol)
 164         self.assertTrue(msg.HasField('messageId'))
 165         self.assertTrue(msg.HasField('serverIdentity'))
 166         self.assertTrue(msg.HasField('id'))
 167         self.assertNotEqual(msg.id, query.id)
 168         self.assertTrue(msg.HasField('inBytes'))
 169         if length is not None:
 170           self.assertEqual(msg.inBytes, length)
 171         else:
 172           # compare inBytes with length of query/response
 173           self.assertEqual(msg.inBytes, len(query.to_wire()))
 174
 175     def checkProtobufQuery(self, msg, protocol, query, qclass, qtype, qname, initiator='127.0.0.1', to='127.0.0.1'):
 176         self.assertEqual(msg.type, dnsmessage_pb2.PBDNSMessage.DNSQueryType)
 177         self.checkProtobufBase(msg, protocol, query, initiator)
 178         # dnsdist doesn't fill the responder field for responses
 179         # because it doesn't keep the information around.
 180         self.assertTrue(msg.HasField('to'))
 181         self.assertEqual(socket.inet_ntop(socket.AF_INET, msg.to), to)
 182         self.assertTrue(msg.HasField('question'))
 183         self.assertTrue(msg.question.HasField('qClass'))
 184         self.assertEqual(msg.question.qClass, qclass)
 185         self.assertTrue(msg.question.HasField('qType'))
 186         self.assertEqual(msg.question.qClass, qtype)
 187         self.assertTrue(msg.question.HasField('qName'))
 188         self.assertEqual(msg.question.qName, qname)
 189
 190     def checkProtobufResponse(self, msg, protocol, response, initiator='127.0.0.1', receivedSize=None, vstate=dnsmessage_pb2.PBDNSMessage.VState.Indeterminate):
 191         self.assertEqual(msg.type, dnsmessage_pb2.PBDNSMessage.DNSResponseType)
 192         self.checkProtobufBase(msg, protocol, response, initiator, receivedSize=receivedSize)
 193         self.assertTrue(msg.HasField('response'))
 194         self.assertTrue(msg.response.HasField('queryTimeSec'))
 195         self.assertTrue(msg.response.HasField('validationState'))
 196         self.assertEqual(msg.response.validationState, vstate)
 197
 198     def checkProtobufResponseRecord(self, record, rclass, rtype, rname, rttl, checkTTL=True):
 199         self.assertTrue(record.HasField('class'))
 200         self.assertEqual(getattr(record, 'class'), rclass)
 201         self.assertTrue(record.HasField('type'))
 202         self.assertEqual(record.type, rtype)
 203         self.assertTrue(record.HasField('name'))
 204         self.assertEqual(record.name, rname)
 205         self.assertTrue(record.HasField('ttl'))
 206         if checkTTL:
 207             self.assertEqual(record.ttl, rttl)
 208         self.assertTrue(record.HasField('rdata'))
 209
 210     def checkProtobufPolicy(self, msg, policyType, reason, trigger, hit, kind):
 211         self.assertEqual(msg.type, dnsmessage_pb2.PBDNSMessage.DNSResponseType)
 212         self.assertTrue(msg.response.HasField('appliedPolicyType'))
 213         self.assertTrue(msg.response.HasField('appliedPolicy'))
 214         self.assertTrue(msg.response.HasField('appliedPolicyTrigger'))
 215         self.assertTrue(msg.response.HasField('appliedPolicyHit'))
 216         self.assertTrue(msg.response.HasField('appliedPolicyKind'))
 217         self.assertEqual(msg.response.appliedPolicy, reason)
 218         self.assertEqual(msg.response.appliedPolicyType, policyType)
 219         self.assertEqual(msg.response.appliedPolicyTrigger, trigger)
 220         self.assertEqual(msg.response.appliedPolicyHit, hit)
 221         self.assertEqual(msg.response.appliedPolicyKind, kind)
 222
 223     def checkProtobufTags(self, msg, tags):
 224         print(tags)
 225         print('---')
 226         print(msg.response.tags)
 227         self.assertEqual(len(msg.response.tags), len(tags))
 228         for tag in msg.response.tags:
 229             self.assertTrue(tag in tags)
 230
 231     def checkProtobufMetas(self, msg, metas):
 232         print(metas)
 233         print('---')
 234         print(msg.meta)
 235         self.assertEqual(len(msg.meta), len(metas))
 236         for m in msg.meta:
 237             self.assertTrue(m.HasField('key'))
 238             self.assertTrue(m.HasField('value'))
 239             self.assertTrue(m.key in metas)
 240             for i in m.value.intVal :
 241               self.assertTrue(i in metas[m.key]['intVal'])
 242             for s in m.value.stringVal :
 243               self.assertTrue(s in metas[m.key]['stringVal'])
 244
 245     def checkProtobufOutgoingQuery(self, msg, protocol, query, qclass, qtype, qname, initiator='127.0.0.1', length=None):
 246         self.assertEqual(msg.type, dnsmessage_pb2.PBDNSMessage.DNSOutgoingQueryType)
 247         self.checkOutgoingProtobufBase(msg, protocol, query, initiator, length=length)
 248         self.assertTrue(msg.HasField('to'))
 249         self.assertTrue(msg.HasField('question'))
 250         self.assertTrue(msg.question.HasField('qClass'))
 251         self.assertEqual(msg.question.qClass, qclass)
 252         self.assertTrue(msg.question.HasField('qType'))
 253         self.assertEqual(msg.question.qType, qtype)
 254         self.assertTrue(msg.question.HasField('qName'))
 255         self.assertEqual(msg.question.qName, qname)
 256
 257     def checkProtobufIncomingResponse(self, msg, protocol, response, initiator='127.0.0.1', length=None):
 258         self.assertEqual(msg.type, dnsmessage_pb2.PBDNSMessage.DNSIncomingResponseType)
 259         self.checkOutgoingProtobufBase(msg, protocol, response, initiator, length=length)
 260         self.assertTrue(msg.HasField('response'))
 261         self.assertTrue(msg.response.HasField('rcode'))
 262         self.assertTrue(msg.response.HasField('queryTimeSec'))
 263
 264     def checkProtobufIncomingNetworkErrorResponse(self, msg, protocol, response, initiator='127.0.0.1'):
 265         self.checkProtobufIncomingResponse(msg, protocol, response, initiator, length=0)
 266         self.assertEqual(msg.response.rcode, 65536)
 267
 268     def checkProtobufIdentity(self, msg, requestorId, deviceId, deviceName):
 269         print(msg)
 270         self.assertTrue((requestorId == '') == (not msg.HasField('requestorId')))
 271         self.assertTrue((deviceId == b'') == (not msg.HasField('deviceId')))
 272         self.assertTrue((deviceName == '') == (not msg.HasField('deviceName')))
 273         self.assertEqual(msg.requestorId, requestorId)
 274         self.assertEqual(msg.deviceId, deviceId)
 275         self.assertEqual(msg.deviceName, deviceName)
 276
 277     def setUp(self):
 278         super(TestRecursorProtobuf, self).setUp()
 279         # Make sure the queue is empty, in case
 280         # a previous test failed
 281         for param in protobufServersParameters:
 282             while not param.queue.empty():
 283                 param.queue.get(False)
 284         # wait long enough to be sure that the housekeeping has
 285         # prime the root NS
 286         time.sleep(1)
 287
 288     @classmethod
 289     def generateRecursorConfig(cls, confdir):
 290         authzonepath = os.path.join(confdir, 'example.zone')
 291         with open(authzonepath, 'w') as authzone:
 292             authzone.write("""$ORIGIN example.
 293 @ 3600 IN SOA {soa}
 294 a 3600 IN A 192.0.2.42
 295 tagged 3600 IN A 192.0.2.84
 296 meta 3600 IN A 192.0.2.85
 297 query-selected 3600 IN A 192.0.2.84
 298 answer-selected 3600 IN A 192.0.2.84
 299 types 3600 IN A 192.0.2.84
 300 types 3600 IN AAAA 2001:DB8::1
 301 types 3600 IN TXT "Lorem ipsum dolor sit amet"
 302 types 3600 IN MX 10 a.example.
 303 types 3600 IN SPF "v=spf1 -all"
 304 types 3600 IN SRV 10 20 443 a.example.
 305 cname 3600 IN CNAME a.example.
 306
 307 """.format(soa=cls._SOA))
 308         super(TestRecursorProtobuf, cls).generateRecursorConfig(confdir)
 309
 310
 311 class ProtobufDefaultTest(TestRecursorProtobuf):
 312     """
 313     This test makes sure that we correctly export queries and response over protobuf.
 314     """
 315
 316     _confdir = 'ProtobufDefault'
 317     _config_template = """
 318 auth-zones=example=configs/%s/example.zone""" % _confdir
 319
 320     def testA(self):
 321         name = 'a.example.'
 322         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 323         query = dns.message.make_query(name, 'A', want_dnssec=True)
 324         query.flags |= dns.flags.CD
 325         res = self.sendUDPQuery(query)
 326
 327         self.assertRRsetInAnswer(res, expected)
 328
 329         # check the protobuf messages corresponding to the UDP query and answer
 330         msg = self.getFirstProtobufMessage()
 331         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
 332         # then the response
 333         msg = self.getFirstProtobufMessage()
 334         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res, '127.0.0.1')
 335         self.assertEqual(len(msg.response.rrs), 1)
 336         rr = msg.response.rrs[0]
 337         # we have max-cache-ttl set to 15
 338         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 339         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
 340         self.checkNoRemainingMessage()
 341
 342     def testCNAME(self):
 343         name = 'cname.example.'
 344         expectedCNAME = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'CNAME', 'a.example.')
 345         expectedA = dns.rrset.from_text('a.example.', 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 346         query = dns.message.make_query(name, 'A', want_dnssec=True)
 347         query.flags |= dns.flags.CD
 348         raw = self.sendUDPQuery(query, decode=False)
 349         res = dns.message.from_wire(raw)
 350         self.assertRRsetInAnswer(res, expectedCNAME)
 351         self.assertRRsetInAnswer(res, expectedA)
 352
 353         # check the protobuf messages corresponding to the UDP query and answer
 354         # but first let the protobuf messages the time to get there
 355         msg = self.getFirstProtobufMessage()
 356         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
 357         # then the response
 358         msg = self.getFirstProtobufMessage()
 359         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res, '127.0.0.1', receivedSize=len(raw))
 360         self.assertEqual(len(msg.response.rrs), 2)
 361         rr = msg.response.rrs[0]
 362         # we don't want to check the TTL for the A record, it has been cached by the previous test
 363         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.CNAME, name, 15)
 364         self.assertEqual(rr.rdata, b'a.example.')
 365         rr = msg.response.rrs[1]
 366         # we have max-cache-ttl set to 15
 367         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, 'a.example.', 15, checkTTL=False)
 368         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
 369         self.checkNoRemainingMessage()
 370
 371 class ProtobufProxyTest(TestRecursorProtobuf):
 372     """
 373     This test makes sure that we correctly export addresses over protobuf when the proxy protocol is used.
 374     """
 375
 376     _confdir = 'ProtobufProxy'
 377     _config_template = """
 378 auth-zones=example=configs/%s/example.zone
 379 proxy-protocol-from=127.0.0.1/32
 380 allow-from=127.0.0.1,6.6.6.6
 381 """ % _confdir
 382
 383     def testA(self):
 384         name = 'a.example.'
 385         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 386         query = dns.message.make_query(name, 'A', want_dnssec=True)
 387         query.flags |= dns.flags.CD
 388         res = self.sendUDPQueryWithProxyProtocol(query, False, '6.6.6.6', '7.7.7.7', 666, 777)
 389
 390         self.assertRRsetInAnswer(res, expected)
 391
 392         # check the protobuf messages corresponding to the UDP query and answer
 393         msg = self.getFirstProtobufMessage()
 394         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name, '6.6.6.6', '7.7.7.7')
 395         # then the response
 396         msg = self.getFirstProtobufMessage()
 397         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res, '6.6.6.6')
 398         self.assertEqual(len(msg.response.rrs), 1)
 399         rr = msg.response.rrs[0]
 400         # we have max-cache-ttl set to 15
 401         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 402         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
 403         self.checkNoRemainingMessage()
 404
 405 class OutgoingProtobufDefaultTest(TestRecursorProtobuf):
 406     """
 407     This test makes sure that we correctly export outgoing queries over protobuf.
 408     It must be improved and setup env so we can check for incoming responses, but makes sure for now
 409     that the recursor at least connects to the protobuf server.
 410     """
 411
 412     _confdir = 'OutgoingProtobufDefault'
 413     _config_template = """
 414     # Switch off QName Minimization, it generates much more protobuf messages
 415     # (or make the test much more smart!)
 416     qname-minimization=no
 417 """
 418     _lua_config_file = """
 419     outgoingProtobufServer({"127.0.0.1:%d", "127.0.0.1:%d"})
 420     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 421
 422     def testA(self):
 423         name = 'host1.secure.example.'
 424         expected = list()
 425
 426         # the root DNSKEY has been learned with priming the root NS already
 427         # ('.', dns.rdatatype.DNSKEY, dnsmessage_pb2.PBDNSMessage.UDP, 201),
 428         for qname, qtype, proto, responseSize in [
 429                 ('host1.secure.example.', dns.rdatatype.A, dnsmessage_pb2.PBDNSMessage.UDP, 248),
 430                 ('host1.secure.example.', dns.rdatatype.A, dnsmessage_pb2.PBDNSMessage.UDP, 221),
 431                 ('example.', dns.rdatatype.DNSKEY, dnsmessage_pb2.PBDNSMessage.UDP, 219),
 432                 ('host1.secure.example.', dns.rdatatype.A, dnsmessage_pb2.PBDNSMessage.UDP, 175),
 433                 ('secure.example.', dns.rdatatype.DNSKEY, dnsmessage_pb2.PBDNSMessage.UDP, 233),
 434         ]:
 435             if not qname:
 436                 expected.append((None, None, None, None, None, None))
 437                 continue
 438             query = dns.message.make_query(qname, qtype, use_edns=True, want_dnssec=True)
 439             resp = dns.message.make_response(query)
 440             expected.append((
 441                 qname, qtype, query, resp, proto, responseSize
 442             ))
 443
 444         query = dns.message.make_query(name, 'A', want_dnssec=True)
 445         query.flags |= dns.flags.RD
 446         res = self.sendUDPQuery(query)
 447
 448         for qname, qtype, qry, ans, proto, responseSize in expected:
 449             if not qname:
 450                 self.getFirstProtobufMessage()
 451                 self.getFirstProtobufMessage()
 452                 continue
 453
 454             msg = self.getFirstProtobufMessage()
 455             self.checkProtobufOutgoingQuery(msg, proto, qry, dns.rdataclass.IN, qtype, qname)
 456
 457             # Check the answer
 458             msg = self.getFirstProtobufMessage()
 459             self.checkProtobufIncomingResponse(msg, proto, ans, length=responseSize)
 460
 461         self.checkNoRemainingMessage()
 462
 463 class OutgoingProtobufNoQueriesTest(TestRecursorProtobuf):
 464     """
 465     This test makes sure that we correctly export incoming responses but not outgoing queries over protobuf.
 466     It must be improved and setup env so we can check for incoming responses, but makes sure for now
 467     that the recursor at least connects to the protobuf server.
 468     """
 469
 470     _confdir = 'OutgoingProtobufNoQueries'
 471     _config_template = """
 472     # Switch off QName Minimization, it generates much more protobuf messages
 473     # (or make the test much more smart!)
 474     qname-minimization=no"""
 475     _lua_config_file = """
 476     outgoingProtobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=false, logResponses=true })
 477     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 478
 479     def testA(self):
 480         name = 'host1.secure.example.'
 481         expected = list()
 482         # the root DNSKEY has been learned with priming the root NS already
 483         # ('.', dns.rdatatype.DNSKEY, dnsmessage_pb2.PBDNSMessage.UDP, 201),
 484         for qname, qtype, proto, size in [
 485                 ('host1.secure.example.', dns.rdatatype.A, dnsmessage_pb2.PBDNSMessage.UDP, 248),
 486                 ('host1.secure.example.', dns.rdatatype.A, dnsmessage_pb2.PBDNSMessage.UDP, 221),
 487                 ('example.', dns.rdatatype.DNSKEY, dnsmessage_pb2.PBDNSMessage.UDP, 219),
 488                 ('host1.secure.example.', dns.rdatatype.A, dnsmessage_pb2.PBDNSMessage.UDP, 175),
 489                 ('secure.example.', dns.rdatatype.DNSKEY, dnsmessage_pb2.PBDNSMessage.UDP, 233),
 490         ]:
 491             if not qname:
 492                 expected.append((None, None, None, None, None, None))
 493                 continue
 494             query = dns.message.make_query(qname, qtype, use_edns=True, want_dnssec=True)
 495             resp = dns.message.make_response(query)
 496             expected.append((
 497                 qname, qtype, query, resp, proto, size
 498             ))
 499
 500         query = dns.message.make_query(name, 'A', want_dnssec=True)
 501         query.flags |= dns.flags.RD
 502         res = self.sendUDPQuery(query)
 503
 504         for qname, qtype, qry, ans, proto, size in expected:
 505             if not qname:
 506                 self.getFirstProtobufMessage()
 507                 continue
 508
 509             # check the response
 510             msg = self.getFirstProtobufMessage()
 511             self.checkProtobufIncomingResponse(msg, proto, ans, length=size)
 512
 513         self.checkNoRemainingMessage()
 514
 515 class ProtobufMasksTest(TestRecursorProtobuf):
 516     """
 517     This test makes sure that we correctly export queries and response over protobuf, respecting the configured initiator masking.
 518     """
 519
 520     _confdir = 'ProtobufMasks'
 521     _config_template = """
 522 auth-zones=example=configs/%s/example.zone""" % _confdir
 523     _protobufMaskV4 = 4
 524     _protobufMaskV6 = 128
 525     _lua_config_file = """
 526     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"})
 527     setProtobufMasks(%d, %d)
 528     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port, _protobufMaskV4, _protobufMaskV6)
 529
 530     def testA(self):
 531         name = 'a.example.'
 532         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 533         query = dns.message.make_query(name, 'A', want_dnssec=True)
 534         query.flags |= dns.flags.CD
 535         res = self.sendUDPQuery(query)
 536         self.assertRRsetInAnswer(res, expected)
 537
 538         # check the protobuf messages corresponding to the UDP query and answer
 539         # but first let the protobuf messages the time to get there
 540         msg = self.getFirstProtobufMessage()
 541         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name, '112.0.0.0')
 542         # then the response
 543         msg = self.getFirstProtobufMessage()
 544         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res, '112.0.0.0')
 545         self.assertEqual(len(msg.response.rrs), 1)
 546         rr = msg.response.rrs[0]
 547         # we have max-cache-ttl set to 15
 548         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 549         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
 550         self.checkNoRemainingMessage()
 551
 552 class ProtobufQueriesOnlyTest(TestRecursorProtobuf):
 553     """
 554     This test makes sure that we correctly export queries but not responses over protobuf.
 555     """
 556
 557     _confdir = 'ProtobufQueriesOnly'
 558     _config_template = """
 559 auth-zones=example=configs/%s/example.zone""" % _confdir
 560     _lua_config_file = """
 561     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=true, logResponses=false } )
 562     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 563
 564     def testA(self):
 565         name = 'a.example.'
 566         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 567         query = dns.message.make_query(name, 'A', want_dnssec=True)
 568         query.flags |= dns.flags.CD
 569         res = self.sendUDPQuery(query)
 570         self.assertRRsetInAnswer(res, expected)
 571
 572         # check the protobuf message corresponding to the UDP query
 573         msg = self.getFirstProtobufMessage()
 574         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
 575         # no response
 576         self.checkNoRemainingMessage()
 577
 578 class ProtobufResponsesOnlyTest(TestRecursorProtobuf):
 579     """
 580     This test makes sure that we correctly export responses but not queries over protobuf.
 581     """
 582
 583     _confdir = 'ProtobufResponsesOnly'
 584     _config_template = """
 585 auth-zones=example=configs/%s/example.zone""" % _confdir
 586     _lua_config_file = """
 587     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=false, logResponses=true } )
 588     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 589
 590     def testA(self):
 591         name = 'a.example.'
 592         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 593         query = dns.message.make_query(name, 'A', want_dnssec=True)
 594         query.flags |= dns.flags.CD
 595         res = self.sendUDPQuery(query)
 596         self.assertRRsetInAnswer(res, expected)
 597
 598         # check the protobuf message corresponding to the UDP response
 599         msg = self.getFirstProtobufMessage()
 600         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res)
 601         self.assertEqual(len(msg.response.rrs), 1)
 602         rr = msg.response.rrs[0]
 603         # we have max-cache-ttl set to 15
 604         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 605         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
 606         # nothing else in the queue
 607         self.checkNoRemainingMessage()
 608
 609 class ProtobufTaggedOnlyTest(TestRecursorProtobuf):
 610     """
 611     This test makes sure that we correctly export queries and responses but only if they have been tagged.
 612     """
 613
 614     _confdir = 'ProtobufTaggedOnly'
 615     _config_template = """
 616 auth-zones=example=configs/%s/example.zone""" % _confdir
 617     _lua_config_file = """
 618     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=true, logResponses=true, taggedOnly=true } )
 619     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 620     _tags = ['tag1', 'tag2']
 621     _tag_from_gettag = 'tag-from-gettag'
 622     _lua_dns_script_file = """
 623     function gettag(remote, ednssubnet, localip, qname, qtype, ednsoptions, tcp)
 624       if qname:equal('tagged.example.') then
 625         return 0, { '%s' }
 626       end
 627       return 0
 628     end
 629     function preresolve(dq)
 630       if dq.qname:equal('tagged.example.') then
 631         dq:addPolicyTag('%s')
 632         dq:addPolicyTag('%s')
 633       end
 634       return false
 635     end
 636     """ % (_tag_from_gettag, _tags[0], _tags[1])
 637
 638     def testA(self):
 639         name = 'a.example.'
 640         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 641         query = dns.message.make_query(name, 'A', want_dnssec=True)
 642         query.flags |= dns.flags.CD
 643         res = self.sendUDPQuery(query)
 644         self.assertRRsetInAnswer(res, expected)
 645
 646         # check the protobuf message corresponding to the UDP response
 647         # the first query and answer are not tagged, so there is nothing in the queue
 648         time.sleep(1)
 649         self.checkNoRemainingMessage()
 650
 651     def testTagged(self):
 652         name = 'tagged.example.'
 653         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.84')
 654         query = dns.message.make_query(name, 'A', want_dnssec=True)
 655         query.flags |= dns.flags.CD
 656         res = self.sendUDPQuery(query)
 657         self.assertRRsetInAnswer(res, expected)
 658
 659         # check the protobuf messages corresponding to the UDP query and answer
 660         msg = self.getFirstProtobufMessage()
 661         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
 662         self.checkProtobufTags(msg, [ self._tag_from_gettag ])
 663         # then the response
 664         msg = self.getFirstProtobufMessage()
 665         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res)
 666         self.assertEqual(len(msg.response.rrs), 1)
 667         rr = msg.response.rrs[0]
 668         # we have max-cache-ttl set to 15
 669         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 670         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.84')
 671         tags = [ self._tag_from_gettag ] + self._tags
 672         self.checkProtobufTags(msg, tags)
 673         self.checkNoRemainingMessage()
 674
 675 class ProtobufSelectedFromLuaTest(TestRecursorProtobuf):
 676     """
 677     This test makes sure that we correctly export queries and responses but only if they have been selected from Lua.
 678     """
 679
 680     _confdir = 'ProtobufSelectedFromLua'
 681     _config_template = """
 682 auth-zones=example=configs/%s/example.zone""" % _confdir
 683     _lua_config_file = """
 684     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=false, logResponses=false } )
 685     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 686     _lua_dns_script_file = """
 687     local ffi = require("ffi")
 688
 689     ffi.cdef[[
 690       typedef struct pdns_ffi_param pdns_ffi_param_t;
 691
 692       const char* pdns_ffi_param_get_qname(pdns_ffi_param_t* ref);
 693       void pdns_ffi_param_set_log_query(pdns_ffi_param_t* ref, bool logQuery);
 694     ]]
 695
 696     function gettag_ffi(obj)
 697       qname = ffi.string(ffi.C.pdns_ffi_param_get_qname(obj))
 698       if qname == 'query-selected.example' then
 699         ffi.C.pdns_ffi_param_set_log_query(obj, true)
 700       end
 701       return 0
 702     end
 703
 704     function preresolve(dq)
 705       if dq.qname:equal('answer-selected.example.') then
 706         dq.logResponse = true
 707       end
 708       return false
 709     end
 710     """
 711
 712     def testA(self):
 713         name = 'a.example.'
 714         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 715         query = dns.message.make_query(name, 'A', want_dnssec=True)
 716         query.flags |= dns.flags.CD
 717         res = self.sendUDPQuery(query)
 718         self.assertRRsetInAnswer(res, expected)
 719
 720         # check the protobuf message corresponding to the UDP response
 721         # the first query and answer are not selected, so there is nothing in the queue
 722         self.checkNoRemainingMessage()
 723
 724     def testQuerySelected(self):
 725         name = 'query-selected.example.'
 726         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.84')
 727         query = dns.message.make_query(name, 'A', want_dnssec=True)
 728         query.flags |= dns.flags.CD
 729         res = self.sendUDPQuery(query)
 730         self.assertRRsetInAnswer(res, expected)
 731
 732         # check the protobuf messages corresponding to the UDP query
 733         msg = self.getFirstProtobufMessage()
 734         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
 735         # there should be no response
 736         self.checkNoRemainingMessage()
 737
 738     def testResponseSelected(self):
 739         name = 'answer-selected.example.'
 740         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.84')
 741         query = dns.message.make_query(name, 'A', want_dnssec=True)
 742         query.flags |= dns.flags.CD
 743         res = self.sendUDPQuery(query)
 744         self.assertRRsetInAnswer(res, expected)
 745
 746         # check the protobuf messages corresponding to the UDP response
 747         msg = self.getFirstProtobufMessage()
 748         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res)
 749         self.assertEqual(len(msg.response.rrs), 1)
 750         rr = msg.response.rrs[0]
 751         # we have max-cache-ttl set to 15
 752         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 753         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.84')
 754         self.checkNoRemainingMessage()
 755
 756 class ProtobufExportTypesTest(TestRecursorProtobuf):
 757     """
 758     This test makes sure that we correctly export other types than A, AAAA and CNAME over protobuf.
 759     """
 760
 761     _confdir = 'ProtobufExportTypes'
 762     _config_template = """
 763 auth-zones=example=configs/%s/example.zone""" % _confdir
 764     _lua_config_file = """
 765     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { exportTypes={"AAAA", "MX", "SPF", "SRV", "TXT"} } )
 766     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 767
 768     def testA(self):
 769         name = 'types.example.'
 770         expected = [dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.84'),
 771                     dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'AAAA', '2001:DB8::1'),
 772                     dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'MX', '10 a.example.'),
 773                     dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'SPF', '"v=spf1 -all"'),
 774                     dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'SRV', '10 20 443 a.example.'),
 775                     dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'TXT', '"Lorem ipsum dolor sit amet"'),
 776                    ]
 777         query = dns.message.make_query(name, 'ANY', want_dnssec=True)
 778         query.flags |= dns.flags.CD
 779         raw = self.sendUDPQuery(query, decode=False)
 780         res = dns.message.from_wire(raw)
 781
 782         for rrset in expected:
 783             self.assertRRsetInAnswer(res, rrset)
 784
 785         # check the protobuf messages corresponding to the UDP query and answer
 786         msg = self.getFirstProtobufMessage()
 787         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
 788         # then the response
 789         msg = self.getFirstProtobufMessage()
 790         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res, '127.0.0.1', receivedSize=len(raw))
 791         self.assertEqual(len(msg.response.rrs), 5)
 792         for rr in msg.response.rrs:
 793             self.assertTrue(rr.type in [dns.rdatatype.AAAA, dns.rdatatype.TXT, dns.rdatatype.MX, dns.rdatatype.SPF, dns.rdatatype.SRV])
 794
 795             if rr.type == dns.rdatatype.AAAA:
 796                 self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.AAAA, name, 15)
 797                 self.assertEqual(socket.inet_ntop(socket.AF_INET6, rr.rdata), '2001:db8::1')
 798             elif rr.type == dns.rdatatype.TXT:
 799                 self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.TXT, name, 15)
 800                 self.assertEqual(rr.rdata, b'"Lorem ipsum dolor sit amet"')
 801             elif rr.type == dns.rdatatype.MX:
 802                 self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.MX, name, 15)
 803                 self.assertEqual(rr.rdata, b'a.example.')
 804             elif rr.type == dns.rdatatype.SPF:
 805                 self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.SPF, name, 15)
 806                 self.assertEqual(rr.rdata, b'"v=spf1 -all"')
 807             elif rr.type == dns.rdatatype.SRV:
 808                 self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.SRV, name, 15)
 809                 self.assertEqual(rr.rdata, b'a.example.')
 810
 811         self.checkNoRemainingMessage()
 812
 813 class ProtobufTaggedExtraFieldsTest(TestRecursorProtobuf):
 814     """
 815     This test makes sure that we correctly export extra fields that may have been set while being tagged.
 816     """
 817
 818     _confdir = 'ProtobufTaggedExtraFields'
 819     _config_template = """
 820 auth-zones=example=configs/%s/example.zone""" % _confdir
 821     _lua_config_file = """
 822     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=true, logResponses=true } )
 823     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 824     _requestorId = 'S-000001727'
 825     _deviceId = 'd1:0a:91:dc:cc:82'
 826     _deviceName = 'Joe'
 827     _lua_dns_script_file = """
 828     function gettag(remote, ednssubnet, localip, qname, qtype, ednsoptions, tcp)
 829       if qname:equal('tagged.example.') then
 830         -- tag number, policy tags, data, requestorId, deviceId, deviceName
 831         return 0, {}, {}, '%s', '%s', '%s'
 832       end
 833       return 0
 834     end
 835     """ % (_requestorId, _deviceId, _deviceName)
 836
 837     def testA(self):
 838         name = 'a.example.'
 839         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 840         query = dns.message.make_query(name, 'A', want_dnssec=True)
 841         query.flags |= dns.flags.CD
 842         res = self.sendUDPQuery(query)
 843         self.assertRRsetInAnswer(res, expected)
 844
 845         # check the protobuf message corresponding to the UDP response
 846         # the first query and answer are not tagged, so there is nothing in the queue
 847         # check the protobuf messages corresponding to the UDP query and answer
 848         msg = self.getFirstProtobufMessage()
 849         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
 850         self.checkProtobufIdentity(msg, '', b'', '')
 851
 852         # then the response
 853         msg = self.getFirstProtobufMessage()
 854         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res, '127.0.0.1')
 855         self.assertEqual(len(msg.response.rrs), 1)
 856         rr = msg.response.rrs[0]
 857         # we have max-cache-ttl set to 15
 858         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 859         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
 860         self.checkProtobufIdentity(msg, '', b'', '')
 861         self.checkNoRemainingMessage()
 862
 863     def testTagged(self):
 864         name = 'tagged.example.'
 865         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.84')
 866         query = dns.message.make_query(name, 'A', want_dnssec=True)
 867         query.flags |= dns.flags.CD
 868         res = self.sendUDPQuery(query)
 869         self.assertRRsetInAnswer(res, expected)
 870
 871         # check the protobuf messages corresponding to the UDP query and answer
 872         msg = self.getFirstProtobufMessage()
 873         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
 874         self.checkProtobufIdentity(msg, self._requestorId, self._deviceId.encode('ascii'), self._deviceName)
 875
 876         # then the response
 877         msg = self.getFirstProtobufMessage()
 878         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res)
 879         self.assertEqual(len(msg.response.rrs), 1)
 880         rr = msg.response.rrs[0]
 881         # we have max-cache-ttl set to 15
 882         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 883         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.84')
 884         self.checkProtobufIdentity(msg, self._requestorId, self._deviceId.encode('ascii'), self._deviceName)
 885         self.checkNoRemainingMessage()
 886
 887 class ProtobufTaggedExtraFieldsFFITest(ProtobufTaggedExtraFieldsTest):
 888     """
 889     This test makes sure that we correctly export extra fields that may have been set while being tagged (FFI version).
 890     """
 891     _confdir = 'ProtobufTaggedExtraFieldsFFI'
 892     _config_template = """
 893 auth-zones=example=configs/%s/example.zone""" % _confdir
 894     _lua_config_file = """
 895     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=true, logResponses=true } )
 896     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 897     _lua_dns_script_file = """
 898     local ffi = require("ffi")
 899
 900     ffi.cdef[[
 901       typedef struct pdns_ffi_param pdns_ffi_param_t;
 902
 903       const char* pdns_ffi_param_get_qname(pdns_ffi_param_t* ref);
 904       void pdns_ffi_param_set_tag(pdns_ffi_param_t* ref, unsigned int tag);
 905       void pdns_ffi_param_set_requestorid(pdns_ffi_param_t* ref, const char* name);
 906       void pdns_ffi_param_set_devicename(pdns_ffi_param_t* ref, const char* name);
 907       void pdns_ffi_param_set_deviceid(pdns_ffi_param_t* ref, size_t len, const void* name);
 908     ]]
 909
 910     function gettag_ffi(obj)
 911       qname = ffi.string(ffi.C.pdns_ffi_param_get_qname(obj))
 912       if qname == 'tagged.example' then
 913         ffi.C.pdns_ffi_param_set_requestorid(obj, "%s")
 914         deviceid = "%s"
 915         ffi.C.pdns_ffi_param_set_deviceid(obj, string.len(deviceid), deviceid)
 916         ffi.C.pdns_ffi_param_set_devicename(obj, "%s")
 917       end
 918       return 0
 919     end
 920     """ % (ProtobufTaggedExtraFieldsTest._requestorId, ProtobufTaggedExtraFieldsTest._deviceId, ProtobufTaggedExtraFieldsTest._deviceName)
 921
 922 class ProtobufRPZTest(TestRecursorProtobuf):
 923     """
 924     This test makes sure that we correctly export the RPZ applied policy in our protobuf messages
 925     """
 926
 927     _confdir = 'ProtobufRPZ'
 928     _config_template = """
 929 auth-zones=example=configs/%s/example.rpz.zone""" % _confdir
 930     _lua_config_file = """
 931     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=true, logResponses=true } )
 932     rpzFile('configs/%s/zone.rpz', { policyName="zone.rpz."})
 933     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port, _confdir)
 934
 935     @classmethod
 936     def generateRecursorConfig(cls, confdir):
 937         authzonepath = os.path.join(confdir, 'example.rpz.zone')
 938         with open(authzonepath, 'w') as authzone:
 939             authzone.write("""$ORIGIN example.
 940 @ 3600 IN SOA {soa}
 941 sub.test 3600 IN A 192.0.2.42
 942 """.format(soa=cls._SOA))
 943
 944         rpzFilePath = os.path.join(confdir, 'zone.rpz')
 945         with open(rpzFilePath, 'w') as rpzZone:
 946             rpzZone.write("""$ORIGIN zone.rpz.
 947 @ 3600 IN SOA {soa}
 948 *.test.example.zone.rpz. 60 IN CNAME rpz-passthru.
 949 """.format(soa=cls._SOA))
 950
 951         super(ProtobufRPZTest, cls).generateRecursorConfig(confdir)
 952
 953     def testA(self):
 954         name = 'sub.test.example.'
 955         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 956         query = dns.message.make_query(name, 'A', want_dnssec=True)
 957         query.flags |= dns.flags.CD
 958         res = self.sendUDPQuery(query)
 959         self.assertRRsetInAnswer(res, expected)
 960
 961         # check the protobuf messages corresponding to the UDP query and answer
 962         msg = self.getFirstProtobufMessage()
 963         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
 964
 965         # then the response
 966         msg = self.getFirstProtobufMessage()
 967         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res)
 968         self.checkProtobufPolicy(msg, dnsmessage_pb2.PBDNSMessage.PolicyType.QNAME, 'zone.rpz.', '*.test.example.', 'sub.test.example', dnsmessage_pb2.PBDNSMessage.PolicyKind.NoAction)
 969         self.assertEqual(len(msg.response.rrs), 1)
 970         rr = msg.response.rrs[0]
 971         # we have max-cache-ttl set to 15
 972         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 973         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
 974         self.checkNoRemainingMessage()
 975
 976 class ProtobufRPZTagsTest(TestRecursorProtobuf):
 977     """
 978     This test makes sure that we correctly export the RPZ tags in our protobuf messages
 979     """
 980
 981     _confdir = 'ProtobufRPZTags'
 982     _config_template = """
 983 auth-zones=example=configs/%s/example.rpz.zone""" % _confdir
 984     _tags = ['tag1', 'tag2']
 985     _tags_from_gettag = ['tag1-from-gettag', 'tag2-from-gettag']
 986     _tags_from_rpz = ['tag1-from-rpz', 'tag2-from-rpz' ]
 987     _lua_config_file = """
 988     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=true, logResponses=true, tags={'tag1', 'tag2'} } )
 989     rpzFile('configs/%s/zone.rpz', { policyName="zone.rpz.", tags={ '%s', '%s'} })
 990     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port, _confdir, _tags_from_rpz[0], _tags_from_rpz[1])
 991     _lua_dns_script_file = """
 992     function gettag(remote, ednssubnet, localip, qname, qtype, ednsoptions, tcp)
 993       return 0, { '%s', '%s' }
 994     end
 995     function preresolve(dq)
 996       dq:addPolicyTag('%s')
 997       dq:addPolicyTag('%s')
 998       return false
 999     end
1000     """ % (_tags_from_gettag[0], _tags_from_gettag[1], _tags[0], _tags[1])
1001
1002     @classmethod
1003     def generateRecursorConfig(cls, confdir):
1004         authzonepath = os.path.join(confdir, 'example.rpz.zone')
1005         with open(authzonepath, 'w') as authzone:
1006             authzone.write("""$ORIGIN example.
1007 @ 3600 IN SOA {soa}
1008 sub.test 3600 IN A 192.0.2.42
1009 """.format(soa=cls._SOA))
1010
1011         rpzFilePath = os.path.join(confdir, 'zone.rpz')
1012         with open(rpzFilePath, 'w') as rpzZone:
1013             rpzZone.write("""$ORIGIN zone.rpz.
1014 @ 3600 IN SOA {soa}
1015 *.test.example.zone.rpz. 60 IN CNAME rpz-passthru.
1016 """.format(soa=cls._SOA))
1017
1018         super(ProtobufRPZTagsTest, cls).generateRecursorConfig(confdir)
1019
1020     def testA(self):
1021         name = 'sub.test.example.'
1022         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
1023         query = dns.message.make_query(name, 'A', want_dnssec=True)
1024         query.flags |= dns.flags.CD
1025         res = self.sendUDPQuery(query)
1026         self.assertRRsetInAnswer(res, expected)
1027
1028         # check the protobuf messages corresponding to the UDP query and answer
1029         msg = self.getFirstProtobufMessage()
1030         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
1031
1032         # then the response
1033         msg = self.getFirstProtobufMessage()
1034         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res)
1035         self.checkProtobufPolicy(msg, dnsmessage_pb2.PBDNSMessage.PolicyType.QNAME, 'zone.rpz.', '*.test.example.', 'sub.test.example', dnsmessage_pb2.PBDNSMessage.PolicyKind.NoAction)
1036         self.checkProtobufTags(msg, self._tags + self._tags_from_gettag + self._tags_from_rpz)
1037         self.assertEqual(len(msg.response.rrs), 1)
1038         rr = msg.response.rrs[0]
1039         # we have max-cache-ttl set to 15
1040         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
1041         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
1042         self.checkNoRemainingMessage()
1043
1044
1045 class ProtobufMetaFFITest(TestRecursorProtobuf):
1046     """
1047     This test makes sure that we can correctly add extra meta fields (FFI version).
1048     """
1049     _confdir = 'ProtobufMetaFFITest'
1050     _config_template = """
1051 auth-zones=example=configs/%s/example.zone""" % _confdir
1052     _lua_config_file = """
1053     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=true, logResponses=true } )
1054     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
1055     _lua_dns_script_file = """
1056     local ffi = require("ffi")
1057
1058     ffi.cdef[[
1059       typedef struct pdns_ffi_param pdns_ffi_param_t;
1060
1061       const char* pdns_ffi_param_get_qname(pdns_ffi_param_t* ref);
1062       void pdns_ffi_param_add_meta_single_string_kv(pdns_ffi_param_t *ref, const char* key, const char* val);
1063       void pdns_ffi_param_add_meta_single_int64_kv(pdns_ffi_param_t *ref, const char* key, int64_t val);
1064     ]]
1065
1066     function gettag_ffi(obj)
1067       qname = ffi.string(ffi.C.pdns_ffi_param_get_qname(obj))
1068       if qname == 'meta.example' then
1069         ffi.C.pdns_ffi_param_add_meta_single_string_kv(obj, "meta-str", "keyword")
1070         ffi.C.pdns_ffi_param_add_meta_single_int64_kv(obj, "meta-int", 42)
1071         ffi.C.pdns_ffi_param_add_meta_single_string_kv(obj, "meta-str", "content")
1072         ffi.C.pdns_ffi_param_add_meta_single_int64_kv(obj, "meta-int", 21)
1073       end
1074       return 0
1075     end
1076     """
1077     def testMeta(self):
1078         name = 'meta.example.'
1079         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.85')
1080         query = dns.message.make_query(name, 'A', want_dnssec=True)
1081         query.flags |= dns.flags.CD
1082         res = self.sendUDPQuery(query)
1083         self.assertRRsetInAnswer(res, expected)
1084
1085         # check the protobuf messages corresponding to the UDP query and answer
1086         msg = self.getFirstProtobufMessage()
1087         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
1088         self.checkProtobufMetas(msg, {'meta-str': { "stringVal" : ["content", "keyword"]}, 'meta-int': {"intVal" : [21, 42]}})
1089
1090         # then the response
1091         msg = self.getFirstProtobufMessage()
1092         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res)
1093         self.assertEqual(len(msg.response.rrs), 1)
1094         rr = msg.response.rrs[0]
1095         # we have max-cache-ttl set to 15
1096         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
1097         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.85')
1098         self.checkProtobufMetas(msg, {'meta-str': { "stringVal" : ["content", "keyword"]}, 'meta-int': {"intVal" : [21, 42]}})
1099
1100         self.checkNoRemainingMessage()