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 import clientsubnetoption
  10
  11 # Python2/3 compatibility hacks
  12 try:
  13   from queue import Queue
  14 except ImportError:
  15   from Queue import Queue
  16
  17 try:
  18   range = xrange
  19 except NameError:
  20   pass
  21
  22 from recursortests import RecursorTest
  23
  24 def ProtobufConnectionHandler(queue, conn):
  25     data = None
  26     while True:
  27         data = conn.recv(2)
  28         if not data:
  29             break
  30         (datalen,) = struct.unpack("!H", data)
  31         data = conn.recv(datalen)
  32         if not data:
  33             break
  34
  35         queue.put(data, True, timeout=2.0)
  36
  37     conn.close()
  38
  39 def ProtobufListener(queue, port):
  40     sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  41     sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
  42     try:
  43         sock.bind(("127.0.0.1", port))
  44     except socket.error as e:
  45         print("Error binding in the protobuf listener: %s" % str(e))
  46         sys.exit(1)
  47
  48     sock.listen(100)
  49     while True:
  50         try:
  51             (conn, _) = sock.accept()
  52             thread = threading.Thread(name='Connection Handler',
  53                                       target=ProtobufConnectionHandler,
  54                                       args=[queue, conn])
  55             thread.setDaemon(True)
  56             thread.start()
  57
  58         except socket.error as e:
  59             print('Error in protobuf socket: %s' % str(e))
  60
  61     sock.close()
  62
  63
  64 class ProtobufServerParams:
  65   def __init__(self, port):
  66     self.queue = Queue()
  67     self.port = port
  68
  69 protobufServersParameters = [ProtobufServerParams(4243), ProtobufServerParams(4244)]
  70 protobufListeners = []
  71 for param in protobufServersParameters:
  72   listener = threading.Thread(name='Protobuf Listener', target=ProtobufListener, args=[param.queue, param.port])
  73   listener.setDaemon(True)
  74   listener.start()
  75   protobufListeners.append(listener)
  76
  77 class TestRecursorProtobuf(RecursorTest):
  78
  79     _lua_config_file = """
  80     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"})
  81     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
  82
  83     _auth_zones = {
  84         '8': {'threads': 1,
  85               'zones': ['ROOT']},
  86         '9': {'threads': 1,
  87               'zones': ['secure.example', 'islandofsecurity.example']},
  88         '10': {'threads': 1,
  89                'zones': ['example']},
  90         '18': {'threads': 1,
  91                'zones': ['example']}
  92     }
  93
  94     def getFirstProtobufMessage(self, retries=1, waitTime=1):
  95         msg = None
  96
  97         #print("in getFirstProtobufMessage")
  98         for param in protobufServersParameters:
  99           #print(param.port)
 100           failed = 0
 101
 102           while param.queue.empty():
 103             #print(failed)
 104             #print(retries)
 105             if failed >= retries:
 106               break
 107
 108             failed = failed + 1
 109             #print("waiting")
 110             time.sleep(waitTime)
 111
 112           self.assertFalse(param.queue.empty())
 113           data = param.queue.get(False)
 114           self.assertTrue(data)
 115           oldmsg = msg
 116           msg = dnsmessage_pb2.PBDNSMessage()
 117           msg.ParseFromString(data)
 118           if oldmsg is not None:
 119             self.assertEqual(msg, oldmsg)
 120
 121         #print(msg)
 122         return msg
 123
 124     def emptyProtoBufQueue(self):
 125         for param in protobufServersParameters:
 126             while not param.queue.empty():
 127                 param.queue.get(False)
 128
 129     def checkNoRemainingMessage(self):
 130         for param in protobufServersParameters:
 131           self.assertTrue(param.queue.empty())
 132
 133     def checkProtobufBase(self, msg, protocol, query, initiator, normalQueryResponse=True, expectedECS=None, receivedSize=None):
 134         self.assertTrue(msg)
 135         self.assertTrue(msg.HasField('timeSec'))
 136         self.assertTrue(msg.HasField('socketFamily'))
 137         self.assertEqual(msg.socketFamily, dnsmessage_pb2.PBDNSMessage.INET)
 138         self.assertTrue(msg.HasField('from'))
 139         fromvalue = getattr(msg, 'from')
 140         self.assertEqual(socket.inet_ntop(socket.AF_INET, fromvalue), initiator)
 141         self.assertTrue(msg.HasField('socketProtocol'))
 142         self.assertEqual(msg.socketProtocol, protocol)
 143         self.assertTrue(msg.HasField('messageId'))
 144         self.assertTrue(msg.HasField('serverIdentity'))
 145         self.assertTrue(msg.HasField('id'))
 146         self.assertEqual(msg.id, query.id)
 147         self.assertTrue(msg.HasField('inBytes'))
 148         if normalQueryResponse:
 149             # compare inBytes with length of query/response
 150             # Note that for responses, the size we received might differ
 151             # because dnspython might compress labels differently from
 152             # the recursor
 153             if receivedSize:
 154                 self.assertEqual(msg.inBytes, receivedSize)
 155             else:
 156                 self.assertEqual(msg.inBytes, len(query.to_wire()))
 157         if expectedECS is not None:
 158             self.assertTrue(msg.HasField('originalRequestorSubnet'))
 159             # v4 only for now
 160             self.assertEqual(len(msg.originalRequestorSubnet), 4)
 161             self.assertEqual(socket.inet_ntop(socket.AF_INET, msg.originalRequestorSubnet), '127.0.0.1')
 162
 163     def checkOutgoingProtobufBase(self, msg, protocol, query, initiator, length=None, expectedECS=None):
 164         self.assertTrue(msg)
 165         self.assertTrue(msg.HasField('timeSec'))
 166         self.assertTrue(msg.HasField('socketFamily'))
 167         self.assertEqual(msg.socketFamily, dnsmessage_pb2.PBDNSMessage.INET)
 168         self.assertTrue(msg.HasField('socketProtocol'))
 169         self.assertEqual(msg.socketProtocol, protocol)
 170         self.assertTrue(msg.HasField('messageId'))
 171         self.assertTrue(msg.HasField('serverIdentity'))
 172         self.assertTrue(msg.HasField('id'))
 173         self.assertNotEqual(msg.id, query.id)
 174         self.assertTrue(msg.HasField('inBytes'))
 175         if length is not None:
 176           self.assertEqual(msg.inBytes, length)
 177         else:
 178           # compare inBytes with length of query/response
 179           self.assertEqual(msg.inBytes, len(query.to_wire()))
 180         if expectedECS is not None:
 181             self.assertTrue(msg.HasField('originalRequestorSubnet'))
 182             # v4 only for now
 183             self.assertEqual(len(msg.originalRequestorSubnet), 4)
 184             self.assertEqual(socket.inet_ntop(socket.AF_INET, msg.originalRequestorSubnet), expectedECS)
 185
 186     def checkProtobufQuery(self, msg, protocol, query, qclass, qtype, qname, initiator='127.0.0.1', to='127.0.0.1'):
 187         self.assertEqual(msg.type, dnsmessage_pb2.PBDNSMessage.DNSQueryType)
 188         self.checkProtobufBase(msg, protocol, query, initiator)
 189         # dnsdist doesn't fill the responder field for responses
 190         # because it doesn't keep the information around.
 191         self.assertTrue(msg.HasField('to'))
 192         self.assertEqual(socket.inet_ntop(socket.AF_INET, msg.to), to)
 193         self.assertTrue(msg.HasField('question'))
 194         self.assertTrue(msg.question.HasField('qClass'))
 195         self.assertEqual(msg.question.qClass, qclass)
 196         self.assertTrue(msg.question.HasField('qType'))
 197         self.assertEqual(msg.question.qClass, qtype)
 198         self.assertTrue(msg.question.HasField('qName'))
 199         self.assertEqual(msg.question.qName, qname)
 200
 201     def checkProtobufResponse(self, msg, protocol, response, initiator='127.0.0.1', receivedSize=None, vstate=dnsmessage_pb2.PBDNSMessage.VState.Indeterminate):
 202         self.assertEqual(msg.type, dnsmessage_pb2.PBDNSMessage.DNSResponseType)
 203         self.checkProtobufBase(msg, protocol, response, initiator, receivedSize=receivedSize)
 204         self.assertTrue(msg.HasField('response'))
 205         self.assertTrue(msg.response.HasField('queryTimeSec'))
 206         self.assertTrue(msg.response.HasField('validationState'))
 207         self.assertEqual(msg.response.validationState, vstate)
 208
 209     def checkProtobufResponseRecord(self, record, rclass, rtype, rname, rttl, checkTTL=True):
 210         self.assertTrue(record.HasField('class'))
 211         self.assertEqual(getattr(record, 'class'), rclass)
 212         self.assertTrue(record.HasField('type'))
 213         self.assertEqual(record.type, rtype)
 214         self.assertTrue(record.HasField('name'))
 215         self.assertEqual(record.name, rname)
 216         self.assertTrue(record.HasField('ttl'))
 217         if checkTTL:
 218             self.assertEqual(record.ttl, rttl)
 219         self.assertTrue(record.HasField('rdata'))
 220
 221     def checkProtobufPolicy(self, msg, policyType, reason, trigger, hit, kind):
 222         self.assertEqual(msg.type, dnsmessage_pb2.PBDNSMessage.DNSResponseType)
 223         self.assertTrue(msg.response.HasField('appliedPolicyType'))
 224         self.assertTrue(msg.response.HasField('appliedPolicy'))
 225         self.assertTrue(msg.response.HasField('appliedPolicyTrigger'))
 226         self.assertTrue(msg.response.HasField('appliedPolicyHit'))
 227         self.assertTrue(msg.response.HasField('appliedPolicyKind'))
 228         self.assertEqual(msg.response.appliedPolicy, reason)
 229         self.assertEqual(msg.response.appliedPolicyType, policyType)
 230         self.assertEqual(msg.response.appliedPolicyTrigger, trigger)
 231         self.assertEqual(msg.response.appliedPolicyHit, hit)
 232         self.assertEqual(msg.response.appliedPolicyKind, kind)
 233
 234     def checkProtobufTags(self, msg, tags):
 235         #print(tags)
 236         #print('---')
 237         #print(msg.response.tags)
 238         self.assertEqual(len(msg.response.tags), len(tags))
 239         for tag in msg.response.tags:
 240             self.assertTrue(tag in tags)
 241
 242     def checkProtobufMetas(self, msg, metas):
 243         #print(metas)
 244         #print('---')
 245         #print(msg.meta)
 246         self.assertEqual(len(msg.meta), len(metas))
 247         for m in msg.meta:
 248             self.assertTrue(m.HasField('key'))
 249             self.assertTrue(m.HasField('value'))
 250             self.assertTrue(m.key in metas)
 251             for i in m.value.intVal :
 252               self.assertTrue(i in metas[m.key]['intVal'])
 253             for s in m.value.stringVal :
 254               self.assertTrue(s in metas[m.key]['stringVal'])
 255
 256     def checkProtobufOutgoingQuery(self, msg, protocol, query, qclass, qtype, qname, initiator='127.0.0.1', length=None, expectedECS=None):
 257         self.assertEqual(msg.type, dnsmessage_pb2.PBDNSMessage.DNSOutgoingQueryType)
 258         self.checkOutgoingProtobufBase(msg, protocol, query, initiator, length=length, expectedECS=expectedECS)
 259         self.assertTrue(msg.HasField('to'))
 260         self.assertTrue(msg.HasField('question'))
 261         self.assertTrue(msg.question.HasField('qClass'))
 262         self.assertEqual(msg.question.qClass, qclass)
 263         self.assertTrue(msg.question.HasField('qType'))
 264         self.assertEqual(msg.question.qType, qtype)
 265         self.assertTrue(msg.question.HasField('qName'))
 266         self.assertEqual(msg.question.qName, qname)
 267
 268     def checkProtobufIncomingResponse(self, msg, protocol, response, initiator='127.0.0.1', length=None):
 269         self.assertEqual(msg.type, dnsmessage_pb2.PBDNSMessage.DNSIncomingResponseType)
 270         self.checkOutgoingProtobufBase(msg, protocol, response, initiator, length=length)
 271         self.assertTrue(msg.HasField('response'))
 272         self.assertTrue(msg.response.HasField('rcode'))
 273         self.assertTrue(msg.response.HasField('queryTimeSec'))
 274
 275     def checkProtobufIncomingNetworkErrorResponse(self, msg, protocol, response, initiator='127.0.0.1'):
 276         self.checkProtobufIncomingResponse(msg, protocol, response, initiator, length=0)
 277         self.assertEqual(msg.response.rcode, 65536)
 278
 279     def checkProtobufIdentity(self, msg, requestorId, deviceId, deviceName):
 280         #print(msg)
 281         self.assertTrue((requestorId == '') == (not msg.HasField('requestorId')))
 282         self.assertTrue((deviceId == b'') == (not msg.HasField('deviceId')))
 283         self.assertTrue((deviceName == '') == (not msg.HasField('deviceName')))
 284         self.assertEqual(msg.requestorId, requestorId)
 285         self.assertEqual(msg.deviceId, deviceId)
 286         self.assertEqual(msg.deviceName, deviceName)
 287
 288     def setUp(self):
 289         super(TestRecursorProtobuf, self).setUp()
 290         # Make sure the queue is empty, in case
 291         # a previous test failed
 292         self.emptyProtoBufQueue()
 293         # wait long enough to be sure that the housekeeping has
 294         # prime the root NS
 295         time.sleep(1)
 296
 297     @classmethod
 298     def generateRecursorConfig(cls, confdir):
 299         authzonepath = os.path.join(confdir, 'example.zone')
 300         with open(authzonepath, 'w') as authzone:
 301             authzone.write("""$ORIGIN example.
 302 @ 3600 IN SOA {soa}
 303 a 3600 IN A 192.0.2.42
 304 tagged 3600 IN A 192.0.2.84
 305 meta 3600 IN A 192.0.2.85
 306 query-selected 3600 IN A 192.0.2.84
 307 answer-selected 3600 IN A 192.0.2.84
 308 types 3600 IN A 192.0.2.84
 309 types 3600 IN AAAA 2001:DB8::1
 310 types 3600 IN TXT "Lorem ipsum dolor sit amet"
 311 types 3600 IN MX 10 a.example.
 312 types 3600 IN SPF "v=spf1 -all"
 313 types 3600 IN SRV 10 20 443 a.example.
 314 cname 3600 IN CNAME a.example.
 315
 316 """.format(soa=cls._SOA))
 317         super(TestRecursorProtobuf, cls).generateRecursorConfig(confdir)
 318
 319
 320 class ProtobufDefaultTest(TestRecursorProtobuf):
 321     """
 322     This test makes sure that we correctly export queries and response over protobuf.
 323     """
 324
 325     _confdir = 'ProtobufDefault'
 326     _config_template = """
 327 auth-zones=example=configs/%s/example.zone""" % _confdir
 328
 329     def testA(self):
 330         name = 'a.example.'
 331         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 332         query = dns.message.make_query(name, 'A', want_dnssec=True)
 333         query.flags |= dns.flags.CD
 334         res = self.sendUDPQuery(query)
 335
 336         self.assertRRsetInAnswer(res, expected)
 337
 338         # check the protobuf messages corresponding to the UDP query and answer
 339         msg = self.getFirstProtobufMessage()
 340         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
 341         # then the response
 342         msg = self.getFirstProtobufMessage()
 343         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res, '127.0.0.1')
 344         self.assertEqual(len(msg.response.rrs), 1)
 345         rr = msg.response.rrs[0]
 346         # we have max-cache-ttl set to 15
 347         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 348         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
 349         self.checkNoRemainingMessage()
 350
 351     def testCNAME(self):
 352         name = 'cname.example.'
 353         expectedCNAME = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'CNAME', 'a.example.')
 354         expectedA = dns.rrset.from_text('a.example.', 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 355         query = dns.message.make_query(name, 'A', want_dnssec=True)
 356         query.flags |= dns.flags.CD
 357         raw = self.sendUDPQuery(query, decode=False)
 358         res = dns.message.from_wire(raw)
 359         self.assertRRsetInAnswer(res, expectedCNAME)
 360         self.assertRRsetInAnswer(res, expectedA)
 361
 362         # check the protobuf messages corresponding to the UDP query and answer
 363         # but first let the protobuf messages the time to get there
 364         msg = self.getFirstProtobufMessage()
 365         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
 366         # then the response
 367         msg = self.getFirstProtobufMessage()
 368         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res, '127.0.0.1', receivedSize=len(raw))
 369         self.assertEqual(len(msg.response.rrs), 2)
 370         rr = msg.response.rrs[0]
 371         # we don't want to check the TTL for the A record, it has been cached by the previous test
 372         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.CNAME, name, 15)
 373         self.assertEqual(rr.rdata, b'a.example.')
 374         rr = msg.response.rrs[1]
 375         # we have max-cache-ttl set to 15
 376         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, 'a.example.', 15, checkTTL=False)
 377         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
 378         self.checkNoRemainingMessage()
 379
 380 class ProtobufProxyMappingTest(TestRecursorProtobuf):
 381     """
 382     This test makes sure that we correctly export queries and response over protobuf with a proxyMapping
 383     """
 384
 385     _confdir = 'ProtobufProxyMappingTest'
 386     _config_template = """
 387     auth-zones=example=configs/%s/example.zone
 388     allow-from=3.4.5.0/24
 389     """ % _confdir
 390
 391     _lua_config_file = """
 392     addProxyMapping("127.0.0.1/24", "3.4.5.6:99")
 393     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"})
 394     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 395
 396     def testA(self):
 397         name = 'a.example.'
 398         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 399         query = dns.message.make_query(name, 'A', want_dnssec=True)
 400         query.flags |= dns.flags.CD
 401         res = self.sendUDPQuery(query)
 402
 403         self.assertRRsetInAnswer(res, expected)
 404
 405         # check the protobuf messages corresponding to the UDP query and answer
 406         msg = self.getFirstProtobufMessage()
 407         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
 408         # then the response
 409         msg = self.getFirstProtobufMessage()
 410         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res, '127.0.0.1')
 411         self.assertEqual(len(msg.response.rrs), 1)
 412         rr = msg.response.rrs[0]
 413         # we have max-cache-ttl set to 15
 414         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 415         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
 416         self.checkNoRemainingMessage()
 417
 418 class ProtobufProxyMappingLogMappedTest(TestRecursorProtobuf):
 419     """
 420     This test makes sure that we correctly export queries and response over protobuf.
 421     """
 422
 423     _confdir = 'ProtobufProxyMappingLogMappedTest'
 424     _config_template = """
 425     auth-zones=example=configs/%s/example.zone
 426     allow-from=3.4.5.0/0"
 427     """ % _confdir
 428
 429     _lua_config_file = """
 430     addProxyMapping("127.0.0.1/24", "3.4.5.6:99")
 431     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logMappedFrom = true })
 432     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 433
 434     def testA(self):
 435         name = 'a.example.'
 436         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 437         query = dns.message.make_query(name, 'A', want_dnssec=True)
 438         query.flags |= dns.flags.CD
 439         res = self.sendUDPQuery(query)
 440
 441         self.assertRRsetInAnswer(res, expected)
 442
 443         # check the protobuf messages corresponding to the UDP query and answer
 444         msg = self.getFirstProtobufMessage()
 445         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name, '3.4.5.6')
 446         # then the response
 447         msg = self.getFirstProtobufMessage()
 448         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res, '3.4.5.6')
 449         self.assertEqual(len(msg.response.rrs), 1)
 450         rr = msg.response.rrs[0]
 451         # we have max-cache-ttl set to 15
 452         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 453         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
 454         self.checkNoRemainingMessage()
 455
 456 class ProtobufProxyTest(TestRecursorProtobuf):
 457     """
 458     This test makes sure that we correctly export addresses over protobuf when the proxy protocol is used.
 459     """
 460
 461     _confdir = 'ProtobufProxy'
 462     _config_template = """
 463 auth-zones=example=configs/%s/example.zone
 464 proxy-protocol-from=127.0.0.1/32
 465 allow-from=127.0.0.1,6.6.6.6
 466 """ % _confdir
 467
 468     def testA(self):
 469         name = 'a.example.'
 470         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 471         query = dns.message.make_query(name, 'A', want_dnssec=True)
 472         query.flags |= dns.flags.CD
 473         res = self.sendUDPQueryWithProxyProtocol(query, False, '6.6.6.6', '7.7.7.7', 666, 777)
 474
 475         self.assertRRsetInAnswer(res, expected)
 476
 477         # check the protobuf messages corresponding to the UDP query and answer
 478         msg = self.getFirstProtobufMessage()
 479         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name, '6.6.6.6', '7.7.7.7')
 480         # then the response
 481         msg = self.getFirstProtobufMessage()
 482         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res, '6.6.6.6')
 483         self.assertEqual(len(msg.response.rrs), 1)
 484         rr = msg.response.rrs[0]
 485         # we have max-cache-ttl set to 15
 486         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 487         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
 488         self.checkNoRemainingMessage()
 489
 490 class ProtobufProxyWithProxyByTableTest(TestRecursorProtobuf):
 491     """
 492     This test makes sure that we correctly export addresses over protobuf when the proxy protocol and a proxy table mapping is used
 493     """
 494
 495     _confdir = 'ProtobufProxyWithProxyByTable'
 496     _config_template = """
 497 auth-zones=example=configs/%s/example.zone
 498 proxy-protocol-from=127.0.0.1/32
 499 allow-from=3.4.5.6
 500 """ % _confdir
 501
 502     _lua_config_file = """
 503     addProxyMapping("6.6.6.6/24", "3.4.5.6:99")
 504     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"})
 505     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 506
 507     def testA(self):
 508         name = 'a.example.'
 509         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 510         query = dns.message.make_query(name, 'A', want_dnssec=True)
 511         query.flags |= dns.flags.CD
 512         res = self.sendUDPQueryWithProxyProtocol(query, False, '6.6.6.6', '7.7.7.7', 666, 777)
 513
 514         self.assertRRsetInAnswer(res, expected)
 515
 516         # check the protobuf messages corresponding to the UDP query and answer
 517         msg = self.getFirstProtobufMessage()
 518         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name, '6.6.6.6', '7.7.7.7')
 519         # then the response
 520         msg = self.getFirstProtobufMessage()
 521         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res, '6.6.6.6')
 522         self.assertEqual(len(msg.response.rrs), 1)
 523         rr = msg.response.rrs[0]
 524         # we have max-cache-ttl set to 15
 525         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 526         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
 527         self.checkNoRemainingMessage()
 528
 529 class ProtobufProxyWithProxyByTableLogMappedTest(TestRecursorProtobuf):
 530     """
 531     This test makes sure that we correctly export addresses over protobuf when the proxy protocol and a proxy table mapping is used
 532     """
 533
 534     _confdir = 'ProtobufProxyWithProxyByTableLogMapped'
 535     _config_template = """
 536 auth-zones=example=configs/%s/example.zone
 537 proxy-protocol-from=127.0.0.1/32
 538 allow-from=3.4.5.6
 539 """ % _confdir
 540
 541     _lua_config_file = """
 542     addProxyMapping("6.6.6.6/24", "3.4.5.6:99")
 543     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logMappedFrom = true })
 544     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 545
 546     def testA(self):
 547         name = 'a.example.'
 548         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 549         query = dns.message.make_query(name, 'A', want_dnssec=True)
 550         query.flags |= dns.flags.CD
 551         res = self.sendUDPQueryWithProxyProtocol(query, False, '6.6.6.6', '7.7.7.7', 666, 777)
 552
 553         self.assertRRsetInAnswer(res, expected)
 554
 555         # check the protobuf messages corresponding to the UDP query and answer
 556         msg = self.getFirstProtobufMessage()
 557         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name, '3.4.5.6', '7.7.7.7')
 558         # then the response
 559         msg = self.getFirstProtobufMessage()
 560         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res, '3.4.5.6')
 561         self.assertEqual(len(msg.response.rrs), 1)
 562         rr = msg.response.rrs[0]
 563         # we have max-cache-ttl set to 15
 564         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 565         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
 566         self.checkNoRemainingMessage()
 567
 568
 569 class OutgoingProtobufDefaultTest(TestRecursorProtobuf):
 570     """
 571     This test makes sure that we correctly export outgoing queries over protobuf.
 572     It must be improved and setup env so we can check for incoming responses, but makes sure for now
 573     that the recursor at least connects to the protobuf server.
 574     """
 575
 576     _confdir = 'OutgoingProtobufDefault'
 577     _config_template = """
 578     # Switch off QName Minimization, it generates much more protobuf messages
 579     # (or make the test much more smart!)
 580     qname-minimization=no
 581     max-cache-ttl=600
 582     loglevel=9
 583 """
 584     _lua_config_file = """
 585     outgoingProtobufServer({"127.0.0.1:%d", "127.0.0.1:%d"})
 586     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 587
 588     def testA(self):
 589         # There is a race in priming (having the . DNSKEY in cache in particular) and this code.
 590         # So make sure we have the . DNSKEY in cache
 591         query = dns.message.make_query('.', 'A', want_dnssec=True)
 592         query.flags |= dns.flags.RD
 593         res = self.sendUDPQuery(query)
 594         time.sleep(1)
 595         self.emptyProtoBufQueue()
 596
 597         name = 'host1.secure.example.'
 598         expected = list()
 599
 600         for qname, qtype, proto, responseSize in [
 601                 ('host1.secure.example.', dns.rdatatype.A, dnsmessage_pb2.PBDNSMessage.UDP, 248),
 602                 ('host1.secure.example.', dns.rdatatype.A, dnsmessage_pb2.PBDNSMessage.UDP, 221),
 603                 ('example.', dns.rdatatype.DNSKEY, dnsmessage_pb2.PBDNSMessage.UDP, 219),
 604                 ('host1.secure.example.', dns.rdatatype.A, dnsmessage_pb2.PBDNSMessage.UDP, 175),
 605                 ('secure.example.', dns.rdatatype.DNSKEY, dnsmessage_pb2.PBDNSMessage.UDP, 233),
 606         ]:
 607             if not qname:
 608                 expected.append((None, None, None, None, None, None))
 609                 continue
 610             query = dns.message.make_query(qname, qtype, use_edns=True, want_dnssec=True)
 611             resp = dns.message.make_response(query)
 612             expected.append((
 613                 qname, qtype, query, resp, proto, responseSize
 614             ))
 615
 616         query = dns.message.make_query(name, 'A', want_dnssec=True)
 617         query.flags |= dns.flags.RD
 618         res = self.sendUDPQuery(query)
 619
 620         for qname, qtype, qry, ans, proto, responseSize in expected:
 621             if not qname:
 622                 self.getFirstProtobufMessage()
 623                 self.getFirstProtobufMessage()
 624                 continue
 625
 626             msg = self.getFirstProtobufMessage()
 627             self.checkProtobufOutgoingQuery(msg, proto, qry, dns.rdataclass.IN, qtype, qname)
 628
 629             # Check the answer
 630             msg = self.getFirstProtobufMessage()
 631             self.checkProtobufIncomingResponse(msg, proto, ans, length=responseSize)
 632
 633         self.checkNoRemainingMessage()
 634
 635 class OutgoingProtobufWithECSMappingTest(TestRecursorProtobuf):
 636     """
 637     This test makes sure that we correctly export outgoing queries over protobuf.
 638     It must be improved and setup env so we can check for incoming responses, but makes sure for now
 639     that the recursor at least connects to the protobuf server.
 640     """
 641
 642     _confdir = 'OutgoingProtobuffWithECSMapping'
 643     _config_template = """
 644     # Switch off QName Minimization, it generates much more protobuf messages
 645     # (or make the test much more smart!)
 646     qname-minimization=no
 647     edns-subnet-allow-list=example
 648     allow-from=1.2.3.4/32
 649     # this is to not let . queries interfere
 650     max-cache-ttl=600
 651     loglevel=9
 652 """
 653     _lua_config_file = """
 654     outgoingProtobufServer({"127.0.0.1:%d", "127.0.0.1:%d"})
 655     addProxyMapping("127.0.0.0/8", "1.2.3.4", { "host1.secure.example." })
 656     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 657
 658     def testA(self):
 659         # There is a race in priming (having the . DNSKEY in cache in particular) and this code.
 660         # So make sure we have the . DNSKEY in cache
 661         query = dns.message.make_query('.', 'A', want_dnssec=True)
 662         query.flags |= dns.flags.RD
 663         res = self.sendUDPQuery(query)
 664         time.sleep(1)
 665         self.emptyProtoBufQueue()
 666
 667         name = 'host1.secure.example.'
 668         expected = list()
 669
 670         for qname, qtype, proto, responseSize, ecs in [
 671                 ('host1.secure.example.', dns.rdatatype.A, dnsmessage_pb2.PBDNSMessage.UDP, 248, "1.2.3.0"),
 672                 ('host1.secure.example.', dns.rdatatype.A, dnsmessage_pb2.PBDNSMessage.UDP, 221, "1.2.3.0"),
 673                 ('example.', dns.rdatatype.DNSKEY, dnsmessage_pb2.PBDNSMessage.UDP, 219, "1.2.3.0"),
 674                 ('host1.secure.example.', dns.rdatatype.A, dnsmessage_pb2.PBDNSMessage.UDP, 175, "1.2.3.0"),
 675                 ('secure.example.', dns.rdatatype.DNSKEY, dnsmessage_pb2.PBDNSMessage.UDP, 233, "1.2.3.0"),
 676         ]:
 677             if not qname:
 678                 expected.append((None, None, None, None, None, None, None))
 679                 continue
 680             ecso = clientsubnetoption.ClientSubnetOption('9.10.11.12', 24)
 681             query = dns.message.make_query(qname, qtype, use_edns=True, want_dnssec=True, options=[ecso], payload=512)
 682             resp = dns.message.make_response(query)
 683             expected.append((
 684                 qname, qtype, query, resp, proto, responseSize, ecs
 685             ))
 686
 687         query = dns.message.make_query(name, 'A', want_dnssec=True)
 688         query.flags |= dns.flags.RD
 689         res = self.sendUDPQuery(query)
 690
 691         for qname, qtype, qry, ans, proto, responseSize, ecs in expected:
 692             if not qname:
 693                 self.getFirstProtobufMessage()
 694                 self.getFirstProtobufMessage()
 695                 continue
 696
 697             msg = self.getFirstProtobufMessage()
 698             self.checkProtobufOutgoingQuery(msg, proto, qry, dns.rdataclass.IN, qtype, qname, "127.0.0.1", None, ecs)
 699             # Check the answer
 700             msg = self.getFirstProtobufMessage()
 701             self.checkProtobufIncomingResponse(msg, proto, ans, length=responseSize)
 702
 703         self.checkNoRemainingMessage()
 704
 705         # this query should use the unmapped ECS
 706         name = 'mx1.secure.example.'
 707         expected = list()
 708
 709         for qname, qtype, proto, responseSize, ecs in [
 710                 ('mx1.secure.example.', dns.rdatatype.A, dnsmessage_pb2.PBDNSMessage.UDP, 173, "127.0.0.1"),
 711         ]:
 712             if not qname:
 713                 expected.append((None, None, None, None, None, None, None))
 714                 continue
 715             ecso = clientsubnetoption.ClientSubnetOption('127.0.0.1', 32)
 716             query = dns.message.make_query(qname, qtype, use_edns=True, want_dnssec=True, options=[ecso], payload=512)
 717             resp = dns.message.make_response(query)
 718             expected.append((
 719                 qname, qtype, query, resp, proto, responseSize, ecs
 720             ))
 721
 722         query = dns.message.make_query(name, 'A', want_dnssec=True)
 723         query.flags |= dns.flags.RD
 724         res = self.sendUDPQuery(query)
 725
 726         for qname, qtype, qry, ans, proto, responseSize, ecs in expected:
 727             if not qname:
 728                 self.getFirstProtobufMessage()
 729                 self.getFirstProtobufMessage()
 730                 continue
 731
 732             msg = self.getFirstProtobufMessage()
 733             self.checkProtobufOutgoingQuery(msg, proto, qry, dns.rdataclass.IN, qtype, qname, "127.0.0.1", None, ecs)
 734             # Check the answer
 735             msg = self.getFirstProtobufMessage()
 736             self.checkProtobufIncomingResponse(msg, proto, ans, length=responseSize)
 737
 738         self.checkNoRemainingMessage()
 739
 740 class OutgoingProtobufNoQueriesTest(TestRecursorProtobuf):
 741     """
 742     This test makes sure that we correctly export incoming responses but not outgoing queries over protobuf.
 743     It must be improved and setup env so we can check for incoming responses, but makes sure for now
 744     that the recursor at least connects to the protobuf server.
 745     """
 746
 747     _confdir = 'OutgoingProtobufNoQueries'
 748     _config_template = """
 749     # Switch off QName Minimization, it generates much more protobuf messages
 750     # (or make the test much more smart!)
 751     qname-minimization=no
 752     max-cache-ttl=600
 753     loglevel=9
 754 """
 755     _lua_config_file = """
 756     outgoingProtobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=false, logResponses=true })
 757     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 758
 759     def testA(self):
 760         # There is a race in priming (having the . DNSKEY in cache in particular) and this code.
 761         # So make sure we have the . DNSKEY in cache
 762         query = dns.message.make_query('.', 'A', want_dnssec=True)
 763         query.flags |= dns.flags.RD
 764         res = self.sendUDPQuery(query)
 765         time.sleep(1)
 766         self.emptyProtoBufQueue()
 767
 768         name = 'host1.secure.example.'
 769         expected = list()
 770         # the root DNSKEY has been learned with priming the root NS already
 771         # ('.', dns.rdatatype.DNSKEY, dnsmessage_pb2.PBDNSMessage.UDP, 201),
 772         for qname, qtype, proto, size in [
 773                 ('host1.secure.example.', dns.rdatatype.A, dnsmessage_pb2.PBDNSMessage.UDP, 248),
 774                 ('host1.secure.example.', dns.rdatatype.A, dnsmessage_pb2.PBDNSMessage.UDP, 221),
 775                 ('example.', dns.rdatatype.DNSKEY, dnsmessage_pb2.PBDNSMessage.UDP, 219),
 776                 ('host1.secure.example.', dns.rdatatype.A, dnsmessage_pb2.PBDNSMessage.UDP, 175),
 777                 ('secure.example.', dns.rdatatype.DNSKEY, dnsmessage_pb2.PBDNSMessage.UDP, 233),
 778         ]:
 779             if not qname:
 780                 expected.append((None, None, None, None, None, None))
 781                 continue
 782             query = dns.message.make_query(qname, qtype, use_edns=True, want_dnssec=True)
 783             resp = dns.message.make_response(query)
 784             expected.append((
 785                 qname, qtype, query, resp, proto, size
 786             ))
 787
 788         query = dns.message.make_query(name, 'A', want_dnssec=True)
 789         query.flags |= dns.flags.RD
 790         res = self.sendUDPQuery(query)
 791
 792         for qname, qtype, qry, ans, proto, size in expected:
 793             if not qname:
 794                 self.getFirstProtobufMessage()
 795                 continue
 796
 797             # check the response
 798             msg = self.getFirstProtobufMessage()
 799             self.checkProtobufIncomingResponse(msg, proto, ans, length=size)
 800
 801         self.checkNoRemainingMessage()
 802
 803 class ProtobufMasksTest(TestRecursorProtobuf):
 804     """
 805     This test makes sure that we correctly export queries and response over protobuf, respecting the configured initiator masking.
 806     """
 807
 808     _confdir = 'ProtobufMasks'
 809     _config_template = """
 810 auth-zones=example=configs/%s/example.zone""" % _confdir
 811     _protobufMaskV4 = 4
 812     _protobufMaskV6 = 128
 813     _lua_config_file = """
 814     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"})
 815     setProtobufMasks(%d, %d)
 816     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port, _protobufMaskV4, _protobufMaskV6)
 817
 818     def testA(self):
 819         name = 'a.example.'
 820         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 821         query = dns.message.make_query(name, 'A', want_dnssec=True)
 822         query.flags |= dns.flags.CD
 823         res = self.sendUDPQuery(query)
 824         self.assertRRsetInAnswer(res, expected)
 825
 826         # check the protobuf messages corresponding to the UDP query and answer
 827         # but first let the protobuf messages the time to get there
 828         msg = self.getFirstProtobufMessage()
 829         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name, '112.0.0.0')
 830         # then the response
 831         msg = self.getFirstProtobufMessage()
 832         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res, '112.0.0.0')
 833         self.assertEqual(len(msg.response.rrs), 1)
 834         rr = msg.response.rrs[0]
 835         # we have max-cache-ttl set to 15
 836         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 837         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
 838         self.checkNoRemainingMessage()
 839
 840 class ProtobufQueriesOnlyTest(TestRecursorProtobuf):
 841     """
 842     This test makes sure that we correctly export queries but not responses over protobuf.
 843     """
 844
 845     _confdir = 'ProtobufQueriesOnly'
 846     _config_template = """
 847 auth-zones=example=configs/%s/example.zone""" % _confdir
 848     _lua_config_file = """
 849     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=true, logResponses=false } )
 850     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 851
 852     def testA(self):
 853         name = 'a.example.'
 854         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 855         query = dns.message.make_query(name, 'A', want_dnssec=True)
 856         query.flags |= dns.flags.CD
 857         res = self.sendUDPQuery(query)
 858         self.assertRRsetInAnswer(res, expected)
 859
 860         # check the protobuf message corresponding to the UDP query
 861         msg = self.getFirstProtobufMessage()
 862         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
 863         # no response
 864         self.checkNoRemainingMessage()
 865
 866 class ProtobufResponsesOnlyTest(TestRecursorProtobuf):
 867     """
 868     This test makes sure that we correctly export responses but not queries over protobuf.
 869     """
 870
 871     _confdir = 'ProtobufResponsesOnly'
 872     _config_template = """
 873 auth-zones=example=configs/%s/example.zone""" % _confdir
 874     _lua_config_file = """
 875     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=false, logResponses=true } )
 876     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 877
 878     def testA(self):
 879         name = 'a.example.'
 880         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 881         query = dns.message.make_query(name, 'A', want_dnssec=True)
 882         query.flags |= dns.flags.CD
 883         res = self.sendUDPQuery(query)
 884         self.assertRRsetInAnswer(res, expected)
 885
 886         # check the protobuf message corresponding to the UDP response
 887         msg = self.getFirstProtobufMessage()
 888         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res)
 889         self.assertEqual(len(msg.response.rrs), 1)
 890         rr = msg.response.rrs[0]
 891         # we have max-cache-ttl set to 15
 892         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 893         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
 894         # nothing else in the queue
 895         self.checkNoRemainingMessage()
 896
 897 class ProtobufTaggedOnlyTest(TestRecursorProtobuf):
 898     """
 899     This test makes sure that we correctly export queries and responses but only if they have been tagged.
 900     """
 901
 902     _confdir = 'ProtobufTaggedOnly'
 903     _config_template = """
 904 auth-zones=example=configs/%s/example.zone""" % _confdir
 905     _lua_config_file = """
 906     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=true, logResponses=true, taggedOnly=true } )
 907     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 908     _tags = ['tag1', 'tag2']
 909     _tag_from_gettag = 'tag-from-gettag'
 910     _lua_dns_script_file = """
 911     function gettag(remote, ednssubnet, localip, qname, qtype, ednsoptions, tcp)
 912       if qname:equal('tagged.example.') then
 913         return 0, { '%s' }
 914       end
 915       return 0
 916     end
 917     function preresolve(dq)
 918       if dq.qname:equal('tagged.example.') then
 919         dq:addPolicyTag('%s')
 920         dq:addPolicyTag('%s')
 921       end
 922       return false
 923     end
 924     """ % (_tag_from_gettag, _tags[0], _tags[1])
 925
 926     def testA(self):
 927         name = 'a.example.'
 928         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 929         query = dns.message.make_query(name, 'A', want_dnssec=True)
 930         query.flags |= dns.flags.CD
 931         res = self.sendUDPQuery(query)
 932         self.assertRRsetInAnswer(res, expected)
 933
 934         # check the protobuf message corresponding to the UDP response
 935         # the first query and answer are not tagged, so there is nothing in the queue
 936         time.sleep(1)
 937         self.checkNoRemainingMessage()
 938
 939     def testTagged(self):
 940         name = 'tagged.example.'
 941         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.84')
 942         query = dns.message.make_query(name, 'A', want_dnssec=True)
 943         query.flags |= dns.flags.CD
 944         res = self.sendUDPQuery(query)
 945         self.assertRRsetInAnswer(res, expected)
 946
 947         # check the protobuf messages corresponding to the UDP query and answer
 948         msg = self.getFirstProtobufMessage()
 949         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
 950         self.checkProtobufTags(msg, [ self._tag_from_gettag ])
 951         # then the response
 952         msg = self.getFirstProtobufMessage()
 953         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res)
 954         self.assertEqual(len(msg.response.rrs), 1)
 955         rr = msg.response.rrs[0]
 956         # we have max-cache-ttl set to 15
 957         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 958         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.84')
 959         tags = [ self._tag_from_gettag ] + self._tags
 960         #print(msg)
 961         self.checkProtobufTags(msg, tags)
 962         self.checkNoRemainingMessage()
 963
 964         # Again to check PC case
 965         res = self.sendUDPQuery(query)
 966         self.assertRRsetInAnswer(res, expected)
 967
 968         # check the protobuf messages corresponding to the UDP query and answer
 969         msg = self.getFirstProtobufMessage()
 970         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
 971         self.checkProtobufTags(msg, [ self._tag_from_gettag ])
 972         # then the response
 973         msg = self.getFirstProtobufMessage()
 974         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res)
 975         self.assertEqual(len(msg.response.rrs), 1)
 976         rr = msg.response.rrs[0]
 977         # time may have passed, so do not check TTL
 978         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15, checkTTL=False)
 979         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.84')
 980         tags = [ self._tag_from_gettag ] + self._tags
 981         self.checkProtobufTags(msg, tags)
 982         self.checkNoRemainingMessage()
 983
 984 class ProtobufTagCacheTest(TestRecursorProtobuf):
 985     """
 986     This test makes sure that we correctly cache tags (actually not cache them)
 987     """
 988
 989     _confdir = 'ProtobufTagCache'
 990     _config_template = """
 991 auth-zones=example=configs/%s/example.zone""" % _confdir
 992     _lua_config_file = """
 993     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=false, logResponses=true } )
 994     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 995     _lua_dns_script_file = """
 996     function gettag(remote, ednssubnet, localip, qname, qtype, ednsoptions, tcp)
 997       if qname:equal('tagged.example.') then
 998         return 0, { '' .. math.random() }
 999       end
