tools/telemetry/third_party/webpagereplay/trafficshaper.py

   1 #!/usr/bin/env python
   2 # Copyright 2010 Google Inc. All Rights Reserved.
   3 #
   4 # Licensed under the Apache License, Version 2.0 (the "License");
   5 # you may not use this file except in compliance with the License.
   6 # You may obtain a copy of the License at
   7 #
   8 #      http://www.apache.org/licenses/LICENSE-2.0
   9 #
  10 # Unless required by applicable law or agreed to in writing, software
  11 # distributed under the License is distributed on an "AS IS" BASIS,
  12 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13 # See the License for the specific language governing permissions and
  14 # limitations under the License.
  15
  16 import logging
  17 import platformsettings
  18 import re
  19
  20
  21 # Mac has broken bandwitdh parsing, so double check the values.
  22 # On Mac OS X 10.6, "KBit/s" actually uses "KByte/s".
  23 BANDWIDTH_PATTERN = r'0|\d+[KM]?(bit|Byte)/s'
  24
  25
  26 class TrafficShaperException(Exception):
  27   pass
  28
  29
  30 class BandwidthValueError(TrafficShaperException):
  31   def __init__(self, value):  # pylint: disable=super-init-not-called
  32     self.value = value
  33
  34   def __str__(self):
  35     return 'Value, "%s", does not match regex: %s' % (
  36         self.value, BANDWIDTH_PATTERN)
  37
  38
  39 class TrafficShaper(object):
  40   """Manages network traffic shaping."""
  41
  42   # Pick webpagetest-compatible values (details: http://goo.gl/oghTg).
  43   _UPLOAD_PIPE = '10'      # Enforces overall upload bandwidth.
  44   _UPLOAD_QUEUE = '10'     # Shares upload bandwidth among source ports.
  45   _UPLOAD_RULE = '5000'    # Specifies when the upload queue is used.
  46   _DOWNLOAD_PIPE = '11'    # Enforces overall download bandwidth.
  47   _DOWNLOAD_QUEUE = '11'   # Shares download bandwidth among destination ports.
  48   _DOWNLOAD_RULE = '5100'  # Specifies when the download queue is used.
  49   _QUEUE_SLOTS = 100       # Number of packets to queue.
  50
  51   _BANDWIDTH_RE = re.compile(BANDWIDTH_PATTERN)
  52
  53   def __init__(self,
  54                dont_use=None,
  55                host='127.0.0.1',
  56                ports=None,
  57                up_bandwidth='0',
  58                down_bandwidth='0',
  59                delay_ms='0',
  60                packet_loss_rate='0',
  61                init_cwnd='0',
  62                use_loopback=True):
  63     """Start shaping traffic.
  64
  65     Args:
  66       host: a host string (name or IP) for the web proxy.
  67       ports: a list of ports to shape traffic on.
  68       up_bandwidth: Upload bandwidth
  69       down_bandwidth: Download bandwidth
  70            Bandwidths measured in [K|M]{bit/s|Byte/s}. '0' means unlimited.
  71       delay_ms: Propagation delay in milliseconds. '0' means no delay.
  72       packet_loss_rate: Packet loss rate in range [0..1]. '0' means no loss.
  73       init_cwnd: the initial cwnd setting. '0' means no change.
  74       use_loopback: True iff shaping is done on the loopback (or equiv) adapter.
  75     """
  76     assert dont_use is None  # Force args to be named.
  77     self.host = host
  78     self.ports = ports
  79     self.up_bandwidth = up_bandwidth
  80     self.down_bandwidth = down_bandwidth
  81     self.delay_ms = delay_ms
  82     self.packet_loss_rate = packet_loss_rate
  83     self.init_cwnd = init_cwnd
  84     self.use_loopback = use_loopback
  85     if not self._BANDWIDTH_RE.match(self.up_bandwidth):
  86       raise BandwidthValueError(self.up_bandwidth)
  87     if not self._BANDWIDTH_RE.match(self.down_bandwidth):
  88       raise BandwidthValueError(self.down_bandwidth)
  89     self.is_shaping = False
  90
  91   def __enter__(self):
  92     if self.use_loopback:
  93       platformsettings.setup_temporary_loopback_config()
  94     if self.init_cwnd != '0':
  95       platformsettings.set_temporary_tcp_init_cwnd(self.init_cwnd)
  96     try:
  97       ipfw_list = platformsettings.ipfw('list')
  98       if not ipfw_list.startswith('65535 '):
  99         logging.warn('ipfw has existing rules:\n%s', ipfw_list)
 100         self._delete_rules(ipfw_list)
 101     except Exception:
 102       pass
 103     if (self.up_bandwidth == '0' and self.down_bandwidth == '0' and
 104         self.delay_ms == '0' and self.packet_loss_rate == '0'):
 105       logging.info('Skipped shaping traffic.')
 106       return
 107     if not self.ports:
 108       raise TrafficShaperException('No ports on which to shape traffic.')
 109
 110     ports = ','.join(str(p) for p in self.ports)
 111     half_delay_ms = int(self.delay_ms) / 2  # split over up/down links
 112
 113     try:
 114       # Configure upload shaping.
 115       platformsettings.ipfw(
 116           'pipe', self._UPLOAD_PIPE,
 117           'config',
 118           'bw', self.up_bandwidth,
 119           'delay', half_delay_ms,
 120           )
 121       platformsettings.ipfw(
 122           'queue', self._UPLOAD_QUEUE,
 123           'config',
 124           'pipe', self._UPLOAD_PIPE,
 125           'plr', self.packet_loss_rate,
 126           'queue', self._QUEUE_SLOTS,
 127           'mask', 'src-port', '0xffff',
 128           )
 129       platformsettings.ipfw(
 130           'add', self._UPLOAD_RULE,
 131           'queue', self._UPLOAD_QUEUE,
 132           'ip',
 133           'from', 'any',
 134           'to', self.host,
 135           self.use_loopback and 'out' or 'in',
 136           'dst-port', ports,
 137           )
 138       self.is_shaping = True
 139
 140       # Configure download shaping.
 141       platformsettings.ipfw(
 142           'pipe', self._DOWNLOAD_PIPE,
 143           'config',
 144           'bw', self.down_bandwidth,
 145           'delay', half_delay_ms,
 146           )
 147       platformsettings.ipfw(
 148           'queue', self._DOWNLOAD_QUEUE,
 149           'config',
 150           'pipe', self._DOWNLOAD_PIPE,
 151           'plr', self.packet_loss_rate,
 152           'queue', self._QUEUE_SLOTS,
 153           'mask', 'dst-port', '0xffff',
 154           )
 155       platformsettings.ipfw(
 156           'add', self._DOWNLOAD_RULE,
 157           'queue', self._DOWNLOAD_QUEUE,
 158           'ip',
 159           'from', self.host,
 160           'to', 'any',
 161           'out',
 162           'src-port', ports,
 163           )
 164       logging.info('Started shaping traffic')
 165     except Exception:
 166       logging.error('Unable to shape traffic.')
 167       raise
 168
 169   def __exit__(self, unused_exc_type, unused_exc_val, unused_exc_tb):
 170     if self.is_shaping:
 171       try:
 172         self._delete_rules()
 173         logging.info('Stopped shaping traffic')
 174       except Exception:
 175         logging.error('Unable to stop shaping traffic.')
 176         raise
 177
 178   def _delete_rules(self, ipfw_list=None):
 179     if ipfw_list is None:
 180       ipfw_list = platformsettings.ipfw('list')
 181     existing_rules = set(
 182         r.split()[0].lstrip('0') for r in ipfw_list.splitlines())
 183     delete_rules = [r for r in (self._DOWNLOAD_RULE, self._UPLOAD_RULE)
 184                     if r in existing_rules]
 185     if delete_rules:
 186       platformsettings.ipfw('delete', *delete_rules)