1 # Copyright (c) 2012 The Chromium 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 """Traffic control library for constraining the network configuration on a port. 6 7 The traffic controller sets up a constrained network configuration on a port. 8 Traffic to the constrained port is forwarded to a specified server port. 9 """ 10 11 import logging 12 import os 13 import re 14 import subprocess 15 16 # The maximum bandwidth limit. 17 _DEFAULT_MAX_BANDWIDTH_KBIT = 1000000 18 19 20 class TrafficControlError(BaseException): 21 """Exception raised for errors in traffic control library. 22 23 Attributes: 24 msg: User defined error message. 25 cmd: Command for which the exception was raised. 26 returncode: Return code of running the command. 27 stdout: Output of running the command. 28 stderr: Error output of running the command. 29 """ 30 31 def __init__(self, msg, cmd=None, returncode=None, output=None, 32 error=None): 33 BaseException.__init__(self, msg) 34 self.msg = msg 35 self.cmd = cmd 36 self.returncode = returncode 37 self.output = output 38 self.error = error 39 40 41 def CheckRequirements(): 42 """Checks if permissions are available to run traffic control commands. 43 44 Raises: 45 TrafficControlError: If permissions to run traffic control commands are not 46 available. 47 """ 48 if os.geteuid() != 0: 49 _Exec(['sudo', '-n', 'tc', '-help'], 50 msg=('Cannot run \'tc\' command. Traffic Control must be run as root ' 51 'or have password-less sudo access to this command.')) 52 _Exec(['sudo', '-n', 'iptables', '-help'], 53 msg=('Cannot run \'iptables\' command. Traffic Control must be run ' 54 'as root or have password-less sudo access to this command.')) 55 56 57 def CreateConstrainedPort(config): 58 """Creates a new constrained port. 59 60 Imposes packet level constraints such as bandwidth, latency, and packet loss 61 on a given port using the specified configuration dictionary. Traffic to that 62 port is forwarded to a specified server port. 63 64 Args: 65 config: Constraint configuration dictionary, format: 66 port: Port to constrain (integer 1-65535). 67 server_port: Port to redirect traffic on [port] to (integer 1-65535). 68 interface: Network interface name (string). 69 latency: Delay added on each packet sent (integer in ms). 70 bandwidth: Maximum allowed upload bandwidth (integer in kbit/s). 71 loss: Percentage of packets to drop (integer 0-100). 72 73 Raises: 74 TrafficControlError: If any operation fails. The message in the exception 75 describes what failed. 76 """ 77 _CheckArgsExist(config, 'interface', 'port', 'server_port') 78 _AddRootQdisc(config['interface']) 79 80 try: 81 _ConfigureClass('add', config) 82 _AddSubQdisc(config) 83 _AddFilter(config['interface'], config['port']) 84 _AddIptableRule(config['interface'], config['port'], config['server_port']) 85 except TrafficControlError as e: 86 logging.debug('Error creating constrained port %d.\nError: %s\n' 87 'Deleting constrained port.', config['port'], e.error) 88 DeleteConstrainedPort(config) 89 raise e 90 91 92 def DeleteConstrainedPort(config): 93 """Deletes an existing constrained port. 94 95 Deletes constraints set on a given port and the traffic forwarding rule from 96 the constrained port to a specified server port. 97 98 The original constrained network configuration used to create the constrained 99 port must be passed in. 100 101 Args: 102 config: Constraint configuration dictionary, format: 103 port: Port to constrain (integer 1-65535). 104 server_port: Port to redirect traffic on [port] to (integer 1-65535). 105 interface: Network interface name (string). 106 bandwidth: Maximum allowed upload bandwidth (integer in kbit/s). 107 108 Raises: 109 TrafficControlError: If any operation fails. The message in the exception 110 describes what failed. 111 """ 112 _CheckArgsExist(config, 'interface', 'port', 'server_port') 113 try: 114 # Delete filters first so it frees the class. 115 _DeleteFilter(config['interface'], config['port']) 116 finally: 117 try: 118 # Deleting the class deletes attached qdisc as well. 119 _ConfigureClass('del', config) 120 finally: 121 _DeleteIptableRule(config['interface'], config['port'], 122 config['server_port']) 123 124 125 def TearDown(config): 126 """Deletes the root qdisc and all iptables rules. 127 128 Args: 129 config: Constraint configuration dictionary, format: 130 interface: Network interface name (string). 131 132 Raises: 133 TrafficControlError: If any operation fails. The message in the exception 134 describes what failed. 135 """ 136 _CheckArgsExist(config, 'interface') 137 138 command = ['sudo', 'tc', 'qdisc', 'del', 'dev', config['interface'], 'root'] 139 try: 140 _Exec(command, msg='Could not delete root qdisc.') 141 finally: 142 _DeleteAllIpTableRules() 143 144 145 def _CheckArgsExist(config, *args): 146 """Check that the args exist in config dictionary and are not None. 147 148 Args: 149 config: Any dictionary. 150 *args: The list of key names to check. 151 152 Raises: 153 TrafficControlError: If any key name does not exist in config or is None. 154 """ 155 for key in args: 156 if key not in config.keys() or config[key] is None: 157 raise TrafficControlError('Missing "%s" parameter.' % key) 158 159 160 def _AddRootQdisc(interface): 161 """Sets up the default root qdisc. 162 163 Args: 164 interface: Network interface name. 165 166 Raises: 167 TrafficControlError: If adding the root qdisc fails for a reason other than 168 it already exists. 169 """ 170 command = ['sudo', 'tc', 'qdisc', 'add', 'dev', interface, 'root', 'handle', 171 '1:', 'htb'] 172 try: 173 _Exec(command, msg=('Error creating root qdisc. ' 174 'Make sure you have root access')) 175 except TrafficControlError as e: 176 # Ignore the error if root already exists. 177 if not 'File exists' in e.error: 178 raise e 179 180 181 def _ConfigureClass(option, config): 182 """Adds or deletes a class and qdisc attached to the root. 183 184 The class specifies bandwidth, and qdisc specifies delay and packet loss. The 185 class ID is based on the config port. 186 187 Args: 188 option: Adds or deletes a class option [add|del]. 189 config: Constraint configuration dictionary, format: 190 port: Port to constrain (integer 1-65535). 191 interface: Network interface name (string). 192 bandwidth: Maximum allowed upload bandwidth (integer in kbit/s). 193 """ 194 # Use constrained port as class ID so we can attach the qdisc and filter to 195 # it, as well as delete the class, using only the port number. 196 class_id = '1:%x' % config['port'] 197 if 'bandwidth' not in config.keys() or not config['bandwidth']: 198 bandwidth = _DEFAULT_MAX_BANDWIDTH_KBIT 199 else: 200 bandwidth = config['bandwidth'] 201 202 bandwidth = '%dkbit' % bandwidth 203 command = ['sudo', 'tc', 'class', option, 'dev', config['interface'], 204 'parent', '1:', 'classid', class_id, 'htb', 'rate', bandwidth, 205 'ceil', bandwidth] 206 _Exec(command, msg=('Error configuring class ID %s using "%s" command.' % 207 (class_id, option))) 208 209 210 def _AddSubQdisc(config): 211 """Adds a qdisc attached to the class identified by the config port. 212 213 Args: 214 config: Constraint configuration dictionary, format: 215 port: Port to constrain (integer 1-65535). 216 interface: Network interface name (string). 217 latency: Delay added on each packet sent (integer in ms). 218 loss: Percentage of packets to drop (integer 0-100). 219 """ 220 port_hex = '%x' % config['port'] 221 class_id = '1:%x' % config['port'] 222 command = ['sudo', 'tc', 'qdisc', 'add', 'dev', config['interface'], 'parent', 223 class_id, 'handle', port_hex + ':0', 'netem'] 224 225 # Check if packet-loss is set in the configuration. 226 if 'loss' in config.keys() and config['loss']: 227 loss = '%d%%' % config['loss'] 228 command.extend(['loss', loss]) 229 # Check if latency is set in the configuration. 230 if 'latency' in config.keys() and config['latency']: 231 latency = '%dms' % config['latency'] 232 command.extend(['delay', latency]) 233 234 _Exec(command, msg='Could not attach qdisc to class ID %s.' % class_id) 235 236 237 def _AddFilter(interface, port): 238 """Redirects packets coming to a specified port into the constrained class. 239 240 Args: 241 interface: Interface name to attach the filter to (string). 242 port: Port number to filter packets with (integer 1-65535). 243 """ 244 class_id = '1:%x' % port 245 246 command = ['sudo', 'tc', 'filter', 'add', 'dev', interface, 'protocol', 'ip', 247 'parent', '1:', 'prio', '1', 'u32', 'match', 'ip', 'sport', port, 248 '0xffff', 'flowid', class_id] 249 _Exec(command, msg='Error adding filter on port %d.' % port) 250 251 252 def _DeleteFilter(interface, port): 253 """Deletes the filter attached to the configured port. 254 255 Args: 256 interface: Interface name the filter is attached to (string). 257 port: Port number being filtered (integer 1-65535). 258 """ 259 handle_id = _GetFilterHandleId(interface, port) 260 command = ['sudo', 'tc', 'filter', 'del', 'dev', interface, 'protocol', 'ip', 261 'parent', '1:0', 'handle', handle_id, 'prio', '1', 'u32'] 262 _Exec(command, msg='Error deleting filter on port %d.' % port) 263 264 265 def _GetFilterHandleId(interface, port): 266 """Searches for the handle ID of the filter identified by the config port. 267 268 Args: 269 interface: Interface name the filter is attached to (string). 270 port: Port number being filtered (integer 1-65535). 271 272 Returns: 273 The handle ID. 274 275 Raises: 276 TrafficControlError: If handle ID was not found. 277 """ 278 command = ['sudo', 'tc', 'filter', 'list', 'dev', interface, 'parent', '1:'] 279 output = _Exec(command, msg='Error listing filters.') 280 # Search for the filter handle ID associated with class ID '1:port'. 281 handle_id_re = re.search( 282 '([0-9a-fA-F]{3}::[0-9a-fA-F]{3}).*(?=flowid 1:%x\s)' % port, output) 283 if handle_id_re: 284 return handle_id_re.group(1) 285 raise TrafficControlError(('Could not find filter handle ID for class ID ' 286 '1:%x.') % port) 287 288 289 def _AddIptableRule(interface, port, server_port): 290 """Forwards traffic from constrained port to a specified server port. 291 292 Args: 293 interface: Interface name to attach the filter to (string). 294 port: Port of incoming packets (integer 1-65535). 295 server_port: Server port to forward the packets to (integer 1-65535). 296 """ 297 # Preroute rules for accessing the port through external connections. 298 command = ['sudo', 'iptables', '-t', 'nat', '-A', 'PREROUTING', '-i', 299 interface, '-p', 'tcp', '--dport', port, '-j', 'REDIRECT', 300 '--to-port', server_port] 301 _Exec(command, msg='Error adding iptables rule for port %d.' % port) 302 303 # Output rules for accessing the rule through localhost or 127.0.0.1 304 command = ['sudo', 'iptables', '-t', 'nat', '-A', 'OUTPUT', '-p', 'tcp', 305 '--dport', port, '-j', 'REDIRECT', '--to-port', server_port] 306 _Exec(command, msg='Error adding iptables rule for port %d.' % port) 307 308 309 def _DeleteIptableRule(interface, port, server_port): 310 """Deletes the iptable rule associated with specified port number. 311 312 Args: 313 interface: Interface name to attach the filter to (string). 314 port: Port of incoming packets (integer 1-65535). 315 server_port: Server port packets are forwarded to (integer 1-65535). 316 """ 317 command = ['sudo', 'iptables', '-t', 'nat', '-D', 'PREROUTING', '-i', 318 interface, '-p', 'tcp', '--dport', port, '-j', 'REDIRECT', 319 '--to-port', server_port] 320 _Exec(command, msg='Error deleting iptables rule for port %d.' % port) 321 322 command = ['sudo', 'iptables', '-t', 'nat', '-D', 'OUTPUT', '-p', 'tcp', 323 '--dport', port, '-j', 'REDIRECT', '--to-port', server_port] 324 _Exec(command, msg='Error adding iptables rule for port %d.' % port) 325 326 327 def _DeleteAllIpTableRules(): 328 """Deletes all iptables rules.""" 329 command = ['sudo', 'iptables', '-t', 'nat', '-F'] 330 _Exec(command, msg='Error deleting all iptables rules.') 331 332 333 def _Exec(command, msg=None): 334 """Executes a command. 335 336 Args: 337 command: Command list to execute. 338 msg: Message describing the error in case the command fails. 339 340 Returns: 341 The standard output from running the command. 342 343 Raises: 344 TrafficControlError: If command fails. Message is set by the msg parameter. 345 """ 346 cmd_list = [str(x) for x in command] 347 cmd = ' '.join(cmd_list) 348 logging.debug('Running command: %s', cmd) 349 350 p = subprocess.Popen(cmd_list, stdout=subprocess.PIPE, stderr=subprocess.PIPE) 351 output, error = p.communicate() 352 if p.returncode != 0: 353 raise TrafficControlError(msg, cmd, p.returncode, output, error) 354 return output.strip() 355