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