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 getattr(proxy, ready_test_name)() 193 successful = False 194 try: 195 logging.info('Waiting %d seconds for XMLRPC server ' 196 'to start.', timeout_seconds) 197 ready_test() 198 successful = True 199 finally: 200 if not successful: 201 logging.error('Failed to start XMLRPC server.') 202 if logfile: 203 with tempfile.NamedTemporaryFile() as temp: 204 self._host.get_file(logfile, temp.name) 205 logging.error('The log of XML RPC server:\n%s', 206 open(temp.name).read()) 207 self.disconnect(port) 208 logging.info('XMLRPC server started successfully.') 209 return proxy 210 211 212 def jsonrpc_connect(self, port): 213 """Creates a jsonrpc proxy connection through an ssh tunnel. 214 215 This method exists to facilitate communication with goofy (which is 216 the default system manager on all factory images) and as such, leaves 217 most of the rpc server sanity checking to the caller. Unlike 218 xmlrpc_connect, this method does not facilitate the creation of a remote 219 jsonrpc server, as the only clients of this code are factory tests, 220 for which the goofy system manager is built in to the image and starts 221 when the target boots. 222 223 One can theoretically create multiple jsonrpc proxies all forwarded 224 to the same remote port, provided the remote port has an rpc server 225 listening. However, in doing so we stand the risk of leaking an 226 existing tunnel process, so we always disconnect any older tunnels 227 we might have through disconnect. 228 229 @param port: port on the remote host that is serving this proxy. 230 231 @return: The client proxy. 232 """ 233 if not jsonrpclib: 234 logging.warning('Jsonrpclib could not be imported. Check that ' 235 'site-packages contains jsonrpclib.') 236 return None 237 238 proxy = jsonrpclib.jsonrpc.ServerProxy(self._setup_rpc(port, None)) 239 240 logging.info('Established a jsonrpc connection through port %s.', port) 241 return proxy 242 243 244 def disconnect(self, port): 245 """Disconnect from an RPC server on the host. 246 247 Terminates the remote RPC server previously started for 248 the given `port`. Also closes the local ssh tunnel created 249 for the connection to the host. This function does not 250 directly alter the state of a previously returned RPC 251 client object; however disconnection will cause all 252 subsequent calls to methods on the object to fail. 253 254 This function does nothing if requested to disconnect a port 255 that was not previously connected via _setup_rpc. 256 257 @param port Port number passed to a previous call to 258 `_setup_rpc()`. 259 """ 260 if port not in self._rpc_proxy_map: 261 return 262 remote_name, tunnel_proc, remote_pid = self._rpc_proxy_map[port] 263 if remote_name: 264 # We use 'pkill' to find our target process rather than 265 # a PID, because the host may have rebooted since 266 # connecting, and we don't want to kill an innocent 267 # process with the same PID. 268 # 269 # 'pkill' helpfully exits with status 1 if no target 270 # process is found, for which run() will throw an 271 # exception. We don't want that, so we the ignore 272 # status. 273 self._host.run("pkill -f '%s'" % remote_name, ignore_status=True) 274 if remote_pid: 275 logging.info('Waiting for RPC server "%s" shutdown', 276 remote_name) 277 start_time = time.time() 278 while (time.time() - start_time < 279 self._RPC_SHUTDOWN_TIMEOUT_SECONDS): 280 running_processes = self._host.run( 281 "pgrep -f '%s'" % remote_name, 282 ignore_status=True).stdout.split() 283 if not remote_pid in running_processes: 284 logging.info('Shut down RPC server.') 285 break 286 time.sleep(self._RPC_SHUTDOWN_POLLING_PERIOD_SECONDS) 287 else: 288 raise error.TestError('Failed to shutdown RPC server %s' % 289 remote_name) 290 291 self._host.disconnect_ssh_tunnel(tunnel_proc, port) 292 del self._rpc_proxy_map[port] 293 294 295 def disconnect_all(self): 296 """Disconnect all known RPC proxy ports.""" 297 for port in self._rpc_proxy_map.keys(): 298 self.disconnect(port) 299 300 301class TimeoutXMLRPCServerProxy(xmlrpclib.ServerProxy): 302 """XMLRPC ServerProxy supporting timeout.""" 303 def __init__(self, uri, timeout=20, *args, **kwargs): 304 """Initializes a TimeoutXMLRPCServerProxy. 305 306 @param uri: URI to a XMLRPC server. 307 @param timeout: Timeout in seconds for a XMLRPC request. 308 @param *args: args to xmlrpclib.ServerProxy. 309 @param **kwargs: kwargs to xmlrpclib.ServerProxy. 310 311 """ 312 if timeout: 313 kwargs['transport'] = TimeoutXMLRPCTransport(timeout=timeout) 314 xmlrpclib.ServerProxy.__init__(self, uri, *args, **kwargs) 315 316 317class TimeoutXMLRPCTransport(xmlrpclib.Transport): 318 """A Transport subclass supporting timeout.""" 319 def __init__(self, timeout=20, *args, **kwargs): 320 """Initializes a TimeoutXMLRPCTransport. 321 322 @param timeout: Timeout in seconds for a HTTP request through this transport layer. 323 @param *args: args to xmlrpclib.Transport. 324 @param **kwargs: kwargs to xmlrpclib.Transport. 325 326 """ 327 xmlrpclib.Transport.__init__(self, *args, **kwargs) 328 self.timeout = timeout 329 330 331 def make_connection(self, host): 332 """Overwrites make_connection in xmlrpclib.Transport with timeout. 333 334 @param host: Host address to connect. 335 336 @return: A httplib.HTTPConnection connecting to host with timeout. 337 338 """ 339 conn = httplib.HTTPConnection(host, timeout=self.timeout) 340 return conn 341