Adding the initial code for Omnibot multi-machine support

testing/legion/client_rpc_server.py

+# Copyright 2015 The Chromium Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+"""The client RPC server code.
+
+This server is an XML-RPC server which serves code from
+client_rpc_methods.RPCMethods.
+
+This server will run until shutdown is called on the server object. This can
+be achieved in 2 ways:
+
+- Calling the Quit RPC method defined in RPCMethods
+- Not receiving any calls within the RPC_IDLE_TIMEOUT time.
+
+The second state is used in a case where the host controller has crashed and
+cannot send a Quit command.
+"""
+import logging
+import threading
+import time
+import xmlrpclib
+import SimpleXMLRPCServer
+import SocketServer
+
+#pylint: disable=relative-import
+import client_rpc_methods
+
+SERVER_ADDRESS = ''
+SERVER_PORT = 31710
+IDLE_TIMEOUT = 30 * 60  # 30 minutes
+
+
+class RequestHandler(SimpleXMLRPCServer.SimpleXMLRPCRequestHandler):
+  """A custom handler class that provides simple authentication.
+
+  Only requests received from specified addresses are fulfilled. All other
+  requests result in a 403 Forbidden result.
+  """
+
+  def Report403(self):
+    """Report a 403 Forbidden error."""
+    self.send_response(403)
+    response = 'Forbidden'
+    self.send_header("Content-type", "text/plain")
+    self.send_header("Content-length", str(len(response)))
+    self.end_headers()
+    self.wfile.write(response)
+
+  def do_POST(self):
+    """Custom do_POST that verifies the client is authorized to perform RPCs."""
+    if self.client_address[0] not in self.server.authorized_addresses:
+      logging.error('Received unauthorized RPC request from %s',
+                    self.client_address[0])
+      self.Report403()
+      return
+    self.server.RPCReceived()
+    return SimpleXMLRPCServer.SimpleXMLRPCRequestHandler.do_POST(self)
+
+
+class RPCServer(SimpleXMLRPCServer.SimpleXMLRPCServer,
+                SocketServer.ThreadingMixIn):
+  """Custom RPC server that supports authorized client addresses."""
+
+  def __init__(self, address=None, port=None):
+    """ctor.
+
+    Args:
+      address: The address to serve on ('', 'localhost', etc).
+      port: The port to serve on.
+    """
+    address = address or SERVER_ADDRESS
+    port = port or SERVER_PORT
+    SimpleXMLRPCServer.SimpleXMLRPCServer.__init__(
+        self, (address, port), allow_none=True, logRequests=False,
+        requestHandler=RequestHandler)
+    self.authorized_addresses = []
+    self._idle_timeout = IDLE_TIMEOUT
+    self.register_instance(client_rpc_methods.RPCMethods(self))
+    self._StartIdleTimeoutThread()
+
+  def _StartIdleTimeoutThread(self):
+    """Create and start the idle timeout monitor thread."""
+    self.shutdown_requested_event = threading.Event()
+    self.rpc_received_event = threading.Event()
+    self.idle_thread = threading.Thread(target=self.CheckForIdleQuit)
+
+  def AddAuthorizedAddress(self, ip_address):
+    """Add an authorized caller address to the server.
+
+    Args:
+      ip_address: The IP address of the caller.
+    """
+    self.authorized_addresses.append(ip_address)
+
+  def SetIdleTimeout(self, timeout):
+    """Sets the idle timeout for the server.
+
+    Args:
+      timeout: The server timeout in seconds.
+    """
+    logging.debug('Idle timeout set to %s', timeout)
+    self._idle_timeout = timeout
+
+  @staticmethod
+  def Connect(client, port=None):
+    """Connect to an RPC server.
+
+    Args:
+      client: The client address which the RPC server is running on.
+      port: The port the client RPC server is running on.
+
+    Returns:
+      An XML-RPC connection to the client RPC server.
+    """
+    port = port or SERVER_PORT
+    server = 'http://%s:%d' % (client, port)
+    logging.debug('Connecting to RPC server at %s', server)
+    return xmlrpclib.Server(server)
+
+  def shutdown(self):
+    """Shutdown the server.
+
+    This overloaded method sets the shutdown_requested_event to allow the
+    idle timeout thread to quit.
+    """
+    self.shutdown_requested_event.set()
+    SimpleXMLRPCServer.SimpleXMLRPCServer.shutdown(self)
+    logging.info('Server shutdown complete')
+
+  def serve_forever(self, poll_interval=0.5):
+    """Serve forever.
+
+    This overloaded method starts the idle timeout thread before calling
+    serve_forever. This ensures the idle timer thread doesn't get started
+    without the server running.
+
+    Args:
+      poll_interval: The interval to poll for shutdown.
+    """
+    logging.info('RPC server starting')
+    self.idle_thread.start()
+    SimpleXMLRPCServer.SimpleXMLRPCServer.serve_forever(self, poll_interval)
+
+  def _dispatch(self, method, params):
+    """Dispatch the call to the correct method with the provided params.
+
+    This overloaded method adds logging to help trace connection and
+    call problems.
+
+    Args:
+      method: The method name to call.
+      params: A tuple of parameters to pass.
+
+    Returns:
+      The result of the parent class' _dispatch method.
+    """
+    logging.debug('Calling %s%s', method, params)
+    return SimpleXMLRPCServer.SimpleXMLRPCServer._dispatch(self, method, params)
+
+  def RPCReceived(self):
+    """Sets the rpc_received_event whenever an RPC is received.
+
+    Setting this event effectively resets the idle timeout thread.
+    """
+    self.rpc_received_event.set()
+
+  def CheckForIdleQuit(self):
+    """Check for, and exit, if the server is idle for too long.
+
+    This method must be run in a separate thread to avoid a deadlock when
+    calling server.shutdown.
+    """
+    timeout = time.time() + self._idle_timeout
+    while time.time() < timeout:
+      if self.shutdown_requested_event.is_set():
+        # An external source called shutdown()
+        return
+      elif self.rpc_received_event.is_set():
+        # An RPC was received, reset the timeout
+        logging.debug('Resetting the idle timeout')
+        timeout = time.time() + self._idle_timeout
+        self.rpc_received_event.clear()
+      time.sleep(1)
+    # We timed out, kill the server
+    logging.warning('Shutting down the server due to the idle timeout')
+    self.shutdown()