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