projects / thirdparty / pdns.git / blame
| Commit | Line | Data |
|---|---|---|
| 5fc8a17f RG |
1 | #!/usr/bin/env python |
| 2 | import unittest | |
| 3 | import dns | |
| 4 | import lmdb | |
| 5 | import socket | |
| 6 | from dnsdisttests import DNSDistTest | |
| 7 | ||
| 8 | class TestLMDB(DNSDistTest): | |
| 9 | ||
| 10 | _lmdbFileName = '/tmp/test-lmdb-db' | |
| 11 | _lmdbDBName = 'db-name' | |
| 12 | _config_template = """ | |
| 13 | newServer{address="127.0.0.1:%d"} | |
| 14 | ||
| 15 | kvs = newLMDBKVStore('%s', '%s') | |
| 16 | ||
| 17 | -- KVS lookups follow | |
| 18 | -- does a lookup in the LMDB database using the source IP as key, and store the result into the 'kvs-sourceip-result' tag | |
| 19 | addAction(AllRule(), KeyValueStoreLookupAction(kvs, KeyValueLookupKeySourceIP(), 'kvs-sourceip-result')) | |
| 20 | ||
| 21 | -- does a lookup in the LMDB database using the qname in wire format as key, and store the result into the 'kvs-qname-result' tag | |
| 22 | addAction(AllRule(), KeyValueStoreLookupAction(kvs, KeyValueLookupKeyQName(), 'kvs-qname-result')) | |
| 23 | ||
| 24 | -- if the value of the 'kvs-qname-result' tag is set to 'this is the value of the qname tag' | |
| 25 | -- does a lookup in the LMDB database using the value of the 'kvs-qname-result' tag as key, and store the result into the 'kvs-tag-result' tag | |
| 26 | addAction(TagRule('kvs-qname-result', 'this is the value of the qname tag'), KeyValueStoreLookupAction(kvs, KeyValueLookupKeyTag('kvs-qname-result'), 'kvs-tag-result')) | |
| 27 | ||
| 28 | -- does a lookup in the LMDB database using the source IP as key, and store the result into the 'kvs-sourceip-result' tag | |
| 29 | addAction(AllRule(), KeyValueStoreLookupAction(kvs, KeyValueLookupKeySourceIP(), 'kvs-sourceip-result')) | |
| 30 | ||
| 31 | -- does a lookup in the LMDB database using the qname in wire format as key, but this time does a suffix lookup, and store the result into the 'kvs-suffix-result' tag | |
| 32 | addAction(AllRule(), KeyValueStoreLookupAction(kvs, KeyValueLookupKeySuffix(), 'kvs-suffix-result')) | |
| 33 | ||
| 34 | -- Now we take action based on the result of the lookups | |
| 35 | -- if the value of the 'kvs-tag-result' is set to 'this is the value of the second tag', spoof a response | |
| 36 | addAction(TagRule('kvs-tag-result', 'this is the value of the second tag'), SpoofAction('1.2.3.4')) | |
| 37 | ||
| 38 | -- if the value of the 'kvs-suffix-result' is set to 'this is the value of the suffix tag', spoof a response | |
| 39 | addAction(TagRule('kvs-suffix-result', 'this is the value of the suffix tag'), SpoofAction('42.42.42.42')) | |
| 40 | ||
| 41 | -- if the value of the 'kvs-sourceip-result' is set to 'this is the value of the source address tag', spoof a response | |
| 42 | addAction(TagRule('kvs-sourceip-result', 'this is the value of the source address tag'), SpoofAction('5.6.7.8')) | |
| 43 | ||
| 44 | -- otherwise, spoof a different response | |
| 45 | addAction(AllRule(), SpoofAction('9.9.9.9')) | |
| 46 | """ | |
| 47 | _config_params = ['_testServerPort', '_lmdbFileName', '_lmdbDBName'] | |
| 48 | ||
| 49 | @classmethod | |
| 50 | def setUpLMDB(cls): | |
| 51 | env = lmdb.open(cls._lmdbFileName, map_size=1014*1024, max_dbs=1024, subdir=False) | |
| 52 | db = env.open_db(key=cls._lmdbDBName.encode()) | |
| 53 | with env.begin(db=db, write=True) as txn: | |
| 54 | txn.put(b'\x05qname\x04lmdb\x05tests\x08powerdns\x03com\x00', b'this is the value of the qname tag') | |
| 55 | txn.put(socket.inet_aton('127.0.0.1'), b'this is the value of the source address tag') | |
| 56 | txn.put(b'this is the value of the qname tag', b'this is the value of the second tag') | |
| 57 | txn.put(b'\x06suffix\x04lmdb\x05tests\x08powerdns\x03com\x00', b'this is the value of the suffix tag') | |
| 58 | ||
| 59 | @classmethod | |
| 60 | def setUpClass(cls): | |
| 61 | ||
| 62 | cls.setUpLMDB() | |
| 63 | cls.startResponders() | |
| 64 | cls.startDNSDist() | |
| 65 | cls.setUpSockets() | |
| 66 | ||
| 67 | print("Launching tests..") | |
| 68 | ||
| 69 | def testLMDBSource(self): | |
| 70 | """ | |
| 71 | LMDB: Match on source address | |
| 72 | """ | |
| 73 | name = 'source-ip.lmdb.tests.powerdns.com.' | |
| 74 | query = dns.message.make_query(name, 'A', 'IN') | |
| 75 | # dnsdist set RA = RD for spoofed responses | |
| 76 | query.flags &= ~dns.flags.RD | |
| 77 | expectedResponse = dns.message.make_response(query) | |
| 78 | rrset = dns.rrset.from_text(name, | |
| 79 | 3600, | |
| 80 | dns.rdataclass.IN, | |
| 81 | dns.rdatatype.A, | |
| 82 | '5.6.7.8') | |
| 83 | expectedResponse.answer.append(rrset) | |
| 84 | ||
| 85 | for method in ("sendUDPQuery", "sendTCPQuery"): | |
| 86 | sender = getattr(self, method) | |
| 87 | (receivedQuery, receivedResponse) = sender(query, response=None, useQueue=False) | |
| 88 | self.assertFalse(receivedQuery) | |
| 89 | self.assertTrue(receivedResponse) | |
| 90 | self.assertEquals(expectedResponse, receivedResponse) | |
| 91 | ||
| 92 | def testLMDBQNamePlusTagLookup(self): | |
| 93 | """ | |
| 94 | LMDB: Match on qname then does a second lookup using the value of the first lookup | |
| 95 | """ | |
| 96 | name = 'qname.lmdb.tests.powerdns.com.' | |
| 97 | query = dns.message.make_query(name, 'A', 'IN') | |
| 98 | # dnsdist set RA = RD for spoofed responses | |
| 99 | query.flags &= ~dns.flags.RD | |
| 100 | expectedResponse = dns.message.make_response(query) | |
| 101 | rrset = dns.rrset.from_text(name, | |
| 102 | 3600, | |
| 103 | dns.rdataclass.IN, | |
| 104 | dns.rdatatype.A, | |
| 105 | '1.2.3.4') | |
| 106 | expectedResponse.answer.append(rrset) | |
| 107 | ||
| 108 | for method in ("sendUDPQuery", "sendTCPQuery"): | |
| 109 | sender = getattr(self, method) | |
| 110 | (receivedQuery, receivedResponse) = sender(query, response=None, useQueue=False) | |
| 111 | self.assertFalse(receivedQuery) | |
| 112 | self.assertTrue(receivedResponse) | |
| 113 | self.assertEquals(expectedResponse, receivedResponse) | |
| 114 | ||
| 115 | def testLMDBSuffixLookup(self): | |
| 116 | """ | |
| 117 | LMDB: Match on the qname via a suffix lookup | |
| 118 | """ | |
| 119 | name = 'sub.sub.suffix.lmdb.tests.powerdns.com.' | |
| 120 | query = dns.message.make_query(name, 'A', 'IN') | |
| 121 | # dnsdist set RA = RD for spoofed responses | |
| 122 | query.flags &= ~dns.flags.RD | |
| 123 | expectedResponse = dns.message.make_response(query) | |
| 124 | rrset = dns.rrset.from_text(name, | |
| 125 | 3600, | |
| 126 | dns.rdataclass.IN, | |
| 127 | dns.rdatatype.A, | |
| 128 | '42.42.42.42') | |
| 129 | expectedResponse.answer.append(rrset) | |
| 130 | ||
| 131 | for method in ("sendUDPQuery", "sendTCPQuery"): | |
| 132 | sender = getattr(self, method) | |
| 133 | (receivedQuery, receivedResponse) = sender(query, response=None, useQueue=False) | |
| 134 | self.assertFalse(receivedQuery) | |
| 135 | self.assertTrue(receivedResponse) | |
| 136 | self.assertEquals(expectedResponse, receivedResponse) | |
| 137 |