regression-tests.dnsdist/test_Lua.py

   1 #!/usr/bin/env python
   2
   3 import base64
   4 import dns
   5 import time
   6 import unittest
   7 from dnsdisttests import DNSDistTest
   8
   9 class TestLuaThread(DNSDistTest):
  10     _consoleKey = DNSDistTest.generateConsoleKey()
  11     _consoleKeyB64 = base64.b64encode(_consoleKey).decode('ascii')
  12
  13     _config_params = ['_consoleKeyB64', '_consolePort']
  14     _config_template = """
  15     setKey("%s")
  16     controlSocket("127.0.0.1:%s")
  17
  18     counter = 0
  19     function threadmessage(cmd, data)
  20         print("counter says", cmd, data.i)
  21         counter = tonumber(data.i)
  22     end
  23
  24     newThread([==[
  25       local socket = require'socket'
  26       local i=1
  27       while true
  28       do
  29           socket.sleep(1)
  30           submitToMainThread("setCounter", {i=i})
  31           i = i + 1
  32       end
  33     ]==])
  34     """
  35
  36     def testLuaThreadCounter(self):
  37         """
  38         LuaThread: Test the lua newThread interface
  39         """
  40         count1 = self.sendConsoleCommand('counter')
  41         time.sleep(3)
  42         count2 = self.sendConsoleCommand('counter')
  43         self.assertTrue(count2 > count1)
  44
  45 class TestLuaDNSHeaderBindings(DNSDistTest):
  46     _config_template = """
  47     newServer{address="127.0.0.1:%s"}
  48
  49     function checkTCSet(dq)
  50       local tc = dq.dh:getTC()
  51       if not tc then
  52         return DNSAction.Spoof, 'tc-not-set.check-tc.lua-dnsheaders.tests.powerdns.com.'
  53       end
  54       return DNSAction.Allow
  55     end
  56
  57     addAction('check-tc.lua-dnsheaders.tests.powerdns.com.', LuaAction(checkTCSet))
  58     """
  59
  60     def testLuaGetTC(self):
  61         """
  62         LuaDNSHeaders: TC
  63         """
  64         name = 'notset.check-tc.lua-dnsheaders.tests.powerdns.com.'
  65         query = dns.message.make_query(name, 'A', 'IN')
  66         # dnsdist set RA = RD for spoofed responses
  67         query.flags &= ~dns.flags.RD
  68         response = dns.message.make_response(query)
  69         rrset = dns.rrset.from_text(name,
  70                                     60,
  71                                     dns.rdataclass.IN,
  72                                     dns.rdatatype.CNAME,
  73                                     'tc-not-set.check-tc.lua-dnsheaders.tests.powerdns.com.')
  74         response.answer.append(rrset)
  75         for method in ("sendUDPQuery", "sendTCPQuery"):
  76             sender = getattr(self, method)
  77             (_, receivedResponse) = sender(query, response=None, useQueue=False)
  78             self.assertEqual(response, receivedResponse)
  79
  80         name = 'set.check-tc.lua-dnsheaders.tests.powerdns.com.'
  81         query = dns.message.make_query(name, 'A', 'IN')
  82         response = dns.message.make_response(query)
  83         rrset = dns.rrset.from_text(name,
  84                                     60,
  85                                     dns.rdataclass.IN,
  86                                     dns.rdatatype.A,
  87                                     '127.0.0.1')
  88         response.answer.append(rrset)
  89         query.flags |= dns.flags.TC
  90         for method in ("sendUDPQuery", "sendTCPQuery"):
  91             sender = getattr(self, method)
  92             (receivedQuery, receivedResponse) = sender(query, response)
  93             receivedQuery.id = query.id
  94             self.assertEqual(query, receivedQuery)
  95             self.assertEqual(response, receivedResponse)