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