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 SocketServer
  21
  22 #pylint: disable=relative-import
  23 import common_lib
  24 import jsonrpclib
  25 import rpc_methods
  26 import SimpleJSONRPCServer
  27
  28
  29 class RequestHandler(SimpleJSONRPCServer.SimpleJSONRPCRequestHandler):
  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 SimpleJSONRPCServer.SimpleJSONRPCRequestHandler.do_POST(self)
  48
  49
  50 class RpcServer(SimpleJSONRPCServer.SimpleJSONRPCServer,
  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     SimpleJSONRPCServer.SimpleJSONRPCServer.__init__(
  57         self, (common_lib.SERVER_ADDRESS, common_lib.SERVER_PORT),
  58         allow_none=True, logRequests=False,
  59         requestHandler=RequestHandler)
  60     self.authorized_address = authorized_address
  61     self.idle_timeout_secs = idle_timeout_secs
  62     self.register_instance(rpc_methods.RPCMethods(self))
  63
  64     self._shutdown_requested_event = threading.Event()
  65     self._rpc_received_event = threading.Event()
  66     self._idle_thread = threading.Thread(target=self._CheckForIdleQuit)
  67
  68   def shutdown(self):
  69     """Shutdown the server.
  70
  71     This overloaded method sets the _shutdown_requested_event to allow the
  72     idle timeout thread to quit.
  73     """
  74     self._shutdown_requested_event.set()
  75     SimpleJSONRPCServer.SimpleJSONRPCServer.shutdown(self)
  76     logging.info('Server shutdown complete')
  77
  78   def serve_forever(self, poll_interval=0.5):
  79     """Serve forever.
  80
  81     This overloaded method starts the idle timeout thread before calling
  82     serve_forever. This ensures the idle timer thread doesn't get started
  83     without the server running.
  84
  85     Args:
  86       poll_interval: The interval to poll for shutdown.
  87     """
  88     logging.info('RPC server starting')
  89     self._idle_thread.start()
  90     SimpleJSONRPCServer.SimpleJSONRPCServer.serve_forever(self, poll_interval)
  91
  92   def _dispatch(self, method, params):
  93     """Dispatch the call to the correct method with the provided params.
  94
  95     This overloaded method adds logging to help trace connection and
  96     call problems.
  97
  98     Args:
  99       method: The method name to call.
 100       params: A tuple of parameters to pass.
 101
 102     Returns:
 103       The result of the parent class' _dispatch method.
 104     """
 105     logging.debug('Calling %s%s', method, params)
 106     self._rpc_received_event.set()
 107     return SimpleJSONRPCServer.SimpleJSONRPCServer._dispatch(
 108         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()
 129
 130   @staticmethod
 131   def Connect(server, port=common_lib.SERVER_PORT):
 132     """Creates and returns a connection to an RPC server."""
 133     addr = 'http://%s:%d' % (server, port)
 134     logging.debug('Connecting to RPC server at %s', addr)
 135     return jsonrpclib.ServerProxy(addr, allow_none=True)