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