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