| # Copyright (c) 2015 The Chromium OS Authors. All rights reserved. | |
| # Use of this source code is governed by a BSD-style license that can be | |
| # found in the LICENSE file. | |
| import httplib | |
| import logging | |
| import socket | |
| import tempfile | |
| import time | |
| import xmlrpclib | |
| import common | |
| from autotest_lib.client.bin import utils | |
| from autotest_lib.client.common_lib import error | |
| from autotest_lib.client.common_lib.cros import retry | |
| try: | |
| import jsonrpclib | |
| except ImportError: | |
| jsonrpclib = None | |
| class RpcServerTracker(object): | |
| """ | |
| This class keeps track of all the RPC server connections started on a remote | |
| host. The caller can use either |xmlrpc_connect| or |jsonrpc_connect| to | |
| start the required type of rpc server on the remote host. | |
| The host will cleanup all the open RPC server connections on disconnect. | |
| """ | |
| _RPC_PROXY_URL_FORMAT = 'http://localhost:%d' | |
| _RPC_HOST_ADDRESS_FORMAT = 'localhost:%d' | |
| _RPC_SHUTDOWN_POLLING_PERIOD_SECONDS = 2 | |
| _RPC_SHUTDOWN_TIMEOUT_SECONDS = 10 | |
| def __init__(self, host): | |
| """ | |
| @param port: The host object associated with this instance of | |
| RpcServerTracker. | |
| """ | |
| self._host = host | |
| self._rpc_proxy_map = {} | |
| def _setup_port(self, port, command_name, remote_pid=None): | |
| """Sets up a tunnel process and register it to rpc_server_tracker. | |
| Chrome OS on the target closes down most external ports for security. | |
| We could open the port, but doing that would conflict with security | |
| tests that check that only expected ports are open. So, to get to | |
| the port on the target we use an ssh tunnel. | |
| This method assumes that xmlrpc and jsonrpc never conflict, since | |
| we can only either have an xmlrpc or a jsonrpc server listening on | |
| a remote port. As such, it enforces a single proxy->remote port | |
| policy, i.e if one starts a jsonrpc proxy/server from port A->B, | |
| and then tries to start an xmlrpc proxy forwarded to the same port, | |
| the xmlrpc proxy will override the jsonrpc tunnel process, however: | |
| 1. None of the methods on the xmlrpc proxy will work because | |
| the server listening on B is jsonrpc. | |
| 2. The xmlrpc client cannot initiate a termination of the JsonRPC | |
| server, as the only use case currently is goofy, which is tied to | |
| the factory image. It is much easier to handle a failed xmlrpc | |
| call on the client than it is to terminate goofy in this scenario, | |
| as doing the latter might leave the DUT in a hard to recover state. | |
| With the current implementation newer rpc proxy connections will | |
| terminate the tunnel processes of older rpc connections tunneling | |
| to the same remote port. If methods are invoked on the client | |
| after this has happened they will fail with connection closed errors. | |
| @param port: The remote forwarding port. | |
| @param command_name: The name of the remote process, to terminate | |
| using pkill. | |
| @param remote_pid: The PID of the remote background process | |
| as a string. | |
| @return the local port which is used for port forwarding on the ssh | |
| client. | |
| """ | |
| self.disconnect(port) | |
| local_port = utils.get_unused_port() | |
| tunnel_proc = self._host.create_ssh_tunnel(port, local_port) | |
| self._rpc_proxy_map[port] = (command_name, tunnel_proc, remote_pid) | |
| return local_port | |
| def _setup_rpc(self, port, command_name, remote_pid=None): | |
| """Construct a URL for an rpc connection using ssh tunnel. | |
| @param port: The remote forwarding port. | |
| @param command_name: The name of the remote process, to terminate | |
| using pkill. | |
| @param remote_pid: The PID of the remote background process | |
| as a string. | |
| @return a url that we can use to initiate the rpc connection. | |
| """ | |
| return self._RPC_PROXY_URL_FORMAT % self._setup_port( | |
| port, command_name, remote_pid=remote_pid) | |
| def tunnel_connect(self, port): | |
| """Construct a host address using ssh tunnel. | |
| @param port: The remote forwarding port. | |
| @return a host address using ssh tunnel. | |
| """ | |
| return self._RPC_HOST_ADDRESS_FORMAT % self._setup_port(port, None) | |
| def xmlrpc_connect(self, command, port, command_name=None, | |
| ready_test_name=None, timeout_seconds=10, | |
| logfile=None, request_timeout_seconds=None): | |
| """Connect to an XMLRPC server on the host. | |
| The `command` argument should be a simple shell command that | |
| starts an XMLRPC server on the given `port`. The command | |
| must not daemonize, and must terminate cleanly on SIGTERM. | |
| The command is started in the background on the host, and a | |
| local XMLRPC client for the server is created and returned | |
| to the caller. | |
| Note that the process of creating an XMLRPC client makes no | |
| attempt to connect to the remote server; the caller is | |
| responsible for determining whether the server is running | |
| correctly, and is ready to serve requests. | |
| Optionally, the caller can pass ready_test_name, a string | |
| containing the name of a method to call on the proxy. This | |
| method should take no parameters and return successfully only | |
| when the server is ready to process client requests. When | |
| ready_test_name is set, xmlrpc_connect will block until the | |
| proxy is ready, and throw a TestError if the server isn't | |
| ready by timeout_seconds. | |
| If a server is already running on the remote port, this | |
| method will kill it and disconnect the tunnel process | |
| associated with the connection before establishing a new one, | |
| by consulting the rpc_proxy_map in disconnect. | |
| @param command Shell command to start the server. | |
| @param port Port number on which the server is expected to | |
| be serving. | |
| @param command_name String to use as input to `pkill` to | |
| terminate the XMLRPC server on the host. | |
| @param ready_test_name String containing the name of a | |
| method defined on the XMLRPC server. | |
| @param timeout_seconds Number of seconds to wait | |
| for the server to become 'ready.' Will throw a | |
| TestFail error if server is not ready in time. | |
| @param logfile Logfile to send output when running | |
| 'command' argument. | |
| @param request_timeout_seconds Timeout in seconds for an XMLRPC request. | |
| """ | |
| # Clean up any existing state. If the caller is willing | |
| # to believe their server is down, we ought to clean up | |
| # any tunnels we might have sitting around. | |
| self.disconnect(port) | |
| remote_pid = None | |
| if command is not None: | |
| if logfile: | |
| remote_cmd = '%s > %s 2>&1' % (command, logfile) | |
| else: | |
| remote_cmd = command | |
| remote_pid = self._host.run_background(remote_cmd) | |
| logging.debug('Started XMLRPC server on host %s, pid = %s', | |
| self._host.hostname, remote_pid) | |
| # Tunnel through SSH to be able to reach that remote port. | |
| rpc_url = self._setup_rpc(port, command_name, remote_pid=remote_pid) | |
| if request_timeout_seconds is not None: | |
| proxy = TimeoutXMLRPCServerProxy( | |
| rpc_url, timeout=request_timeout_seconds, allow_none=True) | |
| else: | |
| proxy = xmlrpclib.ServerProxy(rpc_url, allow_none=True) | |
| if ready_test_name is not None: | |
| # retry.retry logs each attempt; calculate delay_sec to | |
| # keep log spam to a dull roar. | |
| @retry.retry((socket.error, | |
| xmlrpclib.ProtocolError, | |
| httplib.BadStatusLine), | |
| timeout_min=timeout_seconds / 60.0, | |
| delay_sec=min(max(timeout_seconds / 20.0, 0.1), 1)) | |
| def ready_test(): | |
| """ Call proxy.ready_test_name(). """ | |
| getattr(proxy, ready_test_name)() | |
| successful = False | |
| try: | |
| logging.info('Waiting %d seconds for XMLRPC server ' | |
| 'to start.', timeout_seconds) | |
| ready_test() | |
| successful = True | |
| finally: | |
| if not successful: | |
| logging.error('Failed to start XMLRPC server.') | |
| if logfile: | |
| with tempfile.NamedTemporaryFile() as temp: | |
| self._host.get_file(logfile, temp.name) | |
| logging.error('The log of XML RPC server:\n%s', | |
| open(temp.name).read()) | |
| self.disconnect(port) | |
| logging.info('XMLRPC server started successfully.') | |
| return proxy | |
| def jsonrpc_connect(self, port): | |
| """Creates a jsonrpc proxy connection through an ssh tunnel. | |
| This method exists to facilitate communication with goofy (which is | |
| the default system manager on all factory images) and as such, leaves | |
| most of the rpc server sanity checking to the caller. Unlike | |
| xmlrpc_connect, this method does not facilitate the creation of a remote | |
| jsonrpc server, as the only clients of this code are factory tests, | |
| for which the goofy system manager is built in to the image and starts | |
| when the target boots. | |
| One can theoretically create multiple jsonrpc proxies all forwarded | |
| to the same remote port, provided the remote port has an rpc server | |
| listening. However, in doing so we stand the risk of leaking an | |
| existing tunnel process, so we always disconnect any older tunnels | |
| we might have through disconnect. | |
| @param port: port on the remote host that is serving this proxy. | |
| @return: The client proxy. | |
| """ | |
| if not jsonrpclib: | |
| logging.warning('Jsonrpclib could not be imported. Check that ' | |
| 'site-packages contains jsonrpclib.') | |
| return None | |
| proxy = jsonrpclib.jsonrpc.ServerProxy(self._setup_rpc(port, None)) | |
| logging.info('Established a jsonrpc connection through port %s.', port) | |
| return proxy | |
| def disconnect(self, port): | |
| """Disconnect from an RPC server on the host. | |
| Terminates the remote RPC server previously started for | |
| the given `port`. Also closes the local ssh tunnel created | |
| for the connection to the host. This function does not | |
| directly alter the state of a previously returned RPC | |
| client object; however disconnection will cause all | |
| subsequent calls to methods on the object to fail. | |
| This function does nothing if requested to disconnect a port | |
| that was not previously connected via _setup_rpc. | |
| @param port Port number passed to a previous call to | |
| `_setup_rpc()`. | |
| """ | |
| if port not in self._rpc_proxy_map: | |
| return | |
| remote_name, tunnel_proc, remote_pid = self._rpc_proxy_map[port] | |
| if remote_name: | |
| # We use 'pkill' to find our target process rather than | |
| # a PID, because the host may have rebooted since | |
| # connecting, and we don't want to kill an innocent | |
| # process with the same PID. | |
| # | |
| # 'pkill' helpfully exits with status 1 if no target | |
| # process is found, for which run() will throw an | |
| # exception. We don't want that, so we the ignore | |
| # status. | |
| self._host.run("pkill -f '%s'" % remote_name, ignore_status=True) | |
| if remote_pid: | |
| logging.info('Waiting for RPC server "%s" shutdown', | |
| remote_name) | |
| start_time = time.time() | |
| while (time.time() - start_time < | |
| self._RPC_SHUTDOWN_TIMEOUT_SECONDS): | |
| running_processes = self._host.run( | |
| "pgrep -f '%s'" % remote_name, | |
| ignore_status=True).stdout.split() | |
| if not remote_pid in running_processes: | |
| logging.info('Shut down RPC server.') | |
| break | |
| time.sleep(self._RPC_SHUTDOWN_POLLING_PERIOD_SECONDS) | |
| else: | |
| raise error.TestError('Failed to shutdown RPC server %s' % | |
| remote_name) | |
| self._host.disconnect_ssh_tunnel(tunnel_proc, port) | |
| del self._rpc_proxy_map[port] | |
| def disconnect_all(self): | |
| """Disconnect all known RPC proxy ports.""" | |
| for port in self._rpc_proxy_map.keys(): | |
| self.disconnect(port) | |
| class TimeoutXMLRPCServerProxy(xmlrpclib.ServerProxy): | |
| """XMLRPC ServerProxy supporting timeout.""" | |
| def __init__(self, uri, timeout=20, *args, **kwargs): | |
| """Initializes a TimeoutXMLRPCServerProxy. | |
| @param uri: URI to a XMLRPC server. | |
| @param timeout: Timeout in seconds for a XMLRPC request. | |
| @param *args: args to xmlrpclib.ServerProxy. | |
| @param **kwargs: kwargs to xmlrpclib.ServerProxy. | |
| """ | |
| if timeout: | |
| kwargs['transport'] = TimeoutXMLRPCTransport(timeout=timeout) | |
| xmlrpclib.ServerProxy.__init__(self, uri, *args, **kwargs) | |
| class TimeoutXMLRPCTransport(xmlrpclib.Transport): | |
| """A Transport subclass supporting timeout.""" | |
| def __init__(self, timeout=20, *args, **kwargs): | |
| """Initializes a TimeoutXMLRPCTransport. | |
| @param timeout: Timeout in seconds for a HTTP request through this transport layer. | |
| @param *args: args to xmlrpclib.Transport. | |
| @param **kwargs: kwargs to xmlrpclib.Transport. | |
| """ | |
| xmlrpclib.Transport.__init__(self, *args, **kwargs) | |
| self.timeout = timeout | |
| def make_connection(self, host): | |
| """Overwrites make_connection in xmlrpclib.Transport with timeout. | |
| @param host: Host address to connect. | |
| @return: A httplib.HTTPConnection connecting to host with timeout. | |
| """ | |
| conn = httplib.HTTPConnection(host, timeout=self.timeout) | |
| return conn |