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
16import logging
17import platformsettings
18import 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".
23BANDWIDTH_PATTERN = r'0|\d+[KM]?(bit|Byte)/s'
24
25
26class TrafficShaperException(Exception):
27  pass
28
29
30class 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
39class 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)
187