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