1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
|
# Copyright (C) Dnspython Contributors, see LICENSE for text of ISC license
import contextlib
import functools
import socket
import struct
import threading
import trio
import dns.message
import dns.rcode
import dns.trio.query
class Server(threading.Thread):
"""The nanoserver is a nameserver skeleton suitable for faking a DNS
server for various testing purposes. It executes with a trio run
loop in a dedicated thread, and is a context manager. Exiting the
context manager will ensure the server shuts down.
If a port is not specified, random ports will be chosen.
Applications should subclass the server and override the handle()
method to determine how the server responds to queries. The
default behavior is to refuse everything.
If use_thread is set to False in the constructor, then the
server's main() method can be used directly in a trio nursery,
allowing the server's cancellation to be managed in the Trio way.
In this case, no thread creation ever happens even though Server
is a subclass of thread, because the start() method is never
called.
"""
def __init__(self, address='127.0.0.1', port=0, enable_udp=True,
enable_tcp=True, use_thread=True):
super().__init__()
self.address = address
self.port = port
self.enable_udp = enable_udp
self.enable_tcp = enable_tcp
self.use_thread = use_thread
self.left = None
self.right = None
self.udp = None
self.udp_address = None
self.tcp = None
self.tcp_address = None
def __enter__(self):
(self.left, self.right) = socket.socketpair()
# We're making the UDP socket now so it can be sent to by the
# caller immediately (i.e. no race with the listener starting
# in the thread).
if self.enable_udp:
self.udp = socket.socket(socket.AF_INET, socket.SOCK_DGRAM, 0)
self.udp.bind((self.address, self.port))
self.udp_address = self.udp.getsockname()
if self.enable_tcp:
self.tcp = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0)
self.tcp.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
self.tcp.bind((self.address, self.port))
self.tcp.listen()
self.tcp_address = self.udp.getsockname()
if self.use_thread:
self.start()
return self
def __exit__(self, ex_ty, ex_va, ex_tr):
if self.left:
self.left.close()
if self.use_thread and self.is_alive():
self.join()
if self.right:
self.right.close()
if self.udp:
self.udp.close()
if self.tcp:
self.tcp.close()
async def wait_for_input_or_eof(self):
#
# This trio task just waits for input on the right half of the
# socketpair (the left half is owned by the context manager
# returned by launch). As soon as something is read, or the
# socket returns EOF, EOFError is raised, causing a the
# nursery to cancel all other nursery tasks, in particular the
# listeners.
#
try:
with trio.socket.from_stdlib_socket(self.right) as sock:
self.right = None # we own cleanup
await sock.recv(1)
finally:
raise EOFError
def handle(self, message):
#
# Handle message 'message'. Override this method to change
# how the server behaves.
#
# The return value is either a dns.message.Message or a bytes.
# We allow a bytes to be returned for cases where handle wants
# to return an invalid DNS message for testing purposes.
#
r = dns.message.make_response(message)
r.set_rcode(dns.rcode.REFUSED)
return r
async def serve_udp(self):
with trio.socket.from_stdlib_socket(self.udp) as sock:
self.udp = None # we own cleanup
while True:
try:
(wire, from_address) = await sock.recvfrom(65535)
q = dns.message.from_wire(wire)
r = self.handle(q)
if isinstance(r, dns.message.Message):
wire = r.to_wire()
else:
wire = r
await sock.sendto(wire, from_address)
except Exception:
pass
async def serve_tcp(self, stream):
try:
while True:
ldata = await dns.trio.query.read_exactly(stream, 2)
(l,) = struct.unpack("!H", ldata)
wire = await dns.trio.query.read_exactly(stream, l)
q = dns.message.from_wire(wire)
r = self.handle(q)
if isinstance(r, dns.message.Message):
wire = r.to_wire()
else:
wire = r
l = len(wire)
stream_message = struct.pack("!H", l) + wire
await stream.send_all(stream_message)
except Exception:
pass
async def orchestrate_tcp(self):
with trio.socket.from_stdlib_socket(self.tcp) as sock:
self.tcp = None # we own cleanup
listener = trio.SocketListener(sock)
async with trio.open_nursery() as nursery:
serve = functools.partial(trio.serve_listeners, self.serve_tcp,
[listener], handler_nursery=nursery)
nursery.start_soon(serve)
async def main(self):
try:
async with trio.open_nursery() as nursery:
if self.use_thread:
nursery.start_soon(self.wait_for_input_or_eof)
if self.enable_udp:
nursery.start_soon(self.serve_udp)
if self.enable_tcp:
nursery.start_soon(self.orchestrate_tcp)
except Exception:
pass
def run(self):
if not self.use_thread:
raise RuntimeError('start() called on a use_thread=False Server')
trio.run(self.main)
if __name__ == "__main__":
import sys
import time
async def trio_main():
try:
with Server(port=5354, use_thread=False) as server:
print(f'Trio mode: listening on UDP: {server.udp_address}, ' +
f'TCP: {server.tcp_address}')
async with trio.open_nursery() as nursery:
nursery.start_soon(server.main)
except Exception:
pass
def threaded_main():
with Server(port=5354) as server:
print(f'Thread Mode: listening on UDP: {server.udp_address}, ' +
f'TCP: {server.tcp_address}')
time.sleep(300)
if len(sys.argv) > 1 and sys.argv[1] == 'trio':
trio.run(trio_main)
else:
threaded_main()
|
