media/tools/constrained_network_server/traffic_control.py

   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()