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
  83     def getFirstProtobufMessage(self, retries=1, waitTime=1):
  84         msg = None
  85
  86         print("in getFirstProtobufMessage")
  87         for param in protobufServersParameters:
  88           print(param.port)
  89           failed = 0
  90
  91           while param.queue.empty:
  92             print(failed)
  93             print(retries)
  94             if failed >= retries:
  95               break
  96
  97             failed = failed + 1
  98             print("waiting")
  99             time.sleep(waitTime)
 100
 101           self.assertFalse(param.queue.empty())
 102           data = param.queue.get(False)
 103           self.assertTrue(data)
 104           oldmsg = msg
 105           msg = dnsmessage_pb2.PBDNSMessage()
 106           msg.ParseFromString(data)
 107           if oldmsg is not None:
 108             self.assertEquals(msg, oldmsg)
 109
 110         return msg
 111
 112     def checkNoRemainingMessage(self):
 113         for param in protobufServersParameters:
 114           self.assertTrue(param.queue.empty())
 115
 116     def checkProtobufBase(self, msg, protocol, query, initiator, normalQueryResponse=True, expectedECS=None, receivedSize=None):
 117         self.assertTrue(msg)
 118         self.assertTrue(msg.HasField('timeSec'))
 119         self.assertTrue(msg.HasField('socketFamily'))
 120         self.assertEquals(msg.socketFamily, dnsmessage_pb2.PBDNSMessage.INET)
 121         self.assertTrue(msg.HasField('from'))
 122         fromvalue = getattr(msg, 'from')
 123         self.assertEquals(socket.inet_ntop(socket.AF_INET, fromvalue), initiator)
 124         self.assertTrue(msg.HasField('socketProtocol'))
 125         self.assertEquals(msg.socketProtocol, protocol)
 126         self.assertTrue(msg.HasField('messageId'))
 127         self.assertTrue(msg.HasField('serverIdentity'))
 128         self.assertTrue(msg.HasField('id'))
 129         self.assertEquals(msg.id, query.id)
 130         self.assertTrue(msg.HasField('inBytes'))
 131         if normalQueryResponse:
 132             # compare inBytes with length of query/response
 133             # Note that for responses, the size we received might differ
 134             # because dnspython might compress labels differently from
 135             # the recursor
 136             if receivedSize:
 137                 self.assertEquals(msg.inBytes, receivedSize)
 138             else:
 139                 self.assertEquals(msg.inBytes, len(query.to_wire()))
 140         if expectedECS is not None:
 141             self.assertTrue(msg.HasField('originalRequestorSubnet'))
 142             # v4 only for now
 143             self.assertEquals(len(msg.originalRequestorSubnet), 4)
 144             self.assertEquals(socket.inet_ntop(socket.AF_INET, msg.originalRequestorSubnet), '127.0.0.1')
 145
 146     def checkOutgoingProtobufBase(self, msg, protocol, query, initiator, length=None):
 147         self.assertTrue(msg)
 148         self.assertTrue(msg.HasField('timeSec'))
 149         self.assertTrue(msg.HasField('socketFamily'))
 150         self.assertEquals(msg.socketFamily, dnsmessage_pb2.PBDNSMessage.INET)
 151         self.assertTrue(msg.HasField('socketProtocol'))
 152         self.assertEquals(msg.socketProtocol, protocol)
 153         self.assertTrue(msg.HasField('messageId'))
 154         self.assertTrue(msg.HasField('serverIdentity'))
 155         self.assertTrue(msg.HasField('id'))
 156         self.assertNotEquals(msg.id, query.id)
 157         self.assertTrue(msg.HasField('inBytes'))
 158         if length is not None:
 159           self.assertEquals(msg.inBytes, length)
 160         else:
 161           # compare inBytes with length of query/response
 162           self.assertEquals(msg.inBytes, len(query.to_wire()))
 163
 164     def checkProtobufQuery(self, msg, protocol, query, qclass, qtype, qname, initiator='127.0.0.1'):
 165         self.assertEquals(msg.type, dnsmessage_pb2.PBDNSMessage.DNSQueryType)
 166         self.checkProtobufBase(msg, protocol, query, initiator)
 167         # dnsdist doesn't fill the responder field for responses
 168         # because it doesn't keep the information around.
 169         self.assertTrue(msg.HasField('to'))
 170         self.assertEquals(socket.inet_ntop(socket.AF_INET, msg.to), '127.0.0.1')
 171         self.assertTrue(msg.HasField('question'))
 172         self.assertTrue(msg.question.HasField('qClass'))
 173         self.assertEquals(msg.question.qClass, qclass)
 174         self.assertTrue(msg.question.HasField('qType'))
 175         self.assertEquals(msg.question.qClass, qtype)
 176         self.assertTrue(msg.question.HasField('qName'))
 177         self.assertEquals(msg.question.qName, qname)
 178
 179     def checkProtobufResponse(self, msg, protocol, response, initiator='127.0.0.1', receivedSize=None):
 180         self.assertEquals(msg.type, dnsmessage_pb2.PBDNSMessage.DNSResponseType)
 181         self.checkProtobufBase(msg, protocol, response, initiator, receivedSize=receivedSize)
 182         self.assertTrue(msg.HasField('response'))
 183         self.assertTrue(msg.response.HasField('queryTimeSec'))
 184
 185     def checkProtobufResponseRecord(self, record, rclass, rtype, rname, rttl, checkTTL=True):
 186         self.assertTrue(record.HasField('class'))
 187         self.assertEquals(getattr(record, 'class'), rclass)
 188         self.assertTrue(record.HasField('type'))
 189         self.assertEquals(record.type, rtype)
 190         self.assertTrue(record.HasField('name'))
 191         self.assertEquals(record.name, rname)
 192         self.assertTrue(record.HasField('ttl'))
 193         if checkTTL:
 194             self.assertEquals(record.ttl, rttl)
 195         self.assertTrue(record.HasField('rdata'))
 196
 197     def checkProtobufPolicy(self, msg, policyType, reason):
 198         self.assertEquals(msg.type, dnsmessage_pb2.PBDNSMessage.DNSResponseType)
 199         self.assertTrue(msg.response.HasField('appliedPolicyType'))
 200         self.assertTrue(msg.response.HasField('appliedPolicy'))
 201         self.assertEquals(msg.response.appliedPolicy, reason)
 202         self.assertEquals(msg.response.appliedPolicyType, policyType)
 203
 204     def checkProtobufTags(self, msg, tags):
 205         self.assertEquals(len(msg.response.tags), len(tags))
 206         for tag in msg.response.tags:
 207             self.assertTrue(tag in tags)
 208
 209     def checkProtobufOutgoingQuery(self, msg, protocol, query, qclass, qtype, qname, initiator='127.0.0.1', length=None):
 210         self.assertEquals(msg.type, dnsmessage_pb2.PBDNSMessage.DNSOutgoingQueryType)
 211         self.checkOutgoingProtobufBase(msg, protocol, query, initiator, length=length)
 212         self.assertTrue(msg.HasField('to'))
 213         self.assertTrue(msg.HasField('question'))
 214         self.assertTrue(msg.question.HasField('qClass'))
 215         self.assertEquals(msg.question.qClass, qclass)
 216         self.assertTrue(msg.question.HasField('qType'))
 217         self.assertEquals(msg.question.qClass, qtype)
 218         self.assertTrue(msg.question.HasField('qName'))
 219         self.assertEquals(msg.question.qName, qname)
 220
 221     def checkProtobufIncomingResponse(self, msg, protocol, response, initiator='127.0.0.1', length=None):
 222         self.assertEquals(msg.type, dnsmessage_pb2.PBDNSMessage.DNSIncomingResponseType)
 223         self.checkOutgoingProtobufBase(msg, protocol, response, initiator, length=length)
 224         self.assertTrue(msg.HasField('response'))
 225         self.assertTrue(msg.response.HasField('rcode'))
 226         self.assertTrue(msg.response.HasField('queryTimeSec'))
 227
 228     def checkProtobufIncomingNetworkErrorResponse(self, msg, protocol, response, initiator='127.0.0.1'):
 229         self.checkProtobufIncomingResponse(msg, protocol, response, initiator, length=0)
 230         self.assertEquals(msg.response.rcode, 65536)
 231
 232     def checkProtobufIdentity(self, msg, requestorId, deviceId, deviceName):
 233         self.assertTrue(msg.HasField('requestorId'))
 234         self.assertTrue(msg.HasField('deviceId'))
 235         self.assertTrue(msg.HasField('deviceName'))
 236         self.assertEquals(msg.requestorId, requestorId)
 237         self.assertEquals(msg.deviceId, deviceId)
 238         self.assertEquals(msg.deviceName, deviceName)
 239
 240     @classmethod
 241     def setUpClass(cls):
 242
 243         cls.setUpSockets()
 244
 245         cls.startResponders()
 246
 247         confdir = os.path.join('configs', cls._confdir)
 248         cls.createConfigDir(confdir)
 249
 250         cls.generateRecursorConfig(confdir)
 251         cls.startRecursor(confdir, cls._recursorPort)
 252
 253     def setUp(self):
 254       # Make sure the queue is empty, in case
 255       # a previous test failed
 256       for param in protobufServersParameters:
 257         while not param.queue.empty():
 258           param.queue.get(False)
 259
 260     @classmethod
 261     def generateRecursorConfig(cls, confdir):
 262         authzonepath = os.path.join(confdir, 'example.zone')
 263         with open(authzonepath, 'w') as authzone:
 264             authzone.write("""$ORIGIN example.
 265 @ 3600 IN SOA {soa}
 266 a 3600 IN A 192.0.2.42
 267 tagged 3600 IN A 192.0.2.84
 268 query-selected 3600 IN A 192.0.2.84
 269 answer-selected 3600 IN A 192.0.2.84
 270 types 3600 IN A 192.0.2.84
 271 types 3600 IN AAAA 2001:DB8::1
 272 types 3600 IN TXT "Lorem ipsum dolor sit amet"
 273 types 3600 IN MX 10 a.example.
 274 types 3600 IN SPF "v=spf1 -all"
 275 types 3600 IN SRV 10 20 443 a.example.
 276 cname 3600 IN CNAME a.example.
 277
 278 """.format(soa=cls._SOA))
 279         super(TestRecursorProtobuf, cls).generateRecursorConfig(confdir)
 280
 281     @classmethod
 282     def tearDownClass(cls):
 283         cls.tearDownRecursor()
 284
 285 class ProtobufDefaultTest(TestRecursorProtobuf):
 286     """
 287     This test makes sure that we correctly export queries and response over protobuf.
 288     """
 289
 290     _confdir = 'ProtobufDefault'
 291     _config_template = """
 292 auth-zones=example=configs/%s/example.zone""" % _confdir
 293
 294     def testA(self):
 295         name = 'a.example.'
 296         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 297         query = dns.message.make_query(name, 'A', want_dnssec=True)
 298         query.flags |= dns.flags.CD
 299         res = self.sendUDPQuery(query)
 300
 301         self.assertRRsetInAnswer(res, expected)
 302
 303         # check the protobuf messages corresponding to the UDP query and answer
 304         msg = self.getFirstProtobufMessage()
 305         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
 306         # then the response
 307         msg = self.getFirstProtobufMessage()
 308         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res, '127.0.0.1')
 309         self.assertEquals(len(msg.response.rrs), 1)
 310         rr = msg.response.rrs[0]
 311         # we have max-cache-ttl set to 15
 312         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 313         self.assertEquals(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
 314         self.checkNoRemainingMessage()
 315
 316     def testCNAME(self):
 317         name = 'cname.example.'
 318         expectedCNAME = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'CNAME', 'a.example.')
 319         expectedA = dns.rrset.from_text('a.example.', 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 320         query = dns.message.make_query(name, 'A', want_dnssec=True)
 321         query.flags |= dns.flags.CD
 322         raw = self.sendUDPQuery(query, decode=False)
 323         res = dns.message.from_wire(raw)
 324         self.assertRRsetInAnswer(res, expectedCNAME)
 325         self.assertRRsetInAnswer(res, expectedA)
 326
 327         # check the protobuf messages corresponding to the UDP query and answer
 328         # but first let the protobuf messages the time to get there
 329         msg = self.getFirstProtobufMessage()
 330         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
 331         # then the response
 332         msg = self.getFirstProtobufMessage()
 333         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res, '127.0.0.1', receivedSize=len(raw))
 334         self.assertEquals(len(msg.response.rrs), 2)
 335         rr = msg.response.rrs[0]
 336         # we don't want to check the TTL for the A record, it has been cached by the previous test
 337         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.CNAME, name, 15)
 338         self.assertEquals(rr.rdata, 'a.example.')
 339         rr = msg.response.rrs[1]
 340         # we have max-cache-ttl set to 15
 341         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, 'a.example.', 15, checkTTL=False)
 342         self.assertEquals(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
 343         self.checkNoRemainingMessage()
 344
 345 class OutgoingProtobufDefaultTest(TestRecursorProtobuf):
 346     """
 347     This test makes sure that we correctly export outgoing queries over protobuf.
 348     It must be improved and setup env so we can check for incoming responses, but makes sure for now
 349     that the recursor at least connects to the protobuf server.
 350     """
 351
 352     _confdir = 'OutgoingProtobufDefault'
 353     _config_template = """
 354     # Switch off QName Minimization, it generates much more protobuf messages
 355     # (or make the test much more smart!)
 356     qname-minimization=no
 357     auth-zones=example=configs/%s/example.zone""" % _confdir
 358     _lua_config_file = """
 359     outgoingProtobufServer({"127.0.0.1:%d", "127.0.0.1:%d"})
 360     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 361
 362     def testA(self):
 363         name = 'www.example.org.'
 364         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 365         query = dns.message.make_query(name, 'A', want_dnssec=True)
 366         query.flags |= dns.flags.RD
 367         res = self.sendUDPQuery(query)
 368
 369         # check the protobuf messages corresponding to the UDP query and answer
 370         msg = self.getFirstProtobufMessage()
 371         self.checkProtobufOutgoingQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
 372         # then the response
 373         msg = self.getFirstProtobufMessage()
 374         self.checkProtobufIncomingNetworkErrorResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res)
 375         self.checkNoRemainingMessage()
 376
 377 class OutgoingProtobufNoQueriesTest(TestRecursorProtobuf):
 378     """
 379     This test makes sure that we correctly export incoming responses but not outgoing queries over protobuf.
 380     It must be improved and setup env so we can check for incoming responses, but makes sure for now
 381     that the recursor at least connects to the protobuf server.
 382     """
 383
 384     _confdir = 'OutgoingProtobufNoQueries'
 385     _config_template = """
 386     # Switch off QName Minimization, it generates much more protobuf messages
 387     # (or make the test much more smart!)
 388     qname-minimization=no
 389     auth-zones=example=configs/%s/example.zone""" % _confdir
 390     _lua_config_file = """
 391     outgoingProtobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=false, logResponses=true })
 392     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 393
 394     def testA(self):
 395         name = 'www.example.org.'
 396         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 397         query = dns.message.make_query(name, 'A', want_dnssec=True)
 398         query.flags |= dns.flags.RD
 399         res = self.sendUDPQuery(query)
 400
 401         # check the response
 402         msg = self.getFirstProtobufMessage()
 403         self.checkProtobufIncomingNetworkErrorResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res)
 404         self.checkNoRemainingMessage()
 405
 406 class ProtobufMasksTest(TestRecursorProtobuf):
 407     """
 408     This test makes sure that we correctly export queries and response over protobuf, respecting the configured initiator masking.
 409     """
 410
 411     _confdir = 'ProtobufMasks'
 412     _config_template = """
 413 auth-zones=example=configs/%s/example.zone""" % _confdir
 414     _protobufMaskV4 = 4
 415     _protobufMaskV6 = 128
 416     _lua_config_file = """
 417     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"})
 418     setProtobufMasks(%d, %d)
 419     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port, _protobufMaskV4, _protobufMaskV6)
 420
 421     def testA(self):
 422         name = 'a.example.'
 423         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 424         query = dns.message.make_query(name, 'A', want_dnssec=True)
 425         query.flags |= dns.flags.CD
 426         res = self.sendUDPQuery(query)
 427         self.assertRRsetInAnswer(res, expected)
 428
 429         # check the protobuf messages corresponding to the UDP query and answer
 430         # but first let the protobuf messages the time to get there
 431         msg = self.getFirstProtobufMessage()
 432         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name, '112.0.0.0')
 433         # then the response
 434         msg = self.getFirstProtobufMessage()
 435         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res, '112.0.0.0')
 436         self.assertEquals(len(msg.response.rrs), 1)
 437         rr = msg.response.rrs[0]
 438         # we have max-cache-ttl set to 15
 439         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 440         self.assertEquals(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
 441         self.checkNoRemainingMessage()
 442
 443 class ProtobufQueriesOnlyTest(TestRecursorProtobuf):
 444     """
 445     This test makes sure that we correctly export queries but not responses over protobuf.
 446     """
 447
 448     _confdir = 'ProtobufQueriesOnly'
 449     _config_template = """
 450 auth-zones=example=configs/%s/example.zone""" % _confdir
 451     _lua_config_file = """
 452     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=true, logResponses=false } )
 453     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 454
 455     def testA(self):
 456         name = 'a.example.'
 457         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 458         query = dns.message.make_query(name, 'A', want_dnssec=True)
 459         query.flags |= dns.flags.CD
 460         res = self.sendUDPQuery(query)
 461         self.assertRRsetInAnswer(res, expected)
 462
 463         # check the protobuf message corresponding to the UDP query
 464         msg = self.getFirstProtobufMessage()
 465         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
 466         # no response
 467         self.checkNoRemainingMessage()
 468
 469 class ProtobufResponsesOnlyTest(TestRecursorProtobuf):
 470     """
 471     This test makes sure that we correctly export responses but not queries over protobuf.
 472     """
 473
 474     _confdir = 'ProtobufResponsesOnly'
 475     _config_template = """
 476 auth-zones=example=configs/%s/example.zone""" % _confdir
 477     _lua_config_file = """
 478     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=false, logResponses=true } )
 479     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 480
 481     def testA(self):
 482         name = 'a.example.'
 483         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 484         query = dns.message.make_query(name, 'A', want_dnssec=True)
 485         query.flags |= dns.flags.CD
 486         res = self.sendUDPQuery(query)
 487         self.assertRRsetInAnswer(res, expected)
 488
 489         # check the protobuf message corresponding to the UDP response
 490         msg = self.getFirstProtobufMessage()
 491         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res)
 492         self.assertEquals(len(msg.response.rrs), 1)
 493         rr = msg.response.rrs[0]
 494         # we have max-cache-ttl set to 15
 495         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 496         self.assertEquals(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
 497         # nothing else in the queue
 498         self.checkNoRemainingMessage()
 499
 500 class ProtobufTaggedOnlyTest(TestRecursorProtobuf):
 501     """
 502     This test makes sure that we correctly export queries and responses but only if they have been tagged.
 503     """
 504
 505     _confdir = 'ProtobufTaggedOnly'
 506     _config_template = """
 507 auth-zones=example=configs/%s/example.zone""" % _confdir
 508     _lua_config_file = """
 509     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=true, logResponses=true, taggedOnly=true } )
 510     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 511     _tags = ['tag1', 'tag2']
 512     _tag_from_gettag = 'tag-from-gettag'
 513     _lua_dns_script_file = """
 514     function gettag(remote, ednssubnet, localip, qname, qtype, ednsoptions, tcp)
 515       if qname:equal('tagged.example.') then
 516         return 0, { '%s' }
 517       end
 518       return 0
 519     end
 520     function preresolve(dq)
 521       if dq.qname:equal('tagged.example.') then
 522         dq:addPolicyTag('%s')
 523         dq:addPolicyTag('%s')
 524       end
 525       return false
 526     end
 527     """ % (_tag_from_gettag, _tags[0], _tags[1])
 528
 529     def testA(self):
 530         name = 'a.example.'
 531         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 532         query = dns.message.make_query(name, 'A', want_dnssec=True)
 533         query.flags |= dns.flags.CD
 534         res = self.sendUDPQuery(query)
 535         self.assertRRsetInAnswer(res, expected)
 536
 537         # check the protobuf message corresponding to the UDP response
 538         # the first query and answer are not tagged, so there is nothing in the queue
 539         time.sleep(1)
 540         self.checkNoRemainingMessage()
 541
 542     def testTagged(self):
 543         name = 'tagged.example.'
 544         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.84')
 545         query = dns.message.make_query(name, 'A', want_dnssec=True)
 546         query.flags |= dns.flags.CD
 547         res = self.sendUDPQuery(query)
 548         self.assertRRsetInAnswer(res, expected)
 549
 550         # check the protobuf messages corresponding to the UDP query and answer
 551         msg = self.getFirstProtobufMessage()
 552         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
 553         self.checkProtobufTags(msg, [self._tag_from_gettag])
 554         # then the response
 555         msg = self.getFirstProtobufMessage()
 556         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res)
 557         self.assertEquals(len(msg.response.rrs), 1)
 558         rr = msg.response.rrs[0]
 559         # we have max-cache-ttl set to 15
 560         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 561         self.assertEquals(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.84')
 562         tags = [self._tag_from_gettag] + self._tags
 563         self.checkProtobufTags(msg, tags)
 564         self.checkNoRemainingMessage()
 565
 566 class ProtobufSelectedFromLuaTest(TestRecursorProtobuf):
 567     """
 568     This test makes sure that we correctly export queries and responses but only if they have been selected from Lua.
 569     """
 570
 571     _confdir = 'ProtobufSelectedFromLua'
 572     _config_template = """
 573 auth-zones=example=configs/%s/example.zone""" % _confdir
 574     _lua_config_file = """
 575     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=false, logResponses=false } )
 576     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 577     _lua_dns_script_file = """
 578     local ffi = require("ffi")
 579
 580     ffi.cdef[[
 581       typedef struct pdns_ffi_param pdns_ffi_param_t;
 582
 583       const char* pdns_ffi_param_get_qname(pdns_ffi_param_t* ref);
 584       void pdns_ffi_param_set_log_query(pdns_ffi_param_t* ref, bool logQuery);
 585     ]]
 586
 587     function gettag_ffi(obj)
 588       qname = ffi.string(ffi.C.pdns_ffi_param_get_qname(obj))
 589       if qname == 'query-selected.example' then
 590         ffi.C.pdns_ffi_param_set_log_query(obj, true)
 591       end
 592       return 0
 593     end
 594
 595     function preresolve(dq)
 596       if dq.qname:equal('answer-selected.example.') then
 597         dq.logResponse = true
 598       end
 599       return false
 600     end
 601     """
 602
 603     def testA(self):
 604         name = 'a.example.'
 605         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 606         query = dns.message.make_query(name, 'A', want_dnssec=True)
 607         query.flags |= dns.flags.CD
 608         res = self.sendUDPQuery(query)
 609         self.assertRRsetInAnswer(res, expected)
 610
 611         # check the protobuf message corresponding to the UDP response
 612         # the first query and answer are not selected, so there is nothing in the queue
 613         self.checkNoRemainingMessage()
 614
 615     def testQuerySelected(self):
 616         name = 'query-selected.example.'
 617         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.84')
 618         query = dns.message.make_query(name, 'A', want_dnssec=True)
 619         query.flags |= dns.flags.CD
 620         res = self.sendUDPQuery(query)
 621         self.assertRRsetInAnswer(res, expected)
 622
 623         # check the protobuf messages corresponding to the UDP query
 624         msg = self.getFirstProtobufMessage()
 625         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
 626         # there should be no response
 627         self.checkNoRemainingMessage()
 628
 629     def testResponseSelected(self):
 630         name = 'answer-selected.example.'
 631         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.84')
 632         query = dns.message.make_query(name, 'A', want_dnssec=True)
 633         query.flags |= dns.flags.CD
 634         res = self.sendUDPQuery(query)
 635         self.assertRRsetInAnswer(res, expected)
 636
 637         # check the protobuf messages corresponding to the UDP response
 638         msg = self.getFirstProtobufMessage()
 639         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res)
 640         self.assertEquals(len(msg.response.rrs), 1)
 641         rr = msg.response.rrs[0]
 642         # we have max-cache-ttl set to 15
 643         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 644         self.assertEquals(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.84')
 645         self.checkNoRemainingMessage()
 646
 647 class ProtobufExportTypesTest(TestRecursorProtobuf):
 648     """
 649     This test makes sure that we correctly export other types than A, AAAA and CNAME over protobuf.
 650     """
 651
 652     _confdir = 'ProtobufExportTypes'
 653     _config_template = """
 654 auth-zones=example=configs/%s/example.zone""" % _confdir
 655     _lua_config_file = """
 656     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { exportTypes={"AAAA", "MX", "SPF", "SRV", "TXT"} } )
 657     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 658
 659     def testA(self):
 660         name = 'types.example.'
 661         expected = [dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.84'),
 662                     dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'AAAA', '2001:DB8::1'),
 663                     dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'MX', '10 a.example.'),
 664                     dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'SPF', '"v=spf1 -all"'),
 665                     dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'SRV', '10 20 443 a.example.'),
 666                     dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'TXT', '"Lorem ipsum dolor sit amet"'),
 667                    ]
 668         query = dns.message.make_query(name, 'ANY', want_dnssec=True)
 669         query.flags |= dns.flags.CD
 670         raw = self.sendUDPQuery(query, decode=False)
 671         res = dns.message.from_wire(raw)
 672
 673         for rrset in expected:
 674             self.assertRRsetInAnswer(res, rrset)
 675
 676         # check the protobuf messages corresponding to the UDP query and answer
 677         msg = self.getFirstProtobufMessage()
 678         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
 679         # then the response
 680         msg = self.getFirstProtobufMessage()
 681         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res, '127.0.0.1', receivedSize=len(raw))
 682         self.assertEquals(len(msg.response.rrs), 5)
 683         for rr in msg.response.rrs:
 684             self.assertTrue(rr.type in [dns.rdatatype.AAAA, dns.rdatatype.TXT, dns.rdatatype.MX, dns.rdatatype.SPF, dns.rdatatype.SRV])
 685
 686             if rr.type == dns.rdatatype.AAAA:
 687                 self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.AAAA, name, 15)
 688                 self.assertEquals(socket.inet_ntop(socket.AF_INET6, rr.rdata), '2001:db8::1')
 689             elif rr.type == dns.rdatatype.TXT:
 690                 self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.TXT, name, 15)
 691                 self.assertEquals(rr.rdata, '"Lorem ipsum dolor sit amet"')
 692             elif rr.type == dns.rdatatype.MX:
 693                 self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.MX, name, 15)
 694                 self.assertEquals(rr.rdata, 'a.example.')
 695             elif rr.type == dns.rdatatype.SPF:
 696                 self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.SPF, name, 15)
 697                 self.assertEquals(rr.rdata, '"v=spf1 -all"')
 698             elif rr.type == dns.rdatatype.SRV:
 699                 self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.SRV, name, 15)
 700                 self.assertEquals(rr.rdata, 'a.example.')
 701
 702         self.checkNoRemainingMessage()
 703
 704 class ProtobufTaggedExtraFieldsTest(TestRecursorProtobuf):
 705     """
 706     This test makes sure that we correctly export extra fields that may have been set while being tagged.
 707     """
 708
 709     _confdir = 'ProtobufTaggedExtraFields'
 710     _config_template = """
 711 auth-zones=example=configs/%s/example.zone""" % _confdir
 712     _lua_config_file = """
 713     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=true, logResponses=true } )
 714     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 715     _requestorId = 'S-000001727'
 716     _deviceId = 'd1:0a:91:dc:cc:82'
 717     _deviceName = 'Joe'
 718     _lua_dns_script_file = """
 719     function gettag(remote, ednssubnet, localip, qname, qtype, ednsoptions, tcp)
 720       if qname:equal('tagged.example.') then
 721         -- tag number, policy tags, data, requestorId, deviceId, deviceName
 722         return 0, {}, {}, '%s', '%s', '%s'
 723       end
 724       return 0
 725     end
 726     """ % (_requestorId, _deviceId, _deviceName)
 727
 728     def testA(self):
 729         name = 'a.example.'
 730         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 731         query = dns.message.make_query(name, 'A', want_dnssec=True)
 732         query.flags |= dns.flags.CD
 733         res = self.sendUDPQuery(query)
 734         self.assertRRsetInAnswer(res, expected)
 735
 736         # check the protobuf message corresponding to the UDP response
 737         # the first query and answer are not tagged, so there is nothing in the queue
 738         # check the protobuf messages corresponding to the UDP query and answer
 739         msg = self.getFirstProtobufMessage()
 740         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
 741         self.checkProtobufIdentity(msg, '', '', '')
 742
 743         # then the response
 744         msg = self.getFirstProtobufMessage()
 745         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res, '127.0.0.1')
 746         self.assertEquals(len(msg.response.rrs), 1)
 747         rr = msg.response.rrs[0]
 748         # we have max-cache-ttl set to 15
 749         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 750         self.assertEquals(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
 751         self.checkProtobufIdentity(msg, '', '', '')
 752         self.checkNoRemainingMessage()
 753
 754     def testTagged(self):
 755         name = 'tagged.example.'
 756         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.84')
 757         query = dns.message.make_query(name, 'A', want_dnssec=True)
 758         query.flags |= dns.flags.CD
 759         res = self.sendUDPQuery(query)
 760         self.assertRRsetInAnswer(res, expected)
 761
 762         # check the protobuf messages corresponding to the UDP query and answer
 763         msg = self.getFirstProtobufMessage()
 764         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
 765         self.checkProtobufIdentity(msg, self._requestorId, self._deviceId, self._deviceName)
 766
 767         # then the response
 768         msg = self.getFirstProtobufMessage()
 769         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res)
 770         self.assertEquals(len(msg.response.rrs), 1)
 771         rr = msg.response.rrs[0]
 772         # we have max-cache-ttl set to 15
 773         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 774         self.assertEquals(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.84')
 775         self.checkProtobufIdentity(msg, self._requestorId, self._deviceId, self._deviceName)
 776         self.checkNoRemainingMessage()
 777
 778 class ProtobufTaggedExtraFieldsFFITest(ProtobufTaggedExtraFieldsTest):
 779     """
 780     This test makes sure that we correctly export extra fields that may have been set while being tagged (FFI version).
 781     """
 782     _confdir = 'ProtobufTaggedExtraFieldsFFI'
 783     _config_template = """
 784 auth-zones=example=configs/%s/example.zone""" % _confdir
 785     _lua_config_file = """
 786     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=true, logResponses=true } )
 787     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 788     _lua_dns_script_file = """
 789     local ffi = require("ffi")
 790
 791     ffi.cdef[[
 792       typedef struct pdns_ffi_param pdns_ffi_param_t;
 793
 794       const char* pdns_ffi_param_get_qname(pdns_ffi_param_t* ref);
 795       void pdns_ffi_param_set_tag(pdns_ffi_param_t* ref, unsigned int tag);
 796       void pdns_ffi_param_set_requestorid(pdns_ffi_param_t* ref, const char* name);
 797       void pdns_ffi_param_set_devicename(pdns_ffi_param_t* ref, const char* name);
 798       void pdns_ffi_param_set_deviceid(pdns_ffi_param_t* ref, size_t len, const void* name);
 799     ]]
 800
 801     function gettag_ffi(obj)
 802       qname = ffi.string(ffi.C.pdns_ffi_param_get_qname(obj))
 803       if qname == 'tagged.example' then
 804         ffi.C.pdns_ffi_param_set_requestorid(obj, "%s")
 805         deviceid = "%s"
 806         ffi.C.pdns_ffi_param_set_deviceid(obj, string.len(deviceid), deviceid)
 807         ffi.C.pdns_ffi_param_set_devicename(obj, "%s")
 808       end
 809       return 0
 810     end
 811     """ % (ProtobufTaggedExtraFieldsTest._requestorId, ProtobufTaggedExtraFieldsTest._deviceId, ProtobufTaggedExtraFieldsTest._deviceName)