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 time
9import xmlrpclib
10
11import common
12from autotest_lib.client.bin import utils
13from autotest_lib.client.common_lib import error
14from autotest_lib.client.common_lib.cros import retry
15
16try:
17    import jsonrpclib
18except ImportError:
19    jsonrpclib = None
20
21
22class 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