| OLD | NEW |
| (Empty) |
| 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() | |
| OLD | NEW |