testing/legion/rpc_server.py

   1 # Copyright 2015 The Chromium 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 """The task RPC server code.
   6
   7 This server is an XML-RPC server which serves code from
   8 rpc_methods.RPCMethods.
   9
  10 This server will run until shutdown is called on the server object. This can
  11 be achieved in 2 ways:
  12
  13 - Calling the Quit RPC method defined in RPCMethods
  14 - Not receiving any calls within the idle_timeout_secs time.
  15 """
  16
  17 import logging
  18 import threading
  19 import time
  20 import xmlrpclib
  21 import SimpleXMLRPCServer
  22 import SocketServer
  23
  24 #pylint: disable=relative-import
  25 import common_lib
  26 import rpc_methods
  27
  28
  29 class RequestHandler(SimpleXMLRPCServer.SimpleXMLRPCRequestHandler):
  30   """Restricts access to only specified IP address.
  31
  32   This call assumes the server is RPCServer.
  33   """
  34
  35   def do_POST(self):
  36     """Verifies the task is authorized to perform RPCs."""
  37     if self.client_address[0] != self.server.authorized_address:
  38       logging.error('Received unauthorized RPC request from %s',
  39                     self.task_address[0])
  40       self.send_response(403)
  41       response = 'Forbidden'
  42       self.send_header('Content-type', 'text/plain')
  43       self.send_header('Content-length', str(len(response)))
  44       self.end_headers()
  45       self.wfile.write(response)
  46     else:
  47       return SimpleXMLRPCServer.SimpleXMLRPCRequestHandler.do_POST(self)
  48
  49
  50 class RPCServer(SimpleXMLRPCServer.SimpleXMLRPCServer,
  51                 SocketServer.ThreadingMixIn):
  52   """Restricts all endpoints to only specified IP addresses."""
  53
  54   def __init__(self, authorized_address,
  55                idle_timeout_secs=common_lib.DEFAULT_TIMEOUT_SECS):
  56     SimpleXMLRPCServer.SimpleXMLRPCServer.__init__(
  57         self, (common_lib.SERVER_ADDRESS, common_lib.SERVER_PORT),
  58         allow_none=True, logRequests=False,
  59         requestHandler=RequestHandler)
  60
  61     self.authorized_address = authorized_address
  62     self.idle_timeout_secs = idle_timeout_secs
  63     self.register_instance(rpc_methods.RPCMethods(self))
  64
  65     self._shutdown_requested_event = threading.Event()
  66     self._rpc_received_event = threading.Event()
  67     self._idle_thread = threading.Thread(target=self._CheckForIdleQuit)
  68
  69   def shutdown(self):
  70     """Shutdown the server.
  71
  72     This overloaded method sets the _shutdown_requested_event to allow the
  73     idle timeout thread to quit.
  74     """
  75     self._shutdown_requested_event.set()
  76     SimpleXMLRPCServer.SimpleXMLRPCServer.shutdown(self)
  77     logging.info('Server shutdown complete')
  78
  79   def serve_forever(self, poll_interval=0.5):
  80     """Serve forever.
  81
  82     This overloaded method starts the idle timeout thread before calling
  83     serve_forever. This ensures the idle timer thread doesn't get started
  84     without the server running.
  85
  86     Args:
  87       poll_interval: The interval to poll for shutdown.
  88     """
  89     logging.info('RPC server starting')
  90     self._idle_thread.start()
  91     SimpleXMLRPCServer.SimpleXMLRPCServer.serve_forever(self, poll_interval)
  92
  93   def _dispatch(self, method, params):
  94     """Dispatch the call to the correct method with the provided params.
  95
  96     This overloaded method adds logging to help trace connection and
  97     call problems.
  98
  99     Args:
 100       method: The method name to call.
 101       params: A tuple of parameters to pass.
 102
 103     Returns:
 104       The result of the parent class' _dispatch method.
 105     """
 106     logging.debug('Calling %s%s', method, params)
 107     self._rpc_received_event.set()
 108     return SimpleXMLRPCServer.SimpleXMLRPCServer._dispatch(self, method, params)
 109
 110   def _CheckForIdleQuit(self):
 111     """Check for, and exit, if the server is idle for too long.
 112
 113     This method must be run in a separate thread to avoid a deadlock when
 114     calling server.shutdown.
 115     """
 116     timeout = time.time() + self.idle_timeout_secs
 117     while time.time() < timeout:
 118       if self._shutdown_requested_event.is_set():
 119         # An external source called shutdown()
 120         return
 121       elif self._rpc_received_event.is_set():
 122         logging.debug('Resetting the idle timeout')
 123         timeout = time.time() + self.idle_timeout_secs
 124         self._rpc_received_event.clear()
 125       time.sleep(1)
 126     # We timed out, kill the server
 127     logging.warning('Shutting down the server due to the idle timeout')
 128     self.shutdown()