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 ProtobufProxyMappingTest(TestRecursorProtobuf):
 372     """
 373     This test makes sure that we correctly export queries and response over protobuf with a proxyMapping
 374     """
 375
 376     _confdir = 'ProtobufProxyMappingTest'
 377     _config_template = """
 378     auth-zones=example=configs/%s/example.zone
 379     allow-from=3.4.5.0/24
 380     """ % _confdir
 381
 382     _lua_config_file = """
 383     addProxyMapping("127.0.0.1/24", "3.4.5.6:99")
 384     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"})
 385     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 386
 387     def testA(self):
 388         name = 'a.example.'
 389         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 390         query = dns.message.make_query(name, 'A', want_dnssec=True)
 391         query.flags |= dns.flags.CD
 392         res = self.sendUDPQuery(query)
 393
 394         self.assertRRsetInAnswer(res, expected)
 395
 396         # check the protobuf messages corresponding to the UDP query and answer
 397         msg = self.getFirstProtobufMessage()
 398         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
 399         # then the response
 400         msg = self.getFirstProtobufMessage()
 401         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res, '127.0.0.1')
 402         self.assertEqual(len(msg.response.rrs), 1)
 403         rr = msg.response.rrs[0]
 404         # we have max-cache-ttl set to 15
 405         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 406         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
 407         self.checkNoRemainingMessage()
 408
 409 class ProtobufProxyMappingLogMappedTest(TestRecursorProtobuf):
 410     """
 411     This test makes sure that we correctly export queries and response over protobuf.
 412     """
 413
 414     _confdir = 'ProtobufProxyMappingLogMappedTest'
 415     _config_template = """
 416     auth-zones=example=configs/%s/example.zone
 417     allow-from=3.4.5.0/0"
 418     """ % _confdir
 419
 420     _lua_config_file = """
 421     addProxyMapping("127.0.0.1/24", "3.4.5.6:99")
 422     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logMappedFrom = true })
 423     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 424
 425     def testA(self):
 426         name = 'a.example.'
 427         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 428         query = dns.message.make_query(name, 'A', want_dnssec=True)
 429         query.flags |= dns.flags.CD
 430         res = self.sendUDPQuery(query)
 431
 432         self.assertRRsetInAnswer(res, expected)
 433
 434         # check the protobuf messages corresponding to the UDP query and answer
 435         msg = self.getFirstProtobufMessage()
 436         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name, '3.4.5.6')
 437         # then the response
 438         msg = self.getFirstProtobufMessage()
 439         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res, '3.4.5.6')
 440         self.assertEqual(len(msg.response.rrs), 1)
 441         rr = msg.response.rrs[0]
 442         # we have max-cache-ttl set to 15
 443         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 444         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
 445         self.checkNoRemainingMessage()
 446
 447 class ProtobufProxyTest(TestRecursorProtobuf):
 448     """
 449     This test makes sure that we correctly export addresses over protobuf when the proxy protocol is used.
 450     """
 451
 452     _confdir = 'ProtobufProxy'
 453     _config_template = """
 454 auth-zones=example=configs/%s/example.zone
 455 proxy-protocol-from=127.0.0.1/32
 456 allow-from=127.0.0.1,6.6.6.6
 457 """ % _confdir
 458
 459     def testA(self):
 460         name = 'a.example.'
 461         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 462         query = dns.message.make_query(name, 'A', want_dnssec=True)
 463         query.flags |= dns.flags.CD
 464         res = self.sendUDPQueryWithProxyProtocol(query, False, '6.6.6.6', '7.7.7.7', 666, 777)
 465
 466         self.assertRRsetInAnswer(res, expected)
 467
 468         # check the protobuf messages corresponding to the UDP query and answer
 469         msg = self.getFirstProtobufMessage()
 470         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name, '6.6.6.6', '7.7.7.7')
 471         # then the response
 472         msg = self.getFirstProtobufMessage()
 473         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res, '6.6.6.6')
 474         self.assertEqual(len(msg.response.rrs), 1)
 475         rr = msg.response.rrs[0]
 476         # we have max-cache-ttl set to 15
 477         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 478         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
 479         self.checkNoRemainingMessage()
 480
 481 class ProtobufProxyWithProxyByTableTest(TestRecursorProtobuf):
 482     """
 483     This test makes sure that we correctly export addresses over protobuf when the proxy protocol and a proxy table mapping is used
 484     """
 485
 486     _confdir = 'ProtobufProxyWithProxyByTable'
 487     _config_template = """
 488 auth-zones=example=configs/%s/example.zone
 489 proxy-protocol-from=127.0.0.1/32
 490 allow-from=3.4.5.6
 491 """ % _confdir
 492
 493     _lua_config_file = """
 494     addProxyMapping("6.6.6.6/24", "3.4.5.6:99")
 495     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"})
 496     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 497
 498     def testA(self):
 499         name = 'a.example.'
 500         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 501         query = dns.message.make_query(name, 'A', want_dnssec=True)
 502         query.flags |= dns.flags.CD
 503         res = self.sendUDPQueryWithProxyProtocol(query, False, '6.6.6.6', '7.7.7.7', 666, 777)
 504
 505         self.assertRRsetInAnswer(res, expected)
 506
 507         # check the protobuf messages corresponding to the UDP query and answer
 508         msg = self.getFirstProtobufMessage()
 509         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name, '6.6.6.6', '7.7.7.7')
 510         # then the response
 511         msg = self.getFirstProtobufMessage()
 512         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res, '6.6.6.6')
 513         self.assertEqual(len(msg.response.rrs), 1)
 514         rr = msg.response.rrs[0]
 515         # we have max-cache-ttl set to 15
 516         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 517         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
 518         self.checkNoRemainingMessage()
 519
 520 class ProtobufProxyWithProxyByTableLogMappedTest(TestRecursorProtobuf):
 521     """
 522     This test makes sure that we correctly export addresses over protobuf when the proxy protocol and a proxy table mapping is used
 523     """
 524
 525     _confdir = 'ProtobufProxyWithProxyByTableLogMapped'
 526     _config_template = """
 527 auth-zones=example=configs/%s/example.zone
 528 proxy-protocol-from=127.0.0.1/32
 529 allow-from=3.4.5.6
 530 """ % _confdir
 531
 532     _lua_config_file = """
 533     addProxyMapping("6.6.6.6/24", "3.4.5.6:99")
 534     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logMappedFrom = true })
 535     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 536
 537     def testA(self):
 538         name = 'a.example.'
 539         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 540         query = dns.message.make_query(name, 'A', want_dnssec=True)
 541         query.flags |= dns.flags.CD
 542         res = self.sendUDPQueryWithProxyProtocol(query, False, '6.6.6.6', '7.7.7.7', 666, 777)
 543
 544         self.assertRRsetInAnswer(res, expected)
 545
 546         # check the protobuf messages corresponding to the UDP query and answer
 547         msg = self.getFirstProtobufMessage()
 548         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name, '3.4.5.6', '7.7.7.7')
 549         # then the response
 550         msg = self.getFirstProtobufMessage()
 551         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res, '3.4.5.6')
 552         self.assertEqual(len(msg.response.rrs), 1)
 553         rr = msg.response.rrs[0]
 554         # we have max-cache-ttl set to 15
 555         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 556         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
 557         self.checkNoRemainingMessage()
 558
 559
 560 class OutgoingProtobufDefaultTest(TestRecursorProtobuf):
 561     """
 562     This test makes sure that we correctly export outgoing queries over protobuf.
 563     It must be improved and setup env so we can check for incoming responses, but makes sure for now
 564     that the recursor at least connects to the protobuf server.
 565     """
 566
 567     _confdir = 'OutgoingProtobufDefault'
 568     _config_template = """
 569     # Switch off QName Minimization, it generates much more protobuf messages
 570     # (or make the test much more smart!)
 571     qname-minimization=no
 572 """
 573     _lua_config_file = """
 574     outgoingProtobufServer({"127.0.0.1:%d", "127.0.0.1:%d"})
 575     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 576
 577     def testA(self):
 578         name = 'host1.secure.example.'
 579         expected = list()
 580
 581         # the root DNSKEY has been learned with priming the root NS already
 582         # ('.', dns.rdatatype.DNSKEY, dnsmessage_pb2.PBDNSMessage.UDP, 201),
 583         for qname, qtype, proto, responseSize in [
 584                 ('host1.secure.example.', dns.rdatatype.A, dnsmessage_pb2.PBDNSMessage.UDP, 248),
 585                 ('host1.secure.example.', dns.rdatatype.A, dnsmessage_pb2.PBDNSMessage.UDP, 221),
 586                 ('example.', dns.rdatatype.DNSKEY, dnsmessage_pb2.PBDNSMessage.UDP, 219),
 587                 ('host1.secure.example.', dns.rdatatype.A, dnsmessage_pb2.PBDNSMessage.UDP, 175),
 588                 ('secure.example.', dns.rdatatype.DNSKEY, dnsmessage_pb2.PBDNSMessage.UDP, 233),
 589         ]:
 590             if not qname:
 591                 expected.append((None, None, None, None, None, None))
 592                 continue
 593             query = dns.message.make_query(qname, qtype, use_edns=True, want_dnssec=True)
 594             resp = dns.message.make_response(query)
 595             expected.append((
 596                 qname, qtype, query, resp, proto, responseSize
 597             ))
 598
 599         query = dns.message.make_query(name, 'A', want_dnssec=True)
 600         query.flags |= dns.flags.RD
 601         res = self.sendUDPQuery(query)
 602
 603         for qname, qtype, qry, ans, proto, responseSize in expected:
 604             if not qname:
 605                 self.getFirstProtobufMessage()
 606                 self.getFirstProtobufMessage()
 607                 continue
 608
 609             msg = self.getFirstProtobufMessage()
 610             self.checkProtobufOutgoingQuery(msg, proto, qry, dns.rdataclass.IN, qtype, qname)
 611
 612             # Check the answer
 613             msg = self.getFirstProtobufMessage()
 614             self.checkProtobufIncomingResponse(msg, proto, ans, length=responseSize)
 615
 616         self.checkNoRemainingMessage()
 617
 618 class OutgoingProtobufNoQueriesTest(TestRecursorProtobuf):
 619     """
 620     This test makes sure that we correctly export incoming responses but not outgoing queries over protobuf.
 621     It must be improved and setup env so we can check for incoming responses, but makes sure for now
 622     that the recursor at least connects to the protobuf server.
 623     """
 624
 625     _confdir = 'OutgoingProtobufNoQueries'
 626     _config_template = """
 627     # Switch off QName Minimization, it generates much more protobuf messages
 628     # (or make the test much more smart!)
 629     qname-minimization=no"""
 630     _lua_config_file = """
 631     outgoingProtobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=false, logResponses=true })
 632     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 633
 634     def testA(self):
 635         name = 'host1.secure.example.'
 636         expected = list()
 637         # the root DNSKEY has been learned with priming the root NS already
 638         # ('.', dns.rdatatype.DNSKEY, dnsmessage_pb2.PBDNSMessage.UDP, 201),
 639         for qname, qtype, proto, size in [
 640                 ('host1.secure.example.', dns.rdatatype.A, dnsmessage_pb2.PBDNSMessage.UDP, 248),
 641                 ('host1.secure.example.', dns.rdatatype.A, dnsmessage_pb2.PBDNSMessage.UDP, 221),
 642                 ('example.', dns.rdatatype.DNSKEY, dnsmessage_pb2.PBDNSMessage.UDP, 219),
 643                 ('host1.secure.example.', dns.rdatatype.A, dnsmessage_pb2.PBDNSMessage.UDP, 175),
 644                 ('secure.example.', dns.rdatatype.DNSKEY, dnsmessage_pb2.PBDNSMessage.UDP, 233),
 645         ]:
 646             if not qname:
 647                 expected.append((None, None, None, None, None, None))
 648                 continue
 649             query = dns.message.make_query(qname, qtype, use_edns=True, want_dnssec=True)
 650             resp = dns.message.make_response(query)
 651             expected.append((
 652                 qname, qtype, query, resp, proto, size
 653             ))
 654
 655         query = dns.message.make_query(name, 'A', want_dnssec=True)
 656         query.flags |= dns.flags.RD
 657         res = self.sendUDPQuery(query)
 658
 659         for qname, qtype, qry, ans, proto, size in expected:
 660             if not qname:
 661                 self.getFirstProtobufMessage()
 662                 continue
 663
 664             # check the response
 665             msg = self.getFirstProtobufMessage()
 666             self.checkProtobufIncomingResponse(msg, proto, ans, length=size)
 667
 668         self.checkNoRemainingMessage()
 669
 670 class ProtobufMasksTest(TestRecursorProtobuf):
 671     """
 672     This test makes sure that we correctly export queries and response over protobuf, respecting the configured initiator masking.
 673     """
 674
 675     _confdir = 'ProtobufMasks'
 676     _config_template = """
 677 auth-zones=example=configs/%s/example.zone""" % _confdir
 678     _protobufMaskV4 = 4
 679     _protobufMaskV6 = 128
 680     _lua_config_file = """
 681     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"})
 682     setProtobufMasks(%d, %d)
 683     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port, _protobufMaskV4, _protobufMaskV6)
 684
 685     def testA(self):
 686         name = 'a.example.'
 687         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 688         query = dns.message.make_query(name, 'A', want_dnssec=True)
 689         query.flags |= dns.flags.CD
 690         res = self.sendUDPQuery(query)
 691         self.assertRRsetInAnswer(res, expected)
 692
 693         # check the protobuf messages corresponding to the UDP query and answer
 694         # but first let the protobuf messages the time to get there
 695         msg = self.getFirstProtobufMessage()
 696         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name, '112.0.0.0')
 697         # then the response
 698         msg = self.getFirstProtobufMessage()
 699         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res, '112.0.0.0')
 700         self.assertEqual(len(msg.response.rrs), 1)
 701         rr = msg.response.rrs[0]
 702         # we have max-cache-ttl set to 15
 703         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 704         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
 705         self.checkNoRemainingMessage()
 706
 707 class ProtobufQueriesOnlyTest(TestRecursorProtobuf):
 708     """
 709     This test makes sure that we correctly export queries but not responses over protobuf.
 710     """
 711
 712     _confdir = 'ProtobufQueriesOnly'
 713     _config_template = """
 714 auth-zones=example=configs/%s/example.zone""" % _confdir
 715     _lua_config_file = """
 716     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=true, logResponses=false } )
 717     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 718
 719     def testA(self):
 720         name = 'a.example.'
 721         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 722         query = dns.message.make_query(name, 'A', want_dnssec=True)
 723         query.flags |= dns.flags.CD
 724         res = self.sendUDPQuery(query)
 725         self.assertRRsetInAnswer(res, expected)
 726
 727         # check the protobuf message corresponding to the UDP query
 728         msg = self.getFirstProtobufMessage()
 729         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
 730         # no response
 731         self.checkNoRemainingMessage()
 732
 733 class ProtobufResponsesOnlyTest(TestRecursorProtobuf):
 734     """
 735     This test makes sure that we correctly export responses but not queries over protobuf.
 736     """
 737
 738     _confdir = 'ProtobufResponsesOnly'
 739     _config_template = """
 740 auth-zones=example=configs/%s/example.zone""" % _confdir
 741     _lua_config_file = """
 742     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=false, logResponses=true } )
 743     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 744
 745     def testA(self):
 746         name = 'a.example.'
 747         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 748         query = dns.message.make_query(name, 'A', want_dnssec=True)
 749         query.flags |= dns.flags.CD
 750         res = self.sendUDPQuery(query)
 751         self.assertRRsetInAnswer(res, expected)
 752
 753         # check the protobuf message corresponding to the UDP response
 754         msg = self.getFirstProtobufMessage()
 755         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res)
 756         self.assertEqual(len(msg.response.rrs), 1)
 757         rr = msg.response.rrs[0]
 758         # we have max-cache-ttl set to 15
 759         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 760         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
 761         # nothing else in the queue
 762         self.checkNoRemainingMessage()
 763
 764 class ProtobufTaggedOnlyTest(TestRecursorProtobuf):
 765     """
 766     This test makes sure that we correctly export queries and responses but only if they have been tagged.
 767     """
 768
 769     _confdir = 'ProtobufTaggedOnly'
 770     _config_template = """
 771 auth-zones=example=configs/%s/example.zone""" % _confdir
 772     _lua_config_file = """
 773     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=true, logResponses=true, taggedOnly=true } )
 774     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 775     _tags = ['tag1', 'tag2']
 776     _tag_from_gettag = 'tag-from-gettag'
 777     _lua_dns_script_file = """
 778     function gettag(remote, ednssubnet, localip, qname, qtype, ednsoptions, tcp)
 779       if qname:equal('tagged.example.') then
 780         return 0, { '%s' }
 781       end
 782       return 0
 783     end
 784     function preresolve(dq)
 785       if dq.qname:equal('tagged.example.') then
 786         dq:addPolicyTag('%s')
 787         dq:addPolicyTag('%s')
 788       end
 789       return false
 790     end
 791     """ % (_tag_from_gettag, _tags[0], _tags[1])
 792
 793     def testA(self):
 794         name = 'a.example.'
 795         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 796         query = dns.message.make_query(name, 'A', want_dnssec=True)
 797         query.flags |= dns.flags.CD
 798         res = self.sendUDPQuery(query)
 799         self.assertRRsetInAnswer(res, expected)
 800
 801         # check the protobuf message corresponding to the UDP response
 802         # the first query and answer are not tagged, so there is nothing in the queue
 803         time.sleep(1)
 804         self.checkNoRemainingMessage()
 805
 806     def testTagged(self):
 807         name = 'tagged.example.'
 808         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.84')
 809         query = dns.message.make_query(name, 'A', want_dnssec=True)
 810         query.flags |= dns.flags.CD
 811         res = self.sendUDPQuery(query)
 812         self.assertRRsetInAnswer(res, expected)
 813
 814         # check the protobuf messages corresponding to the UDP query and answer
 815         msg = self.getFirstProtobufMessage()
 816         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
 817         self.checkProtobufTags(msg, [ self._tag_from_gettag ])
 818         # then the response
 819         msg = self.getFirstProtobufMessage()
 820         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res)
 821         self.assertEqual(len(msg.response.rrs), 1)
 822         rr = msg.response.rrs[0]
 823         # we have max-cache-ttl set to 15
 824         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 825         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.84')
 826         tags = [ self._tag_from_gettag ] + self._tags
 827         self.checkProtobufTags(msg, tags)
 828         self.checkNoRemainingMessage()
 829
 830 class ProtobufSelectedFromLuaTest(TestRecursorProtobuf):
 831     """
 832     This test makes sure that we correctly export queries and responses but only if they have been selected from Lua.
 833     """
 834
 835     _confdir = 'ProtobufSelectedFromLua'
 836     _config_template = """
 837 auth-zones=example=configs/%s/example.zone""" % _confdir
 838     _lua_config_file = """
 839     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=false, logResponses=false } )
 840     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 841     _lua_dns_script_file = """
 842     local ffi = require("ffi")
 843
 844     ffi.cdef[[
 845       typedef struct pdns_ffi_param pdns_ffi_param_t;
 846
 847       const char* pdns_ffi_param_get_qname(pdns_ffi_param_t* ref);
 848       void pdns_ffi_param_set_log_query(pdns_ffi_param_t* ref, bool logQuery);
 849     ]]
 850
 851     function gettag_ffi(obj)
 852       qname = ffi.string(ffi.C.pdns_ffi_param_get_qname(obj))
 853       if qname == 'query-selected.example' then
 854         ffi.C.pdns_ffi_param_set_log_query(obj, true)
 855       end
 856       return 0
 857     end
 858
 859     function preresolve(dq)
 860       if dq.qname:equal('answer-selected.example.') then
 861         dq.logResponse = true
 862       end
 863       return false
 864     end
 865     """
 866
 867     def testA(self):
 868         name = 'a.example.'
 869         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 870         query = dns.message.make_query(name, 'A', want_dnssec=True)
 871         query.flags |= dns.flags.CD
 872         res = self.sendUDPQuery(query)
 873         self.assertRRsetInAnswer(res, expected)
 874
 875         # check the protobuf message corresponding to the UDP response
 876         # the first query and answer are not selected, so there is nothing in the queue
 877         self.checkNoRemainingMessage()
 878
 879     def testQuerySelected(self):
 880         name = 'query-selected.example.'
 881         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.84')
 882         query = dns.message.make_query(name, 'A', want_dnssec=True)
 883         query.flags |= dns.flags.CD
 884         res = self.sendUDPQuery(query)
 885         self.assertRRsetInAnswer(res, expected)
 886
 887         # check the protobuf messages corresponding to the UDP query
 888         msg = self.getFirstProtobufMessage()
 889         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
 890         # there should be no response
 891         self.checkNoRemainingMessage()
 892
 893     def testResponseSelected(self):
 894         name = 'answer-selected.example.'
 895         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.84')
 896         query = dns.message.make_query(name, 'A', want_dnssec=True)
 897         query.flags |= dns.flags.CD
 898         res = self.sendUDPQuery(query)
 899         self.assertRRsetInAnswer(res, expected)
 900
 901         # check the protobuf messages corresponding to the UDP response
 902         msg = self.getFirstProtobufMessage()
 903         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res)
 904         self.assertEqual(len(msg.response.rrs), 1)
 905         rr = msg.response.rrs[0]
 906         # we have max-cache-ttl set to 15
 907         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 908         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.84')
 909         self.checkNoRemainingMessage()
 910
 911 class ProtobufExportTypesTest(TestRecursorProtobuf):
 912     """
 913     This test makes sure that we correctly export other types than A, AAAA and CNAME over protobuf.
 914     """
 915
 916     _confdir = 'ProtobufExportTypes'
 917     _config_template = """
 918 auth-zones=example=configs/%s/example.zone""" % _confdir
 919     _lua_config_file = """
 920     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { exportTypes={"AAAA", "MX", "SPF", "SRV", "TXT"} } )
 921     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 922
 923     def testA(self):
 924         name = 'types.example.'
 925         expected = [dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.84'),
 926                     dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'AAAA', '2001:DB8::1'),
 927                     dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'MX', '10 a.example.'),
 928                     dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'SPF', '"v=spf1 -all"'),
 929                     dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'SRV', '10 20 443 a.example.'),
 930                     dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'TXT', '"Lorem ipsum dolor sit amet"'),
 931                    ]
 932         query = dns.message.make_query(name, 'ANY', want_dnssec=True)
 933         query.flags |= dns.flags.CD
 934         raw = self.sendUDPQuery(query, decode=False)
 935         res = dns.message.from_wire(raw)
 936
 937         for rrset in expected:
 938             self.assertRRsetInAnswer(res, rrset)
 939
 940         # check the protobuf messages corresponding to the UDP query and answer
 941         msg = self.getFirstProtobufMessage()
 942         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
 943         # then the response
 944         msg = self.getFirstProtobufMessage()
 945         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res, '127.0.0.1', receivedSize=len(raw))
 946         self.assertEqual(len(msg.response.rrs), 5)
 947         for rr in msg.response.rrs:
 948             self.assertTrue(rr.type in [dns.rdatatype.AAAA, dns.rdatatype.TXT, dns.rdatatype.MX, dns.rdatatype.SPF, dns.rdatatype.SRV])
 949
 950             if rr.type == dns.rdatatype.AAAA:
 951                 self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.AAAA, name, 15)
 952                 self.assertEqual(socket.inet_ntop(socket.AF_INET6, rr.rdata), '2001:db8::1')
 953             elif rr.type == dns.rdatatype.TXT:
 954                 self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.TXT, name, 15)
 955                 self.assertEqual(rr.rdata, b'"Lorem ipsum dolor sit amet"')
 956             elif rr.type == dns.rdatatype.MX:
 957                 self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.MX, name, 15)
 958                 self.assertEqual(rr.rdata, b'a.example.')
 959             elif rr.type == dns.rdatatype.SPF:
 960                 self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.SPF, name, 15)
 961                 self.assertEqual(rr.rdata, b'"v=spf1 -all"')
 962             elif rr.type == dns.rdatatype.SRV:
 963                 self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.SRV, name, 15)
 964                 self.assertEqual(rr.rdata, b'a.example.')
 965
 966         self.checkNoRemainingMessage()
 967
 968 class ProtobufTaggedExtraFieldsTest(TestRecursorProtobuf):
 969     """
 970     This test makes sure that we correctly export extra fields that may have been set while being tagged.
 971     """
 972
 973     _confdir = 'ProtobufTaggedExtraFields'
 974     _config_template = """
 975 auth-zones=example=configs/%s/example.zone""" % _confdir
 976     _lua_config_file = """
 977     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=true, logResponses=true } )
 978     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 979     _requestorId = 'S-000001727'
 980     _deviceId = 'd1:0a:91:dc:cc:82'
 981     _deviceName = 'Joe'
 982     _lua_dns_script_file = """
 983     function gettag(remote, ednssubnet, localip, qname, qtype, ednsoptions, tcp)
 984       if qname:equal('tagged.example.') then
 985         -- tag number, policy tags, data, requestorId, deviceId, deviceName
 986         return 0, {}, {}, '%s', '%s', '%s'
 987       end
 988       return 0
 989     end
 990     """ % (_requestorId, _deviceId, _deviceName)
 991
 992     def testA(self):
 993         name = 'a.example.'
 994         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 995         query = dns.message.make_query(name, 'A', want_dnssec=True)
 996         query.flags |= dns.flags.CD
 997         res = self.sendUDPQuery(query)
 998         self.assertRRsetInAnswer(res, expected)
 999
1000         # check the protobuf message corresponding to the UDP response
1001         # the first query and answer are not tagged, so there is nothing in the queue
1002         # check the protobuf messages corresponding to the UDP query and answer
1003         msg = self.getFirstProtobufMessage()
1004         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
1005         self.checkProtobufIdentity(msg, '', b'', '')
1006
1007         # then the response
1008         msg = self.getFirstProtobufMessage()
1009         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res, '127.0.0.1')
1010         self.assertEqual(len(msg.response.rrs), 1)
1011         rr = msg.response.rrs[0]
1012         # we have max-cache-ttl set to 15
1013         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
1014         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
1015         self.checkProtobufIdentity(msg, '', b'', '')
1016         self.checkNoRemainingMessage()
1017
1018     def testTagged(self):
1019         name = 'tagged.example.'
1020         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.84')
1021         query = dns.message.make_query(name, 'A', want_dnssec=True)
1022         query.flags |= dns.flags.CD
1023         res = self.sendUDPQuery(query)
1024         self.assertRRsetInAnswer(res, expected)
1025
1026         # check the protobuf messages corresponding to the UDP query and answer
1027         msg = self.getFirstProtobufMessage()
1028         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
1029         self.checkProtobufIdentity(msg, self._requestorId, self._deviceId.encode('ascii'), self._deviceName)
1030
1031         # then the response
1032         msg = self.getFirstProtobufMessage()
1033         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res)
1034         self.assertEqual(len(msg.response.rrs), 1)
1035         rr = msg.response.rrs[0]
1036         # we have max-cache-ttl set to 15
1037         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
1038         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.84')
1039         self.checkProtobufIdentity(msg, self._requestorId, self._deviceId.encode('ascii'), self._deviceName)
1040         self.checkNoRemainingMessage()
1041
1042 class ProtobufTaggedExtraFieldsFFITest(ProtobufTaggedExtraFieldsTest):
1043     """
1044     This test makes sure that we correctly export extra fields that may have been set while being tagged (FFI version).
1045     """
1046     _confdir = 'ProtobufTaggedExtraFieldsFFI'
1047     _config_template = """
1048 auth-zones=example=configs/%s/example.zone""" % _confdir
1049     _lua_config_file = """
1050     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=true, logResponses=true } )
1051     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
1052     _lua_dns_script_file = """
1053     local ffi = require("ffi")
1054
1055     ffi.cdef[[
1056       typedef struct pdns_ffi_param pdns_ffi_param_t;
1057
1058       const char* pdns_ffi_param_get_qname(pdns_ffi_param_t* ref);
1059       void pdns_ffi_param_set_tag(pdns_ffi_param_t* ref, unsigned int tag);
1060       void pdns_ffi_param_set_requestorid(pdns_ffi_param_t* ref, const char* name);
1061       void pdns_ffi_param_set_devicename(pdns_ffi_param_t* ref, const char* name);
1062       void pdns_ffi_param_set_deviceid(pdns_ffi_param_t* ref, size_t len, const void* name);
1063     ]]
1064
1065     function gettag_ffi(obj)
1066       qname = ffi.string(ffi.C.pdns_ffi_param_get_qname(obj))
1067       if qname == 'tagged.example' then
1068         ffi.C.pdns_ffi_param_set_requestorid(obj, "%s")
1069         deviceid = "%s"
1070         ffi.C.pdns_ffi_param_set_deviceid(obj, string.len(deviceid), deviceid)
1071         ffi.C.pdns_ffi_param_set_devicename(obj, "%s")
1072       end
1073       return 0
1074     end
1075     """ % (ProtobufTaggedExtraFieldsTest._requestorId, ProtobufTaggedExtraFieldsTest._deviceId, ProtobufTaggedExtraFieldsTest._deviceName)
1076
1077 class ProtobufRPZTest(TestRecursorProtobuf):
1078     """
1079     This test makes sure that we correctly export the RPZ applied policy in our protobuf messages
1080     """
1081
1082     _confdir = 'ProtobufRPZ'
1083     _config_template = """
1084 auth-zones=example=configs/%s/example.rpz.zone""" % _confdir
1085     _lua_config_file = """
1086     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=true, logResponses=true } )
1087     rpzFile('configs/%s/zone.rpz', { policyName="zone.rpz."})
1088     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port, _confdir)
1089
1090     @classmethod
1091     def generateRecursorConfig(cls, confdir):
1092         authzonepath = os.path.join(confdir, 'example.rpz.zone')
1093         with open(authzonepath, 'w') as authzone:
1094             authzone.write("""$ORIGIN example.
1095 @ 3600 IN SOA {soa}
1096 sub.test 3600 IN A 192.0.2.42
1097 ip  3600 IN A 33.22.11.99
1098 """.format(soa=cls._SOA))
1099
1100         rpzFilePath = os.path.join(confdir, 'zone.rpz')
1101         with open(rpzFilePath, 'w') as rpzZone:
1102             rpzZone.write("""$ORIGIN zone.rpz.
1103 @ 3600 IN SOA {soa}
1104 *.test.example.zone.rpz. 60 IN CNAME rpz-passthru.
1105 24.0.11.22.33.rpz-ip     60 IN A 1.2.3.4
1106 """.format(soa=cls._SOA))
1107
1108         super(ProtobufRPZTest, cls).generateRecursorConfig(confdir)
1109
1110     def testA(self):
1111         name = 'sub.test.example.'
1112         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
1113         query = dns.message.make_query(name, 'A', want_dnssec=True)
1114         query.flags |= dns.flags.CD
1115         res = self.sendUDPQuery(query)
1116         self.assertRRsetInAnswer(res, expected)
1117
1118         # check the protobuf messages corresponding to the UDP query and answer
1119         msg = self.getFirstProtobufMessage()
1120         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
1121
1122         # then the response
1123         msg = self.getFirstProtobufMessage()
1124         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res)
1125         self.checkProtobufPolicy(msg, dnsmessage_pb2.PBDNSMessage.PolicyType.QNAME, 'zone.rpz.', '*.test.example.', 'sub.test.example', dnsmessage_pb2.PBDNSMessage.PolicyKind.NoAction)
1126         self.assertEqual(len(msg.response.rrs), 1)
1127         rr = msg.response.rrs[0]
1128         # we have max-cache-ttl set to 15
1129         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
1130         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
1131         self.checkNoRemainingMessage()
1132
1133     def testB(self):
1134         name = 'ip.example.'
1135         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '1.2.3.4')
1136         query = dns.message.make_query(name, 'A', want_dnssec=True)
1137         query.flags |= dns.flags.CD
1138         res = self.sendUDPQuery(query)
1139         self.assertRRsetInAnswer(res, expected)
1140
1141         # check the protobuf messages corresponding to the UDP query and answer
1142         msg = self.getFirstProtobufMessage()
1143         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
1144
1145         # then the response
1146         msg = self.getFirstProtobufMessage()
1147         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res)
1148         self.checkProtobufPolicy(msg, dnsmessage_pb2.PBDNSMessage.PolicyType.RESPONSEIP, 'zone.rpz.', '24.0.11.22.33.rpz-ip.', '33.22.11.99', dnsmessage_pb2.PBDNSMessage.PolicyKind.Custom)
1149         self.assertEqual(len(msg.response.rrs), 1)
1150         self.checkNoRemainingMessage()
1151
1152 class ProtobufRPZTagsTest(TestRecursorProtobuf):
1153     """
1154     This test makes sure that we correctly export the RPZ tags in our protobuf messages
1155     """
1156
1157     _confdir = 'ProtobufRPZTags'
1158     _config_template = """
1159 auth-zones=example=configs/%s/example.rpz.zone""" % _confdir
1160     _tags = ['tag1', 'tag2']
1161     _tags_from_gettag = ['tag1-from-gettag', 'tag2-from-gettag']
1162     _tags_from_rpz = ['tag1-from-rpz', 'tag2-from-rpz' ]
1163     _lua_config_file = """
1164     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=true, logResponses=true, tags={'tag1', 'tag2'} } )
1165     rpzFile('configs/%s/zone.rpz', { policyName="zone.rpz.", tags={ '%s', '%s'} })
1166     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port, _confdir, _tags_from_rpz[0], _tags_from_rpz[1])
1167     _lua_dns_script_file = """
1168     function gettag(remote, ednssubnet, localip, qname, qtype, ednsoptions, tcp)
1169       return 0, { '%s', '%s' }
1170     end
1171     function preresolve(dq)
1172       dq:addPolicyTag('%s')
1173       dq:addPolicyTag('%s')
1174       return false
1175     end
1176     """ % (_tags_from_gettag[0], _tags_from_gettag[1], _tags[0], _tags[1])
1177
1178     @classmethod
1179     def generateRecursorConfig(cls, confdir):
1180         authzonepath = os.path.join(confdir, 'example.rpz.zone')
1181         with open(authzonepath, 'w') as authzone:
1182             authzone.write("""$ORIGIN example.
1183 @ 3600 IN SOA {soa}
1184 sub.test 3600 IN A 192.0.2.42
1185 """.format(soa=cls._SOA))
1186
1187         rpzFilePath = os.path.join(confdir, 'zone.rpz')
1188         with open(rpzFilePath, 'w') as rpzZone:
1189             rpzZone.write("""$ORIGIN zone.rpz.
1190 @ 3600 IN SOA {soa}
1191 *.test.example.zone.rpz. 60 IN CNAME rpz-passthru.
1192 """.format(soa=cls._SOA))
1193
1194         super(ProtobufRPZTagsTest, cls).generateRecursorConfig(confdir)
1195
1196     def testA(self):
1197         name = 'sub.test.example.'
1198         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
1199         query = dns.message.make_query(name, 'A', want_dnssec=True)
1200         query.flags |= dns.flags.CD
1201         res = self.sendUDPQuery(query)
1202         self.assertRRsetInAnswer(res, expected)
1203
1204         # check the protobuf messages corresponding to the UDP query and answer
1205         msg = self.getFirstProtobufMessage()
1206         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
1207
1208         # then the response
1209         msg = self.getFirstProtobufMessage()
1210         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res)
1211         self.checkProtobufPolicy(msg, dnsmessage_pb2.PBDNSMessage.PolicyType.QNAME, 'zone.rpz.', '*.test.example.', 'sub.test.example', dnsmessage_pb2.PBDNSMessage.PolicyKind.NoAction)
1212         self.checkProtobufTags(msg, self._tags + self._tags_from_gettag + self._tags_from_rpz)
1213         self.assertEqual(len(msg.response.rrs), 1)
1214         rr = msg.response.rrs[0]
1215         # we have max-cache-ttl set to 15
1216         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
1217         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
1218         self.checkNoRemainingMessage()
1219
1220
1221 class ProtobufMetaFFITest(TestRecursorProtobuf):
1222     """
1223     This test makes sure that we can correctly add extra meta fields (FFI version).
1224     """
1225     _confdir = 'ProtobufMetaFFITest'
1226     _config_template = """
1227 auth-zones=example=configs/%s/example.zone""" % _confdir
1228     _lua_config_file = """
1229     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=true, logResponses=true } )
1230     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
1231     _lua_dns_script_file = """
1232     local ffi = require("ffi")
1233
1234     ffi.cdef[[
1235       typedef struct pdns_ffi_param pdns_ffi_param_t;
1236
1237       const char* pdns_ffi_param_get_qname(pdns_ffi_param_t* ref);
1238       void pdns_ffi_param_add_meta_single_string_kv(pdns_ffi_param_t *ref, const char* key, const char* val);
1239       void pdns_ffi_param_add_meta_single_int64_kv(pdns_ffi_param_t *ref, const char* key, int64_t val);
1240     ]]
1241
1242     function gettag_ffi(obj)
1243       qname = ffi.string(ffi.C.pdns_ffi_param_get_qname(obj))
1244       if qname == 'meta.example' then
1245         ffi.C.pdns_ffi_param_add_meta_single_string_kv(obj, "meta-str", "keyword")
1246         ffi.C.pdns_ffi_param_add_meta_single_int64_kv(obj, "meta-int", 42)
1247         ffi.C.pdns_ffi_param_add_meta_single_string_kv(obj, "meta-str", "content")
1248         ffi.C.pdns_ffi_param_add_meta_single_int64_kv(obj, "meta-int", 21)
1249       end
1250       return 0
1251     end
1252     """
1253     def testMeta(self):
1254         name = 'meta.example.'
1255         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.85')
1256         query = dns.message.make_query(name, 'A', want_dnssec=True)
1257         query.flags |= dns.flags.CD
1258         res = self.sendUDPQuery(query)
1259         self.assertRRsetInAnswer(res, expected)
1260
1261         # check the protobuf messages corresponding to the UDP query and answer
1262         msg = self.getFirstProtobufMessage()
1263         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
1264         self.checkProtobufMetas(msg, {'meta-str': { "stringVal" : ["content", "keyword"]}, 'meta-int': {"intVal" : [21, 42]}})
1265
1266         # then the response
1267         msg = self.getFirstProtobufMessage()
1268         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res)
1269         self.assertEqual(len(msg.response.rrs), 1)
1270         rr = msg.response.rrs[0]
1271         # we have max-cache-ttl set to 15
1272         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
1273         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.85')
1274         self.checkProtobufMetas(msg, {'meta-str': { "stringVal" : ["content", "keyword"]}, 'meta-int': {"intVal" : [21, 42]}})
1275
1276         self.checkNoRemainingMessage()