1# Copyright 2016 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"""Start and stop tsproxy.""" 6 7import logging 8import os 9import re 10import signal 11import subprocess 12import sys 13import time 14 15try: 16 import fcntl 17except ImportError: 18 fcntl = None 19 20import py_utils 21from py_utils import retry_util 22from py_utils import atexit_with_log 23 24_TSPROXY_PATH = os.path.join( 25 py_utils.GetCatapultDir(), 'third_party', 'tsproxy', 'tsproxy.py') 26 27class TsProxyServerError(Exception): 28 """Catch-all exception for tsProxy Server.""" 29 pass 30 31def ParseTsProxyPortFromOutput(output_line): 32 port_re = re.compile( 33 r'Started Socks5 proxy server on ' 34 r'(?P<host>[^:]*):' 35 r'(?P<port>\d+)') 36 m = port_re.match(output_line) 37 if m: 38 return int(m.group('port')) 39 40 41class TsProxyServer(object): 42 """Start and stop tsproxy. 43 44 TsProxy provides basic latency, download and upload traffic shaping. This 45 class provides a programming API to the tsproxy script in 46 catapult/third_party/tsproxy/tsproxy.py 47 48 This class can be used as a context manager. 49 """ 50 51 def __init__(self, host_ip=None, http_port=None, https_port=None): 52 """ 53 Initialize TsProxyServer. 54 55 Args: 56 57 host_ip: A string of the host ip address. 58 http_port: A decimal of the port used for http traffic. 59 https_port: a decimal of the port used for https traffic. 60 61 """ 62 self._proc = None 63 self._port = None 64 self._is_running = False 65 self._host_ip = host_ip 66 assert bool(http_port) == bool(https_port) 67 self._http_port = http_port 68 self._https_port = https_port 69 self._non_blocking = False 70 self._rtt = None 71 self._inbkps = None 72 self._outkbps = None 73 74 @property 75 def port(self): 76 return self._port 77 78 @retry_util.RetryOnException(TsProxyServerError, retries=3) 79 def StartServer(self, timeout=10, retries=None): 80 """Start TsProxy server and verify that it started.""" 81 del retries # Handled by decorator. 82 cmd_line = [sys.executable, _TSPROXY_PATH] 83 # Use port 0 so tsproxy picks a random available port. 84 cmd_line.extend(['--port=0']) 85 if self._host_ip: 86 cmd_line.append('--desthost=%s' % self._host_ip) 87 if self._http_port: 88 cmd_line.append( 89 '--mapports=443:%s,*:%s' % (self._https_port, self._http_port)) 90 logging.info('Tsproxy commandline: %s', cmd_line) 91 self._proc = subprocess.Popen( 92 cmd_line, stdout=subprocess.PIPE, stdin=subprocess.PIPE, 93 stderr=subprocess.PIPE, bufsize=1) 94 self._non_blocking = False 95 if fcntl: 96 logging.info('fcntl is supported, trying to set ' 97 'non blocking I/O for the ts_proxy process') 98 fd = self._proc.stdout.fileno() 99 fl = fcntl.fcntl(fd, fcntl.F_GETFL) 100 fcntl.fcntl(fd, fcntl.F_SETFL, fl | os.O_NONBLOCK) 101 self._non_blocking = True 102 103 atexit_with_log.Register(self.StopServer) 104 try: 105 py_utils.WaitFor(self._IsStarted, timeout) 106 logging.info('TsProxy port: %s', self._port) 107 self._is_running = True 108 except py_utils.TimeoutException: 109 err = self.StopServer() 110 if err: 111 logging.error('Error stopping WPR server:\n%s', err) 112 raise TsProxyServerError( 113 'Error starting tsproxy: timed out after %s seconds' % timeout) 114 115 def _IsStarted(self): 116 assert not self._is_running 117 assert self._proc 118 if self._proc.poll() is not None: 119 return False 120 self._proc.stdout.flush() 121 output_line = self._ReadLineTsProxyStdout(timeout=5) 122 logging.debug('TsProxy output: %s', output_line) 123 self._port = ParseTsProxyPortFromOutput(output_line) 124 return self._port != None 125 126 def _ReadLineTsProxyStdout(self, timeout): 127 def ReadSingleLine(): 128 try: 129 return self._proc.stdout.readline().strip() 130 except IOError: 131 # Add a sleep to avoid trying to read self._proc.stdout too often. 132 if self._non_blocking: 133 time.sleep(0.5) 134 return None 135 return py_utils.WaitFor(ReadSingleLine, timeout) 136 137 @retry_util.RetryOnException(TsProxyServerError, retries=3) 138 def _IssueCommand(self, command_string, timeout, retries=None): 139 del retries # handled by the decorator 140 logging.info('Issuing command to ts_proxy_server: %s', command_string) 141 command_output = [] 142 self._proc.stdin.write('%s\n' % command_string) 143 def CommandStatusIsRead(): 144 self._proc.stdin.flush() 145 self._proc.stdout.flush() 146 command_output.append(self._ReadLineTsProxyStdout(timeout)) 147 return command_output[-1] == 'OK' or command_output[-1] == 'ERROR' 148 149 py_utils.WaitFor(CommandStatusIsRead, timeout) 150 151 success = 'OK' in command_output 152 logging.log(logging.DEBUG if success else logging.ERROR, 153 'TsProxy output:\n%s', '\n'.join(command_output)) 154 if not success: 155 raise TsProxyServerError('Failed to execute command: %s', command_string) 156 157 def UpdateOutboundPorts(self, http_port, https_port, timeout=5): 158 assert http_port and https_port 159 assert http_port != https_port 160 assert isinstance(http_port, int) and isinstance(https_port, int) 161 assert 1 <= http_port <= 65535 162 assert 1 <= https_port <= 65535 163 self._IssueCommand('set mapports 443:%i,*:%i' % (https_port, http_port), 164 timeout) 165 166 def UpdateTrafficSettings( 167 self, round_trip_latency_ms=None, 168 download_bandwidth_kbps=None, upload_bandwidth_kbps=None, timeout=20): 169 """Update traffic settings of the proxy server. 170 171 Notes that this method only updates the specified parameter. 172 """ 173 # Memorize the traffic settings & only execute the command if the traffic 174 # settings are different. 175 if round_trip_latency_ms is not None and self._rtt != round_trip_latency_ms: 176 self._IssueCommand('set rtt %s' % round_trip_latency_ms, timeout) 177 self._rtt = round_trip_latency_ms 178 179 if (download_bandwidth_kbps is not None and 180 self._inbkps != download_bandwidth_kbps): 181 self._IssueCommand('set inkbps %s' % download_bandwidth_kbps, timeout) 182 self._inbkps = download_bandwidth_kbps 183 184 if (upload_bandwidth_kbps is not None and 185 self._outkbps != upload_bandwidth_kbps): 186 self._IssueCommand('set outkbps %s' % upload_bandwidth_kbps, timeout) 187 self._outkbps = upload_bandwidth_kbps 188 189 def StopServer(self): 190 """Stop TsProxy Server.""" 191 if not self._is_running: 192 logging.debug('Attempting to stop TsProxy server that is not running.') 193 return 194 if not self._proc: 195 return 196 try: 197 self._IssueCommand('exit', timeout=10) 198 py_utils.WaitFor(lambda: self._proc.poll() is not None, 10) 199 except py_utils.TimeoutException: 200 # signal.SIGINT is not supported on Windows. 201 if not sys.platform.startswith('win'): 202 try: 203 # Use a SIGNINT so that it can do graceful cleanup 204 self._proc.send_signal(signal.SIGINT) 205 except ValueError: 206 logging.warning('Unable to stop ts_proxy_server gracefully.\n') 207 self._proc.terminate() 208 _, err = self._proc.communicate() 209 210 self._proc = None 211 self._port = None 212 self._is_running = False 213 self._rtt = None 214 self._inbkps = None 215 self._outkbps = None 216 return err 217 218 def __enter__(self): 219 """Add support for with-statement.""" 220 self.StartServer() 221 return self 222 223 def __exit__(self, unused_exc_type, unused_exc_val, unused_exc_tb): 224 """Add support for with-statement.""" 225 self.StopServer() 226