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):
 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         # compare inBytes with length of query/response
 159         self.assertEquals(msg.inBytes, len(query.to_wire()))
 160
 161     def checkProtobufQuery(self, msg, protocol, query, qclass, qtype, qname, initiator='127.0.0.1'):
 162         self.assertEquals(msg.type, dnsmessage_pb2.PBDNSMessage.DNSQueryType)
 163         self.checkProtobufBase(msg, protocol, query, initiator)
 164         # dnsdist doesn't fill the responder field for responses
 165         # because it doesn't keep the information around.
 166         self.assertTrue(msg.HasField('to'))
 167         self.assertEquals(socket.inet_ntop(socket.AF_INET, msg.to), '127.0.0.1')
 168         self.assertTrue(msg.HasField('question'))
 169         self.assertTrue(msg.question.HasField('qClass'))
 170         self.assertEquals(msg.question.qClass, qclass)
 171         self.assertTrue(msg.question.HasField('qType'))
 172         self.assertEquals(msg.question.qClass, qtype)
 173         self.assertTrue(msg.question.HasField('qName'))
 174         self.assertEquals(msg.question.qName, qname)
 175
 176     def checkProtobufResponse(self, msg, protocol, response, initiator='127.0.0.1', receivedSize=None):
 177         self.assertEquals(msg.type, dnsmessage_pb2.PBDNSMessage.DNSResponseType)
 178         self.checkProtobufBase(msg, protocol, response, initiator, receivedSize=receivedSize)
 179         self.assertTrue(msg.HasField('response'))
 180         self.assertTrue(msg.response.HasField('queryTimeSec'))
 181
 182     def checkProtobufResponseRecord(self, record, rclass, rtype, rname, rttl, checkTTL=True):
 183         self.assertTrue(record.HasField('class'))
 184         self.assertEquals(getattr(record, 'class'), rclass)
 185         self.assertTrue(record.HasField('type'))
 186         self.assertEquals(record.type, rtype)
 187         self.assertTrue(record.HasField('name'))
 188         self.assertEquals(record.name, rname)
 189         self.assertTrue(record.HasField('ttl'))
 190         if checkTTL:
 191             self.assertEquals(record.ttl, rttl)
 192         self.assertTrue(record.HasField('rdata'))
 193
 194     def checkProtobufPolicy(self, msg, policyType, reason):
 195         self.assertEquals(msg.type, dnsmessage_pb2.PBDNSMessage.DNSResponseType)
 196         self.assertTrue(msg.response.HasField('appliedPolicyType'))
 197         self.assertTrue(msg.response.HasField('appliedPolicy'))
 198         self.assertEquals(msg.response.appliedPolicy, reason)
 199         self.assertEquals(msg.response.appliedPolicyType, policyType)
 200
 201     def checkProtobufTags(self, msg, tags):
 202         self.assertEquals(len(msg.response.tags), len(tags))
 203         for tag in msg.response.tags:
 204             self.assertTrue(tag in tags)
 205
 206     def checkProtobufOutgoingQuery(self, msg, protocol, query, qclass, qtype, qname, initiator='127.0.0.1'):
 207         self.assertEquals(msg.type, dnsmessage_pb2.PBDNSMessage.DNSOutgoingQueryType)
 208         self.checkOutgoingProtobufBase(msg, protocol, query, initiator)
 209         self.assertTrue(msg.HasField('to'))
 210         self.assertTrue(msg.HasField('question'))
 211         self.assertTrue(msg.question.HasField('qClass'))
 212         self.assertEquals(msg.question.qClass, qclass)
 213         self.assertTrue(msg.question.HasField('qType'))
 214         self.assertEquals(msg.question.qClass, qtype)
 215         self.assertTrue(msg.question.HasField('qName'))
 216         self.assertEquals(msg.question.qName, qname)
 217
 218     def checkProtobufIncomingResponse(self, msg, protocol, response, initiator='127.0.0.1'):
 219         self.assertEquals(msg.type, dnsmessage_pb2.PBDNSMessage.DNSIncomingResponseType)
 220         self.checkOutgoingProtobufBase(msg, protocol, response, initiator)
 221         self.assertTrue(msg.HasField('response'))
 222         self.assertTrue(msg.response.HasField('queryTimeSec'))
 223
 224     @classmethod
 225     def setUpClass(cls):
 226
 227         cls.setUpSockets()
 228
 229         cls.startResponders()
 230
 231         confdir = os.path.join('configs', cls._confdir)
 232         cls.createConfigDir(confdir)
 233
 234         cls.generateRecursorConfig(confdir)
 235         cls.startRecursor(confdir, cls._recursorPort)
 236
 237     def setUp(self):
 238       # Make sure the queue is empty, in case
 239       # a previous test failed
 240       for param in protobufServersParameters:
 241         while not param.queue.empty():
 242           param.queue.get(False)
 243
 244     @classmethod
 245     def generateRecursorConfig(cls, confdir):
 246         authzonepath = os.path.join(confdir, 'example.zone')
 247         with open(authzonepath, 'w') as authzone:
 248             authzone.write("""$ORIGIN example.
 249 @ 3600 IN SOA {soa}
 250 a 3600 IN A 192.0.2.42
 251 tagged 3600 IN A 192.0.2.84
 252 query-selected 3600 IN A 192.0.2.84
 253 answer-selected 3600 IN A 192.0.2.84
 254 types 3600 IN A 192.0.2.84
 255 types 3600 IN AAAA 2001:DB8::1
 256 types 3600 IN TXT "Lorem ipsum dolor sit amet"
 257 types 3600 IN MX 10 a.example.
 258 types 3600 IN SPF "v=spf1 -all"
 259 types 3600 IN SRV 10 20 443 a.example.
 260 cname 3600 IN CNAME a.example.
 261
 262 """.format(soa=cls._SOA))
 263         super(TestRecursorProtobuf, cls).generateRecursorConfig(confdir)
 264
 265     @classmethod
 266     def tearDownClass(cls):
 267         cls.tearDownRecursor()
 268
 269 class ProtobufDefaultTest(TestRecursorProtobuf):
 270     """
 271     This test makes sure that we correctly export queries and response over protobuf.
 272     """
 273
 274     _confdir = 'ProtobufDefault'
 275     _config_template = """
 276 auth-zones=example=configs/%s/example.zone""" % _confdir
 277
 278     def testA(self):
 279         name = 'a.example.'
 280         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 281         query = dns.message.make_query(name, 'A', want_dnssec=True)
 282         query.flags |= dns.flags.CD
 283         res = self.sendUDPQuery(query)
 284
 285         self.assertRRsetInAnswer(res, expected)
 286
 287         # check the protobuf messages corresponding to the UDP query and answer
 288         msg = self.getFirstProtobufMessage()
 289         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
 290         # then the response
 291         msg = self.getFirstProtobufMessage()
 292         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res, '127.0.0.1')
 293         self.assertEquals(len(msg.response.rrs), 1)
 294         rr = msg.response.rrs[0]
 295         # we have max-cache-ttl set to 15
 296         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 297         self.assertEquals(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
 298         self.checkNoRemainingMessage()
 299
 300     def testCNAME(self):
 301         name = 'cname.example.'
 302         expectedCNAME = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'CNAME', 'a.example.')
 303         expectedA = dns.rrset.from_text('a.example.', 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 304         query = dns.message.make_query(name, 'A', want_dnssec=True)
 305         query.flags |= dns.flags.CD
 306         raw = self.sendUDPQuery(query, decode=False)
 307         res = dns.message.from_wire(raw)
 308         self.assertRRsetInAnswer(res, expectedCNAME)
 309         self.assertRRsetInAnswer(res, expectedA)
 310
 311         # check the protobuf messages corresponding to the UDP query and answer
 312         # but first let the protobuf messages the time to get there
 313         msg = self.getFirstProtobufMessage()
 314         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
 315         # then the response
 316         msg = self.getFirstProtobufMessage()
 317         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res, '127.0.0.1', receivedSize=len(raw))
 318         self.assertEquals(len(msg.response.rrs), 2)
 319         rr = msg.response.rrs[0]
 320         # we don't want to check the TTL for the A record, it has been cached by the previous test
 321         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.CNAME, name, 15)
 322         self.assertEquals(rr.rdata, 'a.example.')
 323         rr = msg.response.rrs[1]
 324         # we have max-cache-ttl set to 15
 325         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, 'a.example.', 15, checkTTL=False)
 326         self.assertEquals(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
 327         self.checkNoRemainingMessage()
 328
 329 class OutgoingProtobufDefaultTest(TestRecursorProtobuf):
 330     """
 331     This test makes sure that we correctly export outgoing queries over protobuf.
 332     It must be improved and setup env so we can check for incoming responses, but makes sure for now
 333     that the recursor at least connects to the protobuf server.
 334     """
 335
 336     _confdir = 'OutgoingProtobufDefault'
 337     _config_template = """
 338 auth-zones=example=configs/%s/example.zone""" % _confdir
 339     _lua_config_file = """
 340     outgoingProtobufServer({"127.0.0.1:%d", "127.0.0.1:%d"})
 341     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 342
 343     def testA(self):
 344         name = 'www.example.org.'
 345         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 346         query = dns.message.make_query(name, 'A', want_dnssec=True)
 347         query.flags |= dns.flags.RD
 348         res = self.sendUDPQuery(query)
 349
 350         # check the protobuf messages corresponding to the UDP query and answer
 351         msg = self.getFirstProtobufMessage()
 352         self.checkProtobufOutgoingQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
 353 #        # then the response
 354 #        msg = self.getFirstProtobufMessage()
 355 #        self.checkProtobufIncomingResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res)
 356         self.checkNoRemainingMessage()
 357
 358 class ProtobufMasksTest(TestRecursorProtobuf):
 359     """
 360     This test makes sure that we correctly export queries and response over protobuf, respecting the configured initiator masking.
 361     """
 362
 363     _confdir = 'ProtobufMasks'
 364     _config_template = """
 365 auth-zones=example=configs/%s/example.zone""" % _confdir
 366     _protobufMaskV4 = 4
 367     _protobufMaskV6 = 128
 368     _lua_config_file = """
 369     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"})
 370     setProtobufMasks(%d, %d)
 371     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port, _protobufMaskV4, _protobufMaskV6)
 372
 373     def testA(self):
 374         name = 'a.example.'
 375         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 376         query = dns.message.make_query(name, 'A', want_dnssec=True)
 377         query.flags |= dns.flags.CD
 378         res = self.sendUDPQuery(query)
 379         self.assertRRsetInAnswer(res, expected)
 380
 381         # check the protobuf messages corresponding to the UDP query and answer
 382         # but first let the protobuf messages the time to get there
 383         msg = self.getFirstProtobufMessage()
 384         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name, '112.0.0.0')
 385         # then the response
 386         msg = self.getFirstProtobufMessage()
 387         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res, '112.0.0.0')
 388         self.assertEquals(len(msg.response.rrs), 1)
 389         rr = msg.response.rrs[0]
 390         # we have max-cache-ttl set to 15
 391         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 392         self.assertEquals(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
 393         self.checkNoRemainingMessage()
 394
 395 class ProtobufQueriesOnlyTest(TestRecursorProtobuf):
 396     """
 397     This test makes sure that we correctly export queries but not responses over protobuf.
 398     """
 399
 400     _confdir = 'ProtobufQueriesOnly'
 401     _config_template = """
 402 auth-zones=example=configs/%s/example.zone""" % _confdir
 403     _lua_config_file = """
 404     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=true, logResponses=false } )
 405     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 406
 407     def testA(self):
 408         name = 'a.example.'
 409         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 410         query = dns.message.make_query(name, 'A', want_dnssec=True)
 411         query.flags |= dns.flags.CD
 412         res = self.sendUDPQuery(query)
 413         self.assertRRsetInAnswer(res, expected)
 414
 415         # check the protobuf message corresponding to the UDP query
 416         msg = self.getFirstProtobufMessage()
 417         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
 418         # no response
 419         self.checkNoRemainingMessage()
 420
 421 class ProtobufResponsesOnlyTest(TestRecursorProtobuf):
 422     """
 423     This test makes sure that we correctly export responses but not queries over protobuf.
 424     """
 425
 426     _confdir = 'ProtobufResponsesOnly'
 427     _config_template = """
 428 auth-zones=example=configs/%s/example.zone""" % _confdir
 429     _lua_config_file = """
 430     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=false, logResponses=true } )
 431     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 432
 433     def testA(self):
 434         name = 'a.example.'
 435         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 436         query = dns.message.make_query(name, 'A', want_dnssec=True)
 437         query.flags |= dns.flags.CD
 438         res = self.sendUDPQuery(query)
 439         self.assertRRsetInAnswer(res, expected)
 440
 441         # check the protobuf message corresponding to the UDP response
 442         msg = self.getFirstProtobufMessage()
 443         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res)
 444         self.assertEquals(len(msg.response.rrs), 1)
 445         rr = msg.response.rrs[0]
 446         # we have max-cache-ttl set to 15
 447         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 448         self.assertEquals(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.42')
 449         # nothing else in the queue
 450         self.checkNoRemainingMessage()
 451
 452 class ProtobufTaggedOnlyTest(TestRecursorProtobuf):
 453     """
 454     This test makes sure that we correctly export queries and responses but only if they have been tagged.
 455     """
 456
 457     _confdir = 'ProtobufTaggedOnly'
 458     _config_template = """
 459 auth-zones=example=configs/%s/example.zone""" % _confdir
 460     _lua_config_file = """
 461     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=true, logResponses=true, taggedOnly=true } )
 462     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 463     _tags = ['tag1', 'tag2']
 464     _tag_from_gettag = 'tag-from-gettag'
 465     _lua_dns_script_file = """
 466     function gettag(remote, ednssubnet, localip, qname, qtype, ednsoptions, tcp)
 467       if qname:equal('tagged.example.') then
 468         return 0, { '%s' }
 469       end
 470       return 0
 471     end
 472     function preresolve(dq)
 473       if dq.qname:equal('tagged.example.') then
 474         dq:addPolicyTag('%s')
 475         dq:addPolicyTag('%s')
 476       end
 477       return false
 478     end
 479     """ % (_tag_from_gettag, _tags[0], _tags[1])
 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         # the first query and answer are not tagged, so there is nothing in the queue
 491         time.sleep(1)
 492         self.checkNoRemainingMessage()
 493
 494     def testTagged(self):
 495         name = 'tagged.example.'
 496         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.84')
 497         query = dns.message.make_query(name, 'A', want_dnssec=True)
 498         query.flags |= dns.flags.CD
 499         res = self.sendUDPQuery(query)
 500         self.assertRRsetInAnswer(res, expected)
 501
 502         # check the protobuf messages corresponding to the UDP query and answer
 503         msg = self.getFirstProtobufMessage()
 504         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
 505         self.checkProtobufTags(msg, [self._tag_from_gettag])
 506         # then the response
 507         msg = self.getFirstProtobufMessage()
 508         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res)
 509         self.assertEquals(len(msg.response.rrs), 1)
 510         rr = msg.response.rrs[0]
 511         # we have max-cache-ttl set to 15
 512         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 513         self.assertEquals(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.84')
 514         tags = [self._tag_from_gettag] + self._tags
 515         self.checkProtobufTags(msg, tags)
 516         self.checkNoRemainingMessage()
 517
 518 class ProtobufSelectedFromLuaTest(TestRecursorProtobuf):
 519     """
 520     This test makes sure that we correctly export queries and responses but only if they have been selected from Lua.
 521     """
 522
 523     _confdir = 'ProtobufSelectedFromLua'
 524     _config_template = """
 525 auth-zones=example=configs/%s/example.zone""" % _confdir
 526     _lua_config_file = """
 527     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { logQueries=false, logResponses=false } )
 528     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 529     _lua_dns_script_file = """
 530     local ffi = require("ffi")
 531
 532     ffi.cdef[[
 533       typedef struct pdns_ffi_param pdns_ffi_param_t;
 534
 535       const char* pdns_ffi_param_get_qname(pdns_ffi_param_t* ref);
 536       void pdns_ffi_param_set_log_query(pdns_ffi_param_t* ref, bool logQuery);
 537     ]]
 538
 539     function gettag_ffi(obj)
 540       qname = ffi.string(ffi.C.pdns_ffi_param_get_qname(obj))
 541       if qname == 'query-selected.example' then
 542         ffi.C.pdns_ffi_param_set_log_query(obj, true)
 543       end
 544       return 0
 545     end
 546
 547     function preresolve(dq)
 548       if dq.qname:equal('answer-selected.example.') then
 549         dq.logResponse = true
 550       end
 551       return false
 552     end
 553     """
 554
 555     def testA(self):
 556         name = 'a.example.'
 557         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.42')
 558         query = dns.message.make_query(name, 'A', want_dnssec=True)
 559         query.flags |= dns.flags.CD
 560         res = self.sendUDPQuery(query)
 561         self.assertRRsetInAnswer(res, expected)
 562
 563         # check the protobuf message corresponding to the UDP response
 564         # the first query and answer are not selected, so there is nothing in the queue
 565         self.checkNoRemainingMessage()
 566
 567     def testQuerySelected(self):
 568         name = 'query-selected.example.'
 569         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.84')
 570         query = dns.message.make_query(name, 'A', want_dnssec=True)
 571         query.flags |= dns.flags.CD
 572         res = self.sendUDPQuery(query)
 573         self.assertRRsetInAnswer(res, expected)
 574
 575         # check the protobuf messages corresponding to the UDP query
 576         msg = self.getFirstProtobufMessage()
 577         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
 578         # there should be no response
 579         self.checkNoRemainingMessage()
 580
 581     def testResponseSelected(self):
 582         name = 'answer-selected.example.'
 583         expected = dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.84')
 584         query = dns.message.make_query(name, 'A', want_dnssec=True)
 585         query.flags |= dns.flags.CD
 586         res = self.sendUDPQuery(query)
 587         self.assertRRsetInAnswer(res, expected)
 588
 589         # check the protobuf messages corresponding to the UDP response
 590         msg = self.getFirstProtobufMessage()
 591         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res)
 592         self.assertEquals(len(msg.response.rrs), 1)
 593         rr = msg.response.rrs[0]
 594         # we have max-cache-ttl set to 15
 595         self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.A, name, 15)
 596         self.assertEquals(socket.inet_ntop(socket.AF_INET, rr.rdata), '192.0.2.84')
 597         self.checkNoRemainingMessage()
 598
 599 class ProtobufExportTypesTest(TestRecursorProtobuf):
 600     """
 601     This test makes sure that we correctly export other types than A, AAAA and CNAME over protobuf.
 602     """
 603
 604     _confdir = 'ProtobufExportTypes'
 605     _config_template = """
 606 auth-zones=example=configs/%s/example.zone""" % _confdir
 607     _lua_config_file = """
 608     protobufServer({"127.0.0.1:%d", "127.0.0.1:%d"}, { exportTypes={"AAAA", "MX", "SPF", "SRV", "TXT"} } )
 609     """ % (protobufServersParameters[0].port, protobufServersParameters[1].port)
 610
 611     def testA(self):
 612         name = 'types.example.'
 613         expected = [dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'A', '192.0.2.84'),
 614                     dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'AAAA', '2001:DB8::1'),
 615                     dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'MX', '10 a.example.'),
 616                     dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'SPF', '"v=spf1 -all"'),
 617                     dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'SRV', '10 20 443 a.example.'),
 618                     dns.rrset.from_text(name, 0, dns.rdataclass.IN, 'TXT', '"Lorem ipsum dolor sit amet"'),
 619                    ]
 620         query = dns.message.make_query(name, 'ANY', want_dnssec=True)
 621         query.flags |= dns.flags.CD
 622         raw = self.sendUDPQuery(query, decode=False)
 623         res = dns.message.from_wire(raw)
 624
 625         for rrset in expected:
 626             self.assertRRsetInAnswer(res, rrset)
 627
 628         # check the protobuf messages corresponding to the UDP query and answer
 629         msg = self.getFirstProtobufMessage()
 630         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
 631         # then the response
 632         msg = self.getFirstProtobufMessage()
 633         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, res, '127.0.0.1', receivedSize=len(raw))
 634         self.assertEquals(len(msg.response.rrs), 5)
 635         for rr in msg.response.rrs:
 636             self.assertTrue(rr.type in [dns.rdatatype.AAAA, dns.rdatatype.TXT, dns.rdatatype.MX, dns.rdatatype.SPF, dns.rdatatype.SRV])
 637
 638             if rr.type == dns.rdatatype.AAAA:
 639                 self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.AAAA, name, 15)
 640                 self.assertEquals(socket.inet_ntop(socket.AF_INET6, rr.rdata), '2001:db8::1')
 641             elif rr.type == dns.rdatatype.TXT:
 642                 self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.TXT, name, 15)
 643                 self.assertEquals(rr.rdata, '"Lorem ipsum dolor sit amet"')
 644             elif rr.type == dns.rdatatype.MX:
 645                 self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.MX, name, 15)
 646                 self.assertEquals(rr.rdata, 'a.example.')
 647             elif rr.type == dns.rdatatype.SPF:
 648                 self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.SPF, name, 15)
 649                 self.assertEquals(rr.rdata, '"v=spf1 -all"')
 650             elif rr.type == dns.rdatatype.SRV:
 651                 self.checkProtobufResponseRecord(rr, dns.rdataclass.IN, dns.rdatatype.SRV, name, 15)
 652                 self.assertEquals(rr.rdata, 'a.example.')
 653
 654         self.checkNoRemainingMessage()