1# Copyright (c) 2015 The Chromium OS 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 5import httplib 6import logging 7import socket 8import tempfile 9import time 10import xmlrpclib 11 12import common 13from autotest_lib.client.bin import utils 14from autotest_lib.client.common_lib import error 15from autotest_lib.client.common_lib.cros import retry 16 17try: 18 import jsonrpclib 19except ImportError: 20 jsonrpclib = None 21 22 23class RpcServerTracker(object): 24 """ 25 This class keeps track of all the RPC server connections started on a remote 26 host. The caller can use either |xmlrpc_connect| or |jsonrpc_connect| to 27 start the required type of rpc server on the remote host. 28 The host will cleanup all the open RPC server connections on disconnect. 29 """ 30 31 _RPC_PROXY_URL_FORMAT = 'http://localhost:%d' 32 _RPC_HOST_ADDRESS_FORMAT = 'localhost:%d' 33 _RPC_SHUTDOWN_POLLING_PERIOD_SECONDS = 2 34 _RPC_SHUTDOWN_TIMEOUT_SECONDS = 10 35 36 def __init__(self, host): 37 """ 38 @param port: The host object associated with this instance of 39 RpcServerTracker. 40 """ 41 self._host = host 42 self._rpc_proxy_map = {} 43 44 45 def _setup_port(self, port, command_name, remote_pid=None): 46 """Sets up a tunnel process and register it to rpc_server_tracker. 47 48 Chrome OS on the target closes down most external ports for security. 49 We could open the port, but doing that would conflict with security 50 tests that check that only expected ports are open. So, to get to 51 the port on the target we use an ssh tunnel. 52 53 This method assumes that xmlrpc and jsonrpc never conflict, since 54 we can only either have an xmlrpc or a jsonrpc server listening on 55 a remote port. As such, it enforces a single proxy->remote port 56 policy, i.e if one starts a jsonrpc proxy/server from port A->B, 57 and then tries to start an xmlrpc proxy forwarded to the same port, 58 the xmlrpc proxy will override the jsonrpc tunnel process, however: 59 60 1. None of the methods on the xmlrpc proxy will work because 61 the server listening on B is jsonrpc. 62 63 2. The xmlrpc client cannot initiate a termination of the JsonRPC 64 server, as the only use case currently is goofy, which is tied to 65 the factory image. It is much easier to handle a failed xmlrpc 66 call on the client than it is to terminate goofy in this scenario, 67 as doing the latter might leave the DUT in a hard to recover state. 68 69 With the current implementation newer rpc proxy connections will 70 terminate the tunnel processes of older rpc connections tunneling 71 to the same remote port. If methods are invoked on the client 72 after this has happened they will fail with connection closed errors. 73 74 @param port: The remote forwarding port. 75 @param command_name: The name of the remote process, to terminate 76 using pkill. 77 @param remote_pid: The PID of the remote background process 78 as a string. 79 80 @return the local port which is used for port forwarding on the ssh 81 client. 82 """ 83 self.disconnect(port) 84 local_port = utils.get_unused_port() 85 tunnel_proc = self._host.create_ssh_tunnel(port, local_port) 86 self._rpc_proxy_map[port] = (command_name, tunnel_proc, remote_pid) 87 return local_port 88 89 90 def _setup_rpc(self, port, command_name, remote_pid=None): 91 """Construct a URL for an rpc connection using ssh tunnel. 92 93 @param port: The remote forwarding port. 94 @param command_name: The name of the remote process, to terminate 95 using pkill. 96 @param remote_pid: The PID of the remote background process 97 as a string. 98 99 @return a url that we can use to initiate the rpc connection. 100 """ 101 return self._RPC_PROXY_URL_FORMAT % self._setup_port( 102 port, command_name, remote_pid=remote_pid) 103 104 105 def tunnel_connect(self, port): 106 """Construct a host address using ssh tunnel. 107 108 @param port: The remote forwarding port. 109 110 @return a host address using ssh tunnel. 111 """ 112 return self._RPC_HOST_ADDRESS_FORMAT % self._setup_port(port, None) 113 114 115 def xmlrpc_connect(self, command, port, command_name=None, 116 ready_test_name=None, timeout_seconds=10, 117 logfile=None, request_timeout_seconds=None): 118 """Connect to an XMLRPC server on the host. 119 120 The `command` argument should be a simple shell command that 121 starts an XMLRPC server on the given `port`. The command 122 must not daemonize, and must terminate cleanly on SIGTERM. 123 The command is started in the background on the host, and a 124 local XMLRPC client for the server is created and returned 125 to the caller. 126 127 Note that the process of creating an XMLRPC client makes no 128 attempt to connect to the remote server; the caller is 129 responsible for determining whether the server is running 130 correctly, and is ready to serve requests. 131 132 Optionally, the caller can pass ready_test_name, a string 133 containing the name of a method to call on the proxy. This 134 method should take no parameters and return successfully only 135 when the server is ready to process client requests. When 136 ready_test_name is set, xmlrpc_connect will block until the 137 proxy is ready, and throw a TestError if the server isn't 138 ready by timeout_seconds. 139 140 If a server is already running on the remote port, this 141 method will kill it and disconnect the tunnel process 142 associated with the connection before establishing a new one, 143 by consulting the rpc_proxy_map in disconnect. 144 145 @param command Shell command to start the server. 146 @param port Port number on which the server is expected to 147 be serving. 148 @param command_name String to use as input to `pkill` to 149 terminate the XMLRPC server on the host. 150 @param ready_test_name String containing the name of a 151 method defined on the XMLRPC server. 152 @param timeout_seconds Number of seconds to wait 153 for the server to become 'ready.' Will throw a 154 TestFail error if server is not ready in time. 155 @param logfile Logfile to send output when running 156 'command' argument. 157 @param request_timeout_seconds Timeout in seconds for an XMLRPC request. 158 159 """ 160 # Clean up any existing state. If the caller is willing 161 # to believe their server is down, we ought to clean up 162 # any tunnels we might have sitting around. 163 self.disconnect(port) 164 remote_pid = None 165 if command is not None: 166 if logfile: 167 remote_cmd = '%s > %s 2>&1' % (command, logfile) 168 else: 169 remote_cmd = command 170 remote_pid = self._host.run_background(remote_cmd) 171 logging.debug('Started XMLRPC server on host %s, pid = %s', 172 self._host.hostname, remote_pid) 173 174 # Tunnel through SSH to be able to reach that remote port. 175 rpc_url = self._setup_rpc(port, command_name, remote_pid=remote_pid) 176 if request_timeout_seconds is not None: 177 proxy = TimeoutXMLRPCServerProxy( 178 rpc_url, timeout=request_timeout_seconds, allow_none=True) 179 else: 180 proxy = xmlrpclib.ServerProxy(rpc_url, allow_none=True) 181 182 if ready_test_name is not None: 183 # retry.retry logs each attempt; calculate delay_sec to 184 # keep log spam to a dull roar. 185 @retry.retry((socket.error, 186 xmlrpclib.ProtocolError, 187 httplib.BadStatusLine), 188 timeout_min=timeout_seconds / 60.0, 189 delay_sec=min(max(timeout_seconds / 20.0, 0.1), 1)) 190 def ready_test(): 191 """ Call proxy.ready_test_name(). """ 192 try: 193 getattr(proxy, ready_test_name)() 194 except socket.error as e: 195 e.filename = rpc_url.replace('http://', '') 196 raise 197 successful = False 198 try: 199 logging.info('Waiting %d seconds for XMLRPC server ' 200 'to start.', timeout_seconds) 201 ready_test() 202 successful = True 203 except socket.error as e: 204 e.filename = rpc_url.replace('http://', '') 205 raise 206 finally: 207 if not successful: 208 logging.error('Failed to start XMLRPC server.') 209 if logfile: 210 with tempfile.NamedTemporaryFile() as temp: 211 self._host.get_file(logfile, temp.name) 212 logging.error('The log of XML RPC server:\n%s', 213 open(temp.name).read()) 214 self.disconnect(port) 215 logging.info('XMLRPC server started successfully.') 216 return proxy 217 218 219 def jsonrpc_connect(self, port): 220 """Creates a jsonrpc proxy connection through an ssh tunnel. 221 222 This method exists to facilitate communication with goofy (which is 223 the default system manager on all factory images) and as such, leaves 224 most of the rpc server sanity checking to the caller. Unlike 225 xmlrpc_connect, this method does not facilitate the creation of a remote 226 jsonrpc server, as the only clients of this code are factory tests, 227 for which the goofy system manager is built in to the image and starts 228 when the target boots. 229 230 One can theoretically create multiple jsonrpc proxies all forwarded 231 to the same remote port, provided the remote port has an rpc server 232 listening. However, in doing so we stand the risk of leaking an 233 existing tunnel process, so we always disconnect any older tunnels 234 we might have through disconnect. 235 236 @param port: port on the remote host that is serving this proxy. 237 238 @return: The client proxy. 239 """ 240 if not jsonrpclib: 241 logging.warning('Jsonrpclib could not be imported. Check that ' 242 'site-packages contains jsonrpclib.') 243 return None 244 245 proxy = jsonrpclib.jsonrpc.ServerProxy(self._setup_rpc(port, None)) 246 247 logging.info('Established a jsonrpc connection through port %s.', port) 248 return proxy 249 250 251 def disconnect(self, port): 252 """Disconnect from an RPC server on the host. 253 254 Terminates the remote RPC server previously started for 255 the given `port`. Also closes the local ssh tunnel created 256 for the connection to the host. This function does not 257 directly alter the state of a previously returned RPC 258 client object; however disconnection will cause all 259 subsequent calls to methods on the object to fail. 260 261 This function does nothing if requested to disconnect a port 262 that was not previously connected via _setup_rpc. 263 264 @param port Port number passed to a previous call to 265 `_setup_rpc()`. 266 """ 267 if port not in self._rpc_proxy_map: 268 return 269 remote_name, tunnel_proc, remote_pid = self._rpc_proxy_map[port] 270 if remote_name: 271 # We use 'pkill' to find our target process rather than 272 # a PID, because the host may have rebooted since 273 # connecting, and we don't want to kill an innocent 274 # process with the same PID. 275 # 276 # 'pkill' helpfully exits with status 1 if no target 277 # process is found, for which run() will throw an 278 # exception. We don't want that, so we the ignore 279 # status. 280 self._host.run("pkill -f '%s'" % remote_name, ignore_status=True) 281 if remote_pid: 282 logging.info('Waiting for RPC server "%s" shutdown', 283 remote_name) 284 start_time = time.time() 285 while (time.time() - start_time < 286 self._RPC_SHUTDOWN_TIMEOUT_SECONDS): 287 running_processes = self._host.run( 288 "pgrep -f '%s'" % remote_name, 289 ignore_status=True).stdout.split() 290 if not remote_pid in running_processes: 291 logging.info('Shut down RPC server.') 292 break 293 time.sleep(self._RPC_SHUTDOWN_POLLING_PERIOD_SECONDS) 294 else: 295 raise error.TestError('Failed to shutdown RPC server %s' % 296 remote_name) 297 298 self._host.disconnect_ssh_tunnel(tunnel_proc) 299 del self._rpc_proxy_map[port] 300 301 302 def disconnect_all(self): 303 """Disconnect all known RPC proxy ports.""" 304 for port in self._rpc_proxy_map.keys(): 305 self.disconnect(port) 306 307 308class TimeoutXMLRPCServerProxy(xmlrpclib.ServerProxy): 309 """XMLRPC ServerProxy supporting timeout.""" 310 def __init__(self, uri, timeout=20, *args, **kwargs): 311 """Initializes a TimeoutXMLRPCServerProxy. 312 313 @param uri: URI to a XMLRPC server. 314 @param timeout: Timeout in seconds for a XMLRPC request. 315 @param *args: args to xmlrpclib.ServerProxy. 316 @param **kwargs: kwargs to xmlrpclib.ServerProxy. 317 318 """ 319 if timeout: 320 kwargs['transport'] = TimeoutXMLRPCTransport(timeout=timeout) 321 xmlrpclib.ServerProxy.__init__(self, uri, *args, **kwargs) 322 323 324class TimeoutXMLRPCTransport(xmlrpclib.Transport): 325 """A Transport subclass supporting timeout.""" 326 def __init__(self, timeout=20, *args, **kwargs): 327 """Initializes a TimeoutXMLRPCTransport. 328 329 @param timeout: Timeout in seconds for a HTTP request through this transport layer. 330 @param *args: args to xmlrpclib.Transport. 331 @param **kwargs: kwargs to xmlrpclib.Transport. 332 333 """ 334 xmlrpclib.Transport.__init__(self, *args, **kwargs) 335 self.timeout = timeout 336 337 338 def make_connection(self, host): 339 """Overwrites make_connection in xmlrpclib.Transport with timeout. 340 341 @param host: Host address to connect. 342 343 @return: A httplib.HTTPConnection connecting to host with timeout. 344 345 """ 346 conn = httplib.HTTPConnection(host, timeout=self.timeout) 347 return conn 348