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