1# Lint as: python2, python3
2# Copyright (c) 2012 The Chromium OS Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6"""
7Programmable testing DHCP server.
8
9Simple DHCP server you can program with expectations of future packets and
10responses to those packets.  The server is basically a thin wrapper around a
11server socket with some utility logic to make setting up tests easier.  To write
12a test, you start a server, construct a sequence of handling rules.
13
14Handling rules let you set up expectations of future packets of certain types.
15Handling rules are processed in order, and only the first remaining handler
16handles a given packet.  In theory you could write the entire test into a single
17handling rule and keep an internal state machine for how far that handler has
18gotten through the test.  This would be poor style however.  Correct style is to
19write (or reuse) a handler for each packet the server should see, leading us to
20a happy land where any conceivable packet handler has already been written for
21us.
22
23Example usage:
24
25# Start up the DHCP server, which will ignore packets until a test is started
26server = DhcpTestServer(interface=interface_name)
27server.start()
28
29# Given a list of handling rules, start a test with a 30 sec timeout.
30handling_rules = []
31handling_rules.append(DhcpHandlingRule_RespondToDiscovery(intended_ip,
32                                                          intended_subnet_mask,
33                                                          dhcp_server_ip,
34                                                          lease_time_seconds)
35server.start_test(handling_rules, 30.0)
36
37# Trigger DHCP clients to do various test related actions
38...
39
40# Get results
41server.wait_for_test_to_finish()
42if (server.last_test_passed):
43    ...
44else:
45    ...
46
47
48Note that if you make changes, make sure that the tests in dhcp_unittest.py
49still pass.
50"""
51
52from __future__ import absolute_import
53from __future__ import division
54from __future__ import print_function
55
56import logging
57from six.moves import range
58import socket
59import threading
60import time
61import traceback
62
63from autotest_lib.client.cros import dhcp_packet
64from autotest_lib.client.cros import dhcp_handling_rule
65
66# From socket.h
67SO_BINDTODEVICE = 25
68
69class DhcpTestServer(threading.Thread):
70    def __init__(self,
71                 interface=None,
72                 ingress_address="<broadcast>",
73                 ingress_port=67,
74                 broadcast_address="255.255.255.255",
75                 broadcast_port=68):
76        super(DhcpTestServer, self).__init__()
77        self._mutex = threading.Lock()
78        self._ingress_address = ingress_address
79        self._ingress_port = ingress_port
80        self._broadcast_port = broadcast_port
81        self._broadcast_address = broadcast_address
82        self._socket = None
83        self._interface = interface
84        self._stopped = False
85        self._test_in_progress = False
86        self._last_test_passed = False
87        self._test_timeout = 0
88        self._handling_rules = []
89        self._logger = logging.getLogger("dhcp.test_server")
90        self._exception = None
91        self.daemon = False
92
93    @property
94    def stopped(self):
95        with self._mutex:
96            return self._stopped
97
98    @property
99    def is_healthy(self):
100        with self._mutex:
101            return self._socket is not None
102
103    @property
104    def test_in_progress(self):
105        with self._mutex:
106            return self._test_in_progress
107
108    @property
109    def last_test_passed(self):
110        with self._mutex:
111            return self._last_test_passed
112
113    @property
114    def current_rule(self):
115        """
116        Return the currently active DhcpHandlingRule.
117        """
118        with self._mutex:
119            return self._handling_rules[0]
120
121    def start(self):
122        """
123        Start the DHCP server.  Only call this once.
124        """
125        if self.is_alive():
126            return False
127        self._logger.info("DhcpTestServer started; opening sockets.")
128        try:
129            self._socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
130            self._logger.info("Opening socket on '%s' port %d." %
131                              (self._ingress_address, self._ingress_port))
132            self._socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
133            self._socket.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
134            if self._interface is not None:
135                self._logger.info("Binding to %s" % self._interface)
136                self._socket.setsockopt(socket.SOL_SOCKET,
137                                        SO_BINDTODEVICE,
138                                        self._interface)
139            self._socket.bind((self._ingress_address, self._ingress_port))
140            # Wait 100 ms for a packet, then return, thus keeping the thread
141            # active but mostly idle.
142            self._socket.settimeout(0.1)
143        except socket.error as socket_error:
144            self._logger.error("Socket error: %s." % str(socket_error))
145            self._logger.error(traceback.format_exc())
146            if not self._socket is None:
147                self._socket.close()
148            self._socket = None
149            self._logger.error("Failed to open server socket.  Aborting.")
150            return
151        super(DhcpTestServer, self).start()
152
153    def stop(self):
154        """
155        Stop the DHCP server and free its socket.
156        """
157        with self._mutex:
158            self._stopped = True
159
160    def start_test(self, handling_rules, test_timeout_seconds):
161        """
162        Start a new test using |handling_rules|.  The server will call the
163        test successfull if it receives a RESPONSE_IGNORE_SUCCESS (or
164        RESPONSE_RESPOND_SUCCESS) from a handling_rule before
165        |test_timeout_seconds| passes.  If the timeout passes without that
166        message, the server runs out of handling rules, or a handling rule
167        return RESPONSE_FAIL, the test is ended and marked as not passed.
168
169        All packets received before start_test() is called are received and
170        ignored.
171        """
172        with self._mutex:
173            self._test_timeout = time.time() + test_timeout_seconds
174            self._handling_rules = handling_rules
175            self._test_in_progress = True
176            self._last_test_passed = False
177            self._exception = None
178
179    def wait_for_test_to_finish(self):
180        """
181        Block on the test finishing in a CPU friendly way.  Timeouts, successes,
182        and failures count as finishes.
183        """
184        while self.test_in_progress:
185            time.sleep(0.1)
186        if self._exception:
187            raise self._exception
188
189    def abort_test(self):
190        """
191        Abort a test prematurely, counting the test as a failure.
192        """
193        with self._mutex:
194            self._logger.info("Manually aborting test.")
195            self._end_test_unsafe(False)
196
197    def _teardown(self):
198        with self._mutex:
199            self._socket.close()
200            self._socket = None
201
202    def _end_test_unsafe(self, passed):
203        if not self._test_in_progress:
204            return
205        if passed:
206            self._logger.info("DHCP server says test passed.")
207        else:
208            self._logger.info("DHCP server says test failed.")
209        self._test_in_progress = False
210        self._last_test_passed = passed
211
212    def _send_response_unsafe(self, packet):
213        if packet is None:
214            self._logger.error("Handling rule failed to return a packet.")
215            return False
216        self._logger.debug("Sending response: %s" % packet)
217        binary_string = packet.to_binary_string()
218        if binary_string is None or len(binary_string) < 1:
219            self._logger.error("Packet failed to serialize to binary string.")
220            return False
221
222        self._socket.sendto(binary_string,
223                            (self._broadcast_address, self._broadcast_port))
224        return True
225
226    def _loop_body(self):
227        with self._mutex:
228            if self._test_in_progress and self._test_timeout < time.time():
229                # The test has timed out, so we abort it.  However, we should
230                # continue to accept packets, so we fall through.
231                self._logger.error("Test in progress has timed out.")
232                self._end_test_unsafe(False)
233            try:
234                data, _ = self._socket.recvfrom(1024)
235                self._logger.info("Server received packet of length %d." %
236                                   len(data))
237            except socket.timeout:
238                # No packets available, lets return and see if the server has
239                # been shut down in the meantime.
240                return
241
242            # Receive packets when no test is in progress, just don't process
243            # them.
244            if not self._test_in_progress:
245                return
246
247            packet = dhcp_packet.DhcpPacket(byte_str=data)
248            if not packet.is_valid:
249                self._logger.warning("Server received an invalid packet over a "
250                                     "DHCP port?")
251                return
252
253            logging.debug("Server received a DHCP packet: %s." % packet)
254            if len(self._handling_rules) < 1:
255                self._logger.info("No handling rule for packet: %s." %
256                                  str(packet))
257                self._end_test_unsafe(False)
258                return
259
260            handling_rule = self._handling_rules[0]
261            response_code = handling_rule.handle(packet)
262            logging.info("Handler gave response: %d" % response_code)
263            if response_code & dhcp_handling_rule.RESPONSE_POP_HANDLER:
264                self._handling_rules.pop(0)
265
266            if response_code & dhcp_handling_rule.RESPONSE_HAVE_RESPONSE:
267                for response_instance in range(
268                        handling_rule.response_packet_count):
269                    response = handling_rule.respond(packet)
270                    if not self._send_response_unsafe(response):
271                        self._logger.error(
272                                "Failed to send packet, ending test.")
273                        self._end_test_unsafe(False)
274                        return
275
276            if response_code & dhcp_handling_rule.RESPONSE_TEST_FAILED:
277                self._logger.info("Handling rule %s rejected packet %s." %
278                                  (handling_rule, packet))
279                self._end_test_unsafe(False)
280                return
281
282            if response_code & dhcp_handling_rule.RESPONSE_TEST_SUCCEEDED:
283                self._end_test_unsafe(True)
284                return
285
286    def run(self):
287        """
288        Main method of the thread.  Never call this directly, since it assumes
289        some setup done in start().
290        """
291        with self._mutex:
292            if self._socket is None:
293                self._logger.error("Failed to create server socket, exiting.")
294                return
295
296        self._logger.info("DhcpTestServer entering handling loop.")
297        while not self.stopped:
298            try:
299                self._loop_body()
300                # Python does not have waiting queues on Lock objects.
301                # Give other threads a change to hold the mutex by
302                # forcibly releasing the GIL while we sleep.
303                time.sleep(0.01)
304            except Exception as e:
305                with self._mutex:
306                    self._end_test_unsafe(False)
307                    self._exception = e
308        with self._mutex:
309            self._end_test_unsafe(False)
310        self._logger.info("DhcpTestServer closing sockets.")
311        self._teardown()
312        self._logger.info("DhcpTestServer exiting.")
313