1000       return 0
1001     end
1002     """
1003
1004     def testTagged(self):
1005         name = 'tagged.example.'
1006         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.84')
1007         query = dns.message.make_query(name, 'A', want_dnssec=True)
1008         query.flags |= dns.flags.CD
1009         res = self.sendUDPQuery(query)
1010         self.assertRRsetInAnswer(res, expected)
1011
1012         msg = self.getFirstProtobufMessage()
1013         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res)
1014         self.assertEqual(len(msg.response.rrs), 1)
1015         rr = msg.response.rrs[0]
1016         # we have max-cache-ttl set to 15
1017         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
1018         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.84')
1019         self.checkNoRemainingMessage()
1020         self.assertEqual(len(msg.response.tags), 1)
1021         ts1 = msg.response.tags[0]
1022
1023         # Again to check PC case
1024         res = self.sendUDPQuery(query)
1025         self.assertRRsetInAnswer(res, expected)
1026
1027         msg = self.getFirstProtobufMessage()
1028         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res)
1029         self.assertEqual(len(msg.response.rrs), 1)
1030         rr = msg.response.rrs[0]
1031         # time may have passed, so do not check TTL
1032         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15, checkTTL=False)
1033         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.84')
1034         self.checkNoRemainingMessage()
1035         self.assertEqual(len(msg.response.tags), 1)
1036         ts2 = msg.response.tags[0]
1037         self.assertNotEqual(ts1, ts2)
1038
1039 class ProtobufSelectedFromLuaTest(TestRecursorProtobuf):
1040     """
1041     This test makes sure that we correctly export queries and responses but only if they have been selected from Lua.
1042     """
1043
1044     _confdir = 'ProtobufSelectedFromLua'
1045     _config_template = """
1046 auth-zones=example=configs/%s/example.zone""" % _confdir
1047     _lua_config_file = """
1048     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=false, logResponses=false } )
1049     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
1050     _lua_dns_script_file = """
1051     local ffi = require("ffi")
1052
1053     ffi.cdef[[
1054       typedef struct pdns_ffi_param pdns_ffi_param_t;
1055
1056       const char* pdns_ffi_param_get_qname(pdns_ffi_param_t* ref);
1057       void pdns_ffi_param_set_log_query(pdns_ffi_param_t* ref, bool logQuery);
1058     ]]
1059
1060     function gettag_ffi(obj)
1061       qname = ffi.string(ffi.C.pdns_ffi_param_get_qname(obj))
1062       if qname == 'query-selected.example' then
1063         ffi.C.pdns_ffi_param_set_log_query(obj, true)
1064       end
1065       return 0
1066     end
1067
1068     function preresolve(dq)
1069       if dq.qname:equal('answer-selected.example.') then
1070         dq.logResponse = true
1071       end
1072       return false
1073     end
1074     """
1075
1076     def testA(self):
1077         name = 'a.example.'
1078         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
1079         query = dns.message.make_query(name, 'A', want_dnssec=True)
1080         query.flags |= dns.flags.CD
1081         res = self.sendUDPQuery(query)
1082         self.assertRRsetInAnswer(res, expected)
1083
1084         # check the protobuf message corresponding to the UDP response
1085         # the first query and answer are not selected, so there is nothing in the queue
1086         self.checkNoRemainingMessage()
1087
1088     def testQuerySelected(self):
1089         name = 'query-selected.example.'
1090         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.84')
1091         query = dns.message.make_query(name, 'A', want_dnssec=True)
1092         query.flags |= dns.flags.CD
1093         res = self.sendUDPQuery(query)
1094         self.assertRRsetInAnswer(res, expected)
1095
1096         # check the protobuf messages corresponding to the UDP query
1097         msg = self.getFirstProtobufMessage()
1098         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
1099         # there should be no response
1100         self.checkNoRemainingMessage()
1101
1102     def testResponseSelected(self):
1103         name = 'answer-selected.example.'
1104         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.84')
1105         query = dns.message.make_query(name, 'A', want_dnssec=True)
1106         query.flags |= dns.flags.CD
1107         res = self.sendUDPQuery(query)
1108         self.assertRRsetInAnswer(res, expected)
1109
1110         # check the protobuf messages corresponding to the UDP response
1111         msg = self.getFirstProtobufMessage()
1112         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res)
1113         self.assertEqual(len(msg.response.rrs), 1)
1114         rr = msg.response.rrs[0]
1115         # we have max-cache-ttl set to 15
1116         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
1117         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.84')
1118         self.checkNoRemainingMessage()
1119
1120 class ProtobufExportTypesTest(TestRecursorProtobuf):
1121     """
1122     This test makes sure that we correctly export other types than A, AAAA and CNAME over protobuf.
1123     """
1124
1125     _confdir = 'ProtobufExportTypes'
1126     _config_template = """
1127 auth-zones=example=configs/%s/example.zone""" % _confdir
1128     _lua_config_file = """
1129     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { exportTypes={"AAAA", "MX", "SPF", "SRV", "TXT"} } )
1130     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
1131
1132     def testA(self):
1133         name = 'types.example.'
1134         expected = [dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.84'),
1135                     dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'AAAA', '2001:DB8::1'),
1136                     dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'MX', '10 a.example.'),
1137                     dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'SPF', '"v=spf1 -all"'),
1138                     dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'SRV', '10 20 443 a.example.'),
1139                     dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'TXT', '"Lorem ipsum dolor sit amet"'),
1140                    ]
1141         query = dns.message.make_query(name, 'ANY', want_dnssec=True)
1142         query.flags |= dns.flags.CD
1143         raw = self.sendUDPQuery(query, decode=False)
1144         res = dns.message.from_wire(raw)
1145
1146         for rrset in expected:
1147             self.assertRRsetInAnswer(res, rrset)
1148
1149         # check the protobuf messages corresponding to the UDP query and answer
1150         msg = self.getFirstProtobufMessage()
1151         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
1152         # then the response
1153         msg = self.getFirstProtobufMessage()
1154         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res, '127.0.0.1', receivedSize=len(raw))
1155         self.assertEqual(len(msg.response.rrs), 5)
1156         for rr in msg.response.rrs:
1157             self.assertTrue(rr.type in [dns.rdatatype.AAAA, dns.rdatatype.TXT, dns.rdatatype.MX, dns.rdatatype.SPF, dns.rdatatype.SRV])
1158
1159             if rr.type == dns.rdatatype.AAAA:
1160                 self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.AAAA, name, 15)
1161                 self.assertEqual(socket.inet_ntop(socket.AF_INET6, rr.rdata), '2001:db8::1')
1162             elif rr.type == dns.rdatatype.TXT:
1163                 self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.TXT, name, 15)
1164                 self.assertEqual(rr.rdata, b'"Lorem ipsum dolor sit amet"')
1165             elif rr.type == dns.rdatatype.MX:
1166                 self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.MX, name, 15)
1167                 self.assertEqual(rr.rdata, b'a.example.')
1168             elif rr.type == dns.rdatatype.SPF:
1169                 self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.SPF, name, 15)
1170                 self.assertEqual(rr.rdata, b'"v=spf1 -all"')
1171             elif rr.type == dns.rdatatype.SRV:
1172                 self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.SRV, name, 15)
1173                 self.assertEqual(rr.rdata, b'a.example.')
1174
1175         self.checkNoRemainingMessage()
1176
1177 class ProtobufTaggedExtraFieldsTest(TestRecursorProtobuf):
1178     """
1179     This test makes sure that we correctly export extra fields that may have been set while being tagged.
1180     """
1181
1182     _confdir = 'ProtobufTaggedExtraFields'
1183     _config_template = """
1184 auth-zones=example=configs/%s/example.zone""" % _confdir
1185     _lua_config_file = """
1186     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=true, logResponses=true } )
1187     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
1188     _requestorId = 'S-000001727'
1189     _deviceId = 'd1:0a:91:dc:cc:82'
1190     _deviceName = 'Joe'
1191     _lua_dns_script_file = """
1192     function gettag(remote, ednssubnet, localip, qname, qtype, ednsoptions, tcp)
1193       if qname:equal('tagged.example.') then
1194         -- tag number, policy tags, data, requestorId, deviceId, deviceName
1195         return 0, {}, {}, '%s', '%s', '%s'
1196       end
1197       return 0
1198     end
1199     """ % (_requestorId, _deviceId, _deviceName)
1200
1201     def testA(self):
1202         name = 'a.example.'
1203         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
1204         query = dns.message.make_query(name, 'A', want_dnssec=True)
1205         query.flags |= dns.flags.CD
1206         res = self.sendUDPQuery(query)
1207         self.assertRRsetInAnswer(res, expected)
1208
1209         # check the protobuf message corresponding to the UDP response
1210         # the first query and answer are not tagged, so there is nothing in the queue
1211         # check the protobuf messages corresponding to the UDP query and answer
1212         msg = self.getFirstProtobufMessage()
1213         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
1214         self.checkProtobufIdentity(msg, '', b'', '')
1215
1216         # then the response
1217         msg = self.getFirstProtobufMessage()
1218         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res, '127.0.0.1')
1219         self.assertEqual(len(msg.response.rrs), 1)
1220         rr = msg.response.rrs[0]
1221         # we have max-cache-ttl set to 15
1222         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
1223         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
1224         self.checkProtobufIdentity(msg, '', b'', '')
1225         self.checkNoRemainingMessage()
1226
1227     def testTagged(self):
1228         name = 'tagged.example.'
1229         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.84')
1230         query = dns.message.make_query(name, 'A', want_dnssec=True)
1231         query.flags |= dns.flags.CD
1232         res = self.sendUDPQuery(query)
1233         self.assertRRsetInAnswer(res, expected)
1234
1235         # check the protobuf messages corresponding to the UDP query and answer
1236         msg = self.getFirstProtobufMessage()
1237         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
1238         self.checkProtobufIdentity(msg, self._requestorId, self._deviceId.encode('ascii'), self._deviceName)
1239
1240         # then the response
1241         msg = self.getFirstProtobufMessage()
1242         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res)
1243         self.assertEqual(len(msg.response.rrs), 1)
1244         rr = msg.response.rrs[0]
1245         # we have max-cache-ttl set to 15
1246         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
1247         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.84')
1248         self.checkProtobufIdentity(msg, self._requestorId, self._deviceId.encode('ascii'), self._deviceName)
1249         self.checkNoRemainingMessage()
1250
1251 class ProtobufTaggedExtraFieldsFFITest(ProtobufTaggedExtraFieldsTest):
1252     """
1253     This test makes sure that we correctly export extra fields that may have been set while being tagged (FFI version).
1254     """
1255     _confdir = 'ProtobufTaggedExtraFieldsFFI'
1256     _config_template = """
1257 auth-zones=example=configs/%s/example.zone""" % _confdir
1258     _lua_config_file = """
1259     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=true, logResponses=true } )
1260     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
1261     _lua_dns_script_file = """
1262     local ffi = require("ffi")
1263
1264     ffi.cdef[[
1265       typedef struct pdns_ffi_param pdns_ffi_param_t;
1266
1267       const char* pdns_ffi_param_get_qname(pdns_ffi_param_t* ref);
1268       void pdns_ffi_param_set_tag(pdns_ffi_param_t* ref, unsigned int tag);
1269       void pdns_ffi_param_set_requestorid(pdns_ffi_param_t* ref, const char* name);
1270       void pdns_ffi_param_set_devicename(pdns_ffi_param_t* ref, const char* name);
1271       void pdns_ffi_param_set_deviceid(pdns_ffi_param_t* ref, size_t len, const void* name);
1272     ]]
1273
1274     function gettag_ffi(obj)
1275       qname = ffi.string(ffi.C.pdns_ffi_param_get_qname(obj))
1276       if qname == 'tagged.example' then
1277         ffi.C.pdns_ffi_param_set_requestorid(obj, "%s")
1278         deviceid = "%s"
1279         ffi.C.pdns_ffi_param_set_deviceid(obj, string.len(deviceid), deviceid)
1280         ffi.C.pdns_ffi_param_set_devicename(obj, "%s")
1281       end
1282       return 0
1283     end
1284     """ % (ProtobufTaggedExtraFieldsTest._requestorId, ProtobufTaggedExtraFieldsTest._deviceId, ProtobufTaggedExtraFieldsTest._deviceName)
1285
1286 class ProtobufRPZTest(TestRecursorProtobuf):
1287     """
1288     This test makes sure that we correctly export the RPZ applied policy in our protobuf messages
1289     """
1290
1291     _confdir = 'ProtobufRPZ'
1292     _config_template = """
1293 auth-zones=example=configs/%s/example.rpz.zone""" % _confdir
1294     _lua_config_file = """
1295     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=true, logResponses=true } )
1296     rpzFile('configs/%s/zone.rpz', { policyName="zone.rpz."})
1297     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port, _confdir)
1298
1299     @classmethod
1300     def generateRecursorConfig(cls, confdir):
1301         authzonepath = os.path.join(confdir, 'example.rpz.zone')
1302         with open(authzonepath, 'w') as authzone:
1303             authzone.write("""$ORIGIN example.
1304 @ 3600 IN SOA {soa}
1305 sub.test 3600 IN A 192.0.2.42
1306 ip  3600 IN A 33.22.11.99
1307 """.format(soa=cls._SOA))
1308
1309         rpzFilePath = os.path.join(confdir, 'zone.rpz')
1310         with open(rpzFilePath, 'w') as rpzZone:
1311             rpzZone.write("""$ORIGIN zone.rpz.
1312 @ 3600 IN SOA {soa}
1313 *.test.example.zone.rpz. 60 IN CNAME rpz-passthru.
1314 24.0.11.22.33.rpz-ip     60 IN A 1.2.3.4
1315 """.format(soa=cls._SOA))
1316
1317         super(ProtobufRPZTest, cls).generateRecursorConfig(confdir)
1318
1319     def testA(self):
1320         name = 'sub.test.example.'
1321         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
1322         query = dns.message.make_query(name, 'A', want_dnssec=True)
1323         query.flags |= dns.flags.CD
1324         res = self.sendUDPQuery(query)
1325         self.assertRRsetInAnswer(res, expected)
1326
1327         # check the protobuf messages corresponding to the UDP query and answer
1328         msg = self.getFirstProtobufMessage()
1329         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
1330
1331         # then the response
1332         msg = self.getFirstProtobufMessage()
1333         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res)
1334         self.checkProtobufPolicy(msg, dnsmessage_pb2.PBDNSMessage.PolicyType.QNAME, 'zone.rpz.', '*.test.example.', 'sub.test.example', dnsmessage_pb2.PBDNSMessage.PolicyKind.NoAction)
1335         self.assertEqual(len(msg.response.rrs), 1)
1336         rr = msg.response.rrs[0]
1337         # we have max-cache-ttl set to 15
1338         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
1339         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
1340         self.checkNoRemainingMessage()
1341
1342     def testB(self):
1343         name = 'ip.example.'
1344         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '1.2.3.4')
1345         query = dns.message.make_query(name, 'A', want_dnssec=True)
1346         query.flags |= dns.flags.CD
1347         res = self.sendUDPQuery(query)
1348         self.assertRRsetInAnswer(res, expected)
1349
1350         # check the protobuf messages corresponding to the UDP query and answer
1351         msg = self.getFirstProtobufMessage()
1352         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
1353
1354         # then the response
1355         msg = self.getFirstProtobufMessage()
1356         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res)
1357         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)
1358         self.assertEqual(len(msg.response.rrs), 1)
1359         self.checkNoRemainingMessage()
1360
1361 class ProtobufRPZTagsTest(TestRecursorProtobuf):
1362     """
1363     This test makes sure that we correctly export the RPZ tags in our protobuf messages
1364     """
1365
1366     _confdir = 'ProtobufRPZTags'
1367     _config_template = """
1368 auth-zones=example=configs/%s/example.rpz.zone""" % _confdir
1369     _tags = ['tag1', 'tag2']
1370     _tags_from_gettag = ['tag1-from-gettag', 'tag2-from-gettag']
1371     _tags_from_rpz = ['tag1-from-rpz', 'tag2-from-rpz' ]
1372     _lua_config_file = """
1373     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=true, logResponses=true, tags={'tag1', 'tag2'} } )
1374     rpzFile('configs/%s/zone.rpz', { policyName="zone.rpz.", tags={ '%s', '%s'} })
1375     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port, _confdir, _tags_from_rpz[0], _tags_from_rpz[1])
1376     _lua_dns_script_file = """
1377     function gettag(remote, ednssubnet, localip, qname, qtype, ednsoptions, tcp)
1378       return 0, { '%s', '%s' }
1379     end
1380     function preresolve(dq)
1381       dq:addPolicyTag('%s')
1382       dq:addPolicyTag('%s')
1383       return false
1384     end
1385     """ % (_tags_from_gettag[0], _tags_from_gettag[1], _tags[0], _tags[1])
1386
1387     @classmethod
1388     def generateRecursorConfig(cls, confdir):
1389         authzonepath = os.path.join(confdir, 'example.rpz.zone')
1390         with open(authzonepath, 'w') as authzone:
1391             authzone.write("""$ORIGIN example.
1392 @ 3600 IN SOA {soa}
1393 sub.test 3600 IN A 192.0.2.42
1394 """.format(soa=cls._SOA))
1395
1396         rpzFilePath = os.path.join(confdir, 'zone.rpz')
1397         with open(rpzFilePath, 'w') as rpzZone:
1398             rpzZone.write("""$ORIGIN zone.rpz.
1399 @ 3600 IN SOA {soa}
1400 *.test.example.zone.rpz. 60 IN CNAME rpz-passthru.
1401 """.format(soa=cls._SOA))
1402
1403         super(ProtobufRPZTagsTest, cls).generateRecursorConfig(confdir)
1404
1405     def testA(self):
1406         name = 'sub.test.example.'
1407         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
1408         query = dns.message.make_query(name, 'A', want_dnssec=True)
1409         query.flags |= dns.flags.CD
1410         res = self.sendUDPQuery(query)
1411         self.assertRRsetInAnswer(res, expected)
1412
1413         # check the protobuf messages corresponding to the UDP query and answer
1414         msg = self.getFirstProtobufMessage()
1415         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
1416
1417         # then the response
1418         msg = self.getFirstProtobufMessage()
1419         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res)
1420         self.checkProtobufPolicy(msg, dnsmessage_pb2.PBDNSMessage.PolicyType.QNAME, 'zone.rpz.', '*.test.example.', 'sub.test.example', dnsmessage_pb2.PBDNSMessage.PolicyKind.NoAction)
1421         self.checkProtobufTags(msg, self._tags + self._tags_from_gettag + self._tags_from_rpz)
1422         self.assertEqual(len(msg.response.rrs), 1)
1423         rr = msg.response.rrs[0]
1424         # we have max-cache-ttl set to 15
1425         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
1426         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
1427         self.checkNoRemainingMessage()
1428
1429
1430 class ProtobufMetaFFITest(TestRecursorProtobuf):
1431     """
1432     This test makes sure that we can correctly add extra meta fields (FFI version).
1433     """
1434     _confdir = 'ProtobufMetaFFITest'
1435     _config_template = """
1436 auth-zones=example=configs/%s/example.zone""" % _confdir
1437     _lua_config_file = """
1438     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=true, logResponses=true } )
1439     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
1440     _lua_dns_script_file = """
1441     local ffi = require("ffi")
1442
1443     ffi.cdef[[
1444       typedef struct pdns_ffi_param pdns_ffi_param_t;
1445
1446       const char* pdns_ffi_param_get_qname(pdns_ffi_param_t* ref);
1447       void pdns_ffi_param_add_meta_single_string_kv(pdns_ffi_param_t *ref, const char* key, const char* val);
1448       void pdns_ffi_param_add_meta_single_int64_kv(pdns_ffi_param_t *ref, const char* key, int64_t val);
1449     ]]
1450
1451     function gettag_ffi(obj)
1452       qname = ffi.string(ffi.C.pdns_ffi_param_get_qname(obj))
1453       if qname == 'meta.example' then
1454         ffi.C.pdns_ffi_param_add_meta_single_string_kv(obj, "meta-str", "keyword")
1455         ffi.C.pdns_ffi_param_add_meta_single_int64_kv(obj, "meta-int", 42)
1456         ffi.C.pdns_ffi_param_add_meta_single_string_kv(obj, "meta-str", "content")
1457         ffi.C.pdns_ffi_param_add_meta_single_int64_kv(obj, "meta-int", 21)
1458       end
1459       return 0
1460     end
1461     """
1462     def testMeta(self):
1463         name = 'meta.example.'
1464         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.85')
1465         query = dns.message.make_query(name, 'A', want_dnssec=True)
1466         query.flags |= dns.flags.CD
1467         res = self.sendUDPQuery(query)
1468         self.assertRRsetInAnswer(res, expected)
1469
1470         # check the protobuf messages corresponding to the UDP query and answer
1471         msg = self.getFirstProtobufMessage()
1472         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
1473         self.checkProtobufMetas(msg, {'meta-str': { "stringVal" : ["content", "keyword"]}, 'meta-int': {"intVal" : [21, 42]}})
1474
1475         # then the response
1476         msg = self.getFirstProtobufMessage()
1477         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res)
1478         self.assertEqual(len(msg.response.rrs), 1)
1479         rr = msg.response.rrs[0]
1480         # we have max-cache-ttl set to 15
1481         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
1482         self.assertEqual(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.85')
1483         self.checkProtobufMetas(msg, {'meta-str': { "stringVal" : ["content", "keyword"]}, 'meta-int': {"intVal" : [21, 42]}})
1484
1485         self.checkNoRemainingMessage()