1#!/usr/bin/env python2.7 2# Copyright 2015 gRPC authors. 3# 4# Licensed under the Apache License, Version 2.0 (the "License"); 5# you may not use this file except in compliance with the License. 6# You may obtain a copy of the License at 7# 8# http://www.apache.org/licenses/LICENSE-2.0 9# 10# Unless required by applicable law or agreed to in writing, software 11# distributed under the License is distributed on an "AS IS" BASIS, 12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13# See the License for the specific language governing permissions and 14# limitations under the License. 15"""Starts a local DNS server for use in tests""" 16 17import argparse 18import os 19import platform 20import signal 21import sys 22import threading 23import time 24 25import twisted 26import twisted.internet 27import twisted.internet.defer 28import twisted.internet.protocol 29import twisted.internet.reactor 30import twisted.internet.threads 31import twisted.names 32from twisted.names import authority 33from twisted.names import client 34from twisted.names import common 35from twisted.names import dns 36from twisted.names import server 37import twisted.names.client 38import twisted.names.dns 39import twisted.names.server 40import yaml 41 42_SERVER_HEALTH_CHECK_RECORD_NAME = 'health-check-local-dns-server-is-alive.resolver-tests.grpctestingexp' # missing end '.' for twisted syntax 43_SERVER_HEALTH_CHECK_RECORD_DATA = '123.123.123.123' 44 45 46class NoFileAuthority(authority.FileAuthority): 47 48 def __init__(self, soa, records): 49 # skip FileAuthority 50 common.ResolverBase.__init__(self) 51 self.soa = soa 52 self.records = records 53 54 55def start_local_dns_server(args): 56 all_records = {} 57 58 def _push_record(name, r): 59 name = name.encode('ascii') 60 print('pushing record: |%s|' % name) 61 if all_records.get(name) is not None: 62 all_records[name].append(r) 63 return 64 all_records[name] = [r] 65 66 def _maybe_split_up_txt_data(name, txt_data, r_ttl): 67 txt_data = txt_data.encode('ascii') 68 start = 0 69 txt_data_list = [] 70 while len(txt_data[start:]) > 0: 71 next_read = len(txt_data[start:]) 72 if next_read > 255: 73 next_read = 255 74 txt_data_list.append(txt_data[start:start + next_read]) 75 start += next_read 76 _push_record(name, dns.Record_TXT(*txt_data_list, ttl=r_ttl)) 77 78 with open(args.records_config_path) as config: 79 test_records_config = yaml.load(config) 80 common_zone_name = test_records_config['resolver_tests_common_zone_name'] 81 for group in test_records_config['resolver_component_tests']: 82 for name in group['records'].keys(): 83 for record in group['records'][name]: 84 r_type = record['type'] 85 r_data = record['data'] 86 r_ttl = int(record['TTL']) 87 record_full_name = '%s.%s' % (name, common_zone_name) 88 assert record_full_name[-1] == '.' 89 record_full_name = record_full_name[:-1] 90 if r_type == 'A': 91 _push_record(record_full_name, 92 dns.Record_A(r_data, ttl=r_ttl)) 93 if r_type == 'AAAA': 94 _push_record(record_full_name, 95 dns.Record_AAAA(r_data, ttl=r_ttl)) 96 if r_type == 'SRV': 97 p, w, port, target = r_data.split(' ') 98 p = int(p) 99 w = int(w) 100 port = int(port) 101 target_full_name = ( 102 '%s.%s' % (target, common_zone_name)).encode('ascii') 103 _push_record( 104 record_full_name, 105 dns.Record_SRV(p, w, port, target_full_name, ttl=r_ttl)) 106 if r_type == 'TXT': 107 _maybe_split_up_txt_data(record_full_name, r_data, r_ttl) 108 # Add an optional IPv4 record is specified 109 if args.add_a_record: 110 extra_host, extra_host_ipv4 = args.add_a_record.split(':') 111 _push_record(extra_host, dns.Record_A(extra_host_ipv4, ttl=0)) 112 # Server health check record 113 _push_record(_SERVER_HEALTH_CHECK_RECORD_NAME, 114 dns.Record_A(_SERVER_HEALTH_CHECK_RECORD_DATA, ttl=0)) 115 soa_record = dns.Record_SOA(mname=common_zone_name.encode('ascii')) 116 test_domain_com = NoFileAuthority( 117 soa=(common_zone_name.encode('ascii'), soa_record), 118 records=all_records, 119 ) 120 server = twisted.names.server.DNSServerFactory( 121 authorities=[test_domain_com], verbose=2) 122 server.noisy = 2 123 twisted.internet.reactor.listenTCP(args.port, server) 124 dns_proto = twisted.names.dns.DNSDatagramProtocol(server) 125 dns_proto.noisy = 2 126 twisted.internet.reactor.listenUDP(args.port, dns_proto) 127 print('starting local dns server on 127.0.0.1:%s' % args.port) 128 print('starting twisted.internet.reactor') 129 twisted.internet.reactor.suggestThreadPoolSize(1) 130 twisted.internet.reactor.run() 131 132 133def _quit_on_signal(signum, _frame): 134 print('Received SIGNAL %d. Quitting with exit code 0' % signum) 135 twisted.internet.reactor.stop() 136 sys.stdout.flush() 137 sys.exit(0) 138 139 140def flush_stdout_loop(): 141 num_timeouts_so_far = 0 142 sleep_time = 1 143 # Prevent zombies. Tests that use this server are short-lived. 144 max_timeouts = 60 * 10 145 while num_timeouts_so_far < max_timeouts: 146 sys.stdout.flush() 147 time.sleep(sleep_time) 148 num_timeouts_so_far += 1 149 print('Process timeout reached, or cancelled. Exitting 0.') 150 os.kill(os.getpid(), signal.SIGTERM) 151 152 153def main(): 154 argp = argparse.ArgumentParser( 155 description='Local DNS Server for resolver tests') 156 argp.add_argument('-p', 157 '--port', 158 default=None, 159 type=int, 160 help='Port for DNS server to listen on for TCP and UDP.') 161 argp.add_argument( 162 '-r', 163 '--records_config_path', 164 default=None, 165 type=str, 166 help=('Directory of resolver_test_record_groups.yaml file. ' 167 'Defaults to path needed when the test is invoked as part ' 168 'of run_tests.py.')) 169 argp.add_argument( 170 '--add_a_record', 171 default=None, 172 type=str, 173 help=('Add an A record via the command line. Useful for when we ' 174 'need to serve a one-off A record that is under a ' 175 'different domain then the rest the records configured in ' 176 '--records_config_path (which all need to be under the ' 177 'same domain). Format: <name>:<ipv4 address>')) 178 args = argp.parse_args() 179 signal.signal(signal.SIGTERM, _quit_on_signal) 180 signal.signal(signal.SIGINT, _quit_on_signal) 181 output_flush_thread = threading.Thread(target=flush_stdout_loop) 182 output_flush_thread.setDaemon(True) 183 output_flush_thread.start() 184 start_local_dns_server(args) 185 186 187if __name__ == '__main__': 188 main() 189