Switching the RPC protocol from XML to JSON

This change allows browser_test-based tests to communicate with the RPC servers
more easily (using Javascript JSON libraries).

BUG=509731

Review URL: https://codereview.chromium.org/1236003002

Cr-Commit-Position: refs/heads/master@{#338565}

8 files changed, 581 insertions, 16 deletions

diff --git a/testing/legion/SimpleJSONRPCServer.py b/testing/legion/SimpleJSONRPCServer.py
new file mode 100644
index 0000000..4dbfd6e
--- /dev/null
+++ b/testing/legion/SimpleJSONRPCServer.py
@@ -0,0 +1,195 @@
+# Copyright 2014 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.
+
+"""Module to implement the SimpleXMLRPCServer module using JSON-RPC.
+
+This module uses SimpleXMLRPCServer as the base and only overrides those
+portions that implement the XML-RPC protocol. These portions are rewritten
+to use the JSON-RPC protocol instead.
+
+When large portions of code need to be rewritten the original code and
+comments are preserved. The intention here is to keep the amount of code
+change to a minimum.
+
+This module only depends on default Python modules, as well as jsonrpclib
+which also uses only default modules. No third party code is required to
+use this module.
+"""
+import fcntl
+import json
+import SimpleXMLRPCServer as _base
+import SocketServer
+import sys
+import traceback
+try:
+  import gzip
+except ImportError:
+  gzip = None #python can be built without zlib/gzip support
+
+#pylint: disable=relative-import
+import jsonrpclib
+
+
+class SimpleJSONRPCRequestHandler(_base.SimpleXMLRPCRequestHandler):
+  """Request handler class for received requests.
+
+  This class extends the functionality of SimpleXMLRPCRequestHandler and only
+  overrides the operations needed to change the protocol from XML-RPC to
+  JSON-RPC.
+  """
+
+  def do_POST(self):
+    """Handles the HTTP POST request.
+
+    Attempts to interpret all HTTP POST requests as JSON-RPC calls,
+    which are forwarded to the server's _dispatch method for handling.
+    """
+    # Check that the path is legal
+    if not self.is_rpc_path_valid():
+      self.report_404()
+      return
+
+    try:
+      # Get arguments by reading body of request.
+      # We read this in chunks to avoid straining
+      # socket.read(); around the 10 or 15Mb mark, some platforms
+      # begin to have problems (bug #792570).
+      max_chunk_size = 10*1024*1024
+      size_remaining = int(self.headers['content-length'])
+      data = []
+      while size_remaining:
+        chunk_size = min(size_remaining, max_chunk_size)
+        chunk = self.rfile.read(chunk_size)
+        if not chunk:
+          break
+        data.append(chunk)
+        size_remaining -= len(data[-1])
+      data = ''.join(data)
+      data = self.decode_request_content(data)
+
+      if data is None:
+        return  # response has been sent
+
+      # In previous versions of SimpleXMLRPCServer, _dispatch
+      # could be overridden in this class, instead of in
+      # SimpleXMLRPCDispatcher. To maintain backwards compatibility,
+      # check to see if a subclass implements _dispatch and dispatch
+      # using that method if present.
+      response = self.server._marshaled_dispatch(
+          data, getattr(self, '_dispatch', None), self.path)
+
+    except Exception, e: # This should only happen if the module is buggy
+      # internal error, report as HTTP server error
+      self.send_response(500)
+      # Send information about the exception if requested
+      if (hasattr(self.server, '_send_traceback_header') and
+          self.server._send_traceback_header):
+        self.send_header('X-exception', str(e))
+        self.send_header('X-traceback', traceback.format_exc())
+
+      self.send_header('Content-length', '0')
+      self.end_headers()
+    else:
+      # got a valid JSON RPC response
+      self.send_response(200)
+      self.send_header('Content-type', 'application/json')
+
+      if self.encode_threshold is not None:
+        if len(response) > self.encode_threshold:
+          q = self.accept_encodings().get('gzip', 0)
+          if q:
+            try:
+              response = jsonrpclib.gzip_encode(response)
+              self.send_header('Content-Encoding', 'gzip')
+            except NotImplementedError:
+              pass
+
+      self.send_header('Content-length', str(len(response)))
+      self.end_headers()
+      self.wfile.write(response)
+
+
+class SimpleJSONRPCDispatcher(_base.SimpleXMLRPCDispatcher):
+  """Dispatcher for received JSON-RPC requests.
+
+  This class extends the functionality of SimpleXMLRPCDispatcher and only
+  overrides the operations needed to change the protocol from XML-RPC to
+  JSON-RPC.
+  """
+
+  def _marshaled_dispatch(self, data, dispatch_method=None, path=None):
+    """Dispatches an JSON-RPC method from marshalled (JSON) data.
+
+    JSON-RPC methods are dispatched from the marshalled (JSON) data
+    using the _dispatch method and the result is returned as
+    marshalled data. For backwards compatibility, a dispatch
+    function can be provided as an argument (see comment in
+    SimpleJSONRPCRequestHandler.do_POST) but overriding the
+    existing method through subclassing is the preferred means
+    of changing method dispatch behavior.
+
+    Returns:
+      The JSON-RPC string to return.
+    """
+    method = ''
+    params = []
+    ident = ''
+    try:
+      request = json.loads(data)
+      jsonrpclib.ValidateRequest(request)
+      method = request['method']
+      params = request['params']
+      ident = request['id']
+
+      # generate response
+      if dispatch_method is not None:
+        response = dispatch_method(method, params)
+      else:
+        response = self._dispatch(method, params)
+      response = jsonrpclib.CreateResponseString(response, ident)
+
+    except jsonrpclib.Fault as fault:
+      response = jsonrpclib.CreateResponseString(fault, ident)
+
+    # Catch all exceptions here as they should be raised on the caller side.
+    except:  #pylint: disable=bare-except
+      # report exception back to server
+      exc_type, exc_value, _ = sys.exc_info()
+      response = jsonrpclib.CreateResponseString(
+          jsonrpclib.Fault(1, '%s:%s' % (exc_type, exc_value)), ident)
+    return response
+
+
+class SimpleJSONRPCServer(SocketServer.TCPServer,
+                          SimpleJSONRPCDispatcher):
+  """Simple JSON-RPC server.
+
+  This class mimics the functionality of SimpleXMLRPCServer and only
+  overrides the operations needed to change the protocol from XML-RPC to
+  JSON-RPC.
+  """
+
+  allow_reuse_address = True
+
+  # Warning: this is for debugging purposes only! Never set this to True in
+  # production code, as will be sending out sensitive information (exception
+  # and stack trace details) when exceptions are raised inside
+  # SimpleJSONRPCRequestHandler.do_POST
+  _send_traceback_header = False
+
+  def __init__(self, addr, requestHandler=SimpleJSONRPCRequestHandler,
+               logRequests=True, allow_none=False, encoding=None,
+               bind_and_activate=True):
+    self.logRequests = logRequests
+    SimpleJSONRPCDispatcher.__init__(self, allow_none, encoding)
+    SocketServer.TCPServer.__init__(self, addr, requestHandler,
+                                    bind_and_activate)
+
+    # [Bug #1222790] If possible, set close-on-exec flag; if a
+    # method spawns a subprocess, the subprocess shouldn't have
+    # the listening socket open.
+    if fcntl is not None and hasattr(fcntl, 'FD_CLOEXEC'):
+      flags = fcntl.fcntl(self.fileno(), fcntl.F_GETFD)
+      flags |= fcntl.FD_CLOEXEC
+      fcntl.fcntl(self.fileno(), fcntl.F_SETFD, flags)
diff --git a/testing/legion/common_lib.py b/testing/legion/common_lib.py
index bc3222c..f8169d5 100644
--- a/testing/legion/common_lib.py
+++ b/testing/legion/common_lib.py
@@ -8,7 +8,6 @@ import argparse
 import logging
 import os
 import socket
-import xmlrpclib
 
 LOGGING_LEVELS = ['DEBUG', 'INFO', 'WARNING', 'WARN', 'ERROR']
 MY_IP = socket.gethostbyname(socket.gethostname())
diff --git a/testing/legion/jsonrpclib.py b/testing/legion/jsonrpclib.py
new file mode 100644
index 0000000..56b30e8
--- /dev/null
+++ b/testing/legion/jsonrpclib.py
@@ -0,0 +1,370 @@
+# Copyright 2014 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.
+
+"""Module to implement the JSON-RPC protocol.
+
+This module uses xmlrpclib as the base and only overrides those
+portions that implement the XML-RPC protocol. These portions are rewritten
+to use the JSON-RPC protocol instead.
+
+When large portions of code need to be rewritten the original code and
+comments are preserved. The intention here is to keep the amount of code
+change to a minimum.
+
+This module only depends on default Python modules. No third party code is
+required to use this module.
+"""
+import json
+import urllib
+import xmlrpclib as _base
+
+__version__ = '1.0.0'
+gzip_encode = _base.gzip_encode
+gzip = _base.gzip
+
+
+class Error(Exception):
+
+  def __str__(self):
+    return repr(self)
+
+
+class ProtocolError(Error):
+  """Indicates a JSON protocol error."""
+
+  def __init__(self, url, errcode, errmsg, headers):
+    Error.__init__(self)
+    self.url = url
+    self.errcode = errcode
+    self.errmsg = errmsg
+    self.headers = headers
+
+  def __repr__(self):
+    return (
+        '<ProtocolError for %s: %s %s>' %
+        (self.url, self.errcode, self.errmsg))
+
+
+class ResponseError(Error):
+  """Indicates a broken response package."""
+  pass
+
+
+class Fault(Error):
+  """Indicates a JSON-RPC fault package."""
+
+  def __init__(self, code, message):
+    Error.__init__(self)
+    if not isinstance(code, int):
+      raise ProtocolError('Fault code must be an integer.')
+    self.code = code
+    self.message = message
+
+  def __repr__(self):
+      return (
+          '<Fault %s: %s>' %
+          (self.code, repr(self.message))
+          )
+
+
+def CreateRequest(methodname, params, ident=''):
+  """Create a valid JSON-RPC request.
+
+  Args:
+    methodname: The name of the remote method to invoke.
+    params: The parameters to pass to the remote method. This should be a
+        list or tuple and able to be encoded by the default JSON parser.
+
+  Returns:
+    A valid JSON-RPC request object.
+  """
+  request = {
+      'jsonrpc': '2.0',
+      'method': methodname,
+      'params': params,
+      'id': ident
+      }
+
+  return request
+
+
+def CreateRequestString(methodname, params, ident=''):
+  """Create a valid JSON-RPC request string.
+
+  Args:
+    methodname: The name of the remote method to invoke.
+    params: The parameters to pass to the remote method.
+        These parameters need to be encode-able by the default JSON parser.
+    ident: The request identifier.
+
+  Returns:
+    A valid JSON-RPC request string.
+  """
+  return json.dumps(CreateRequest(methodname, params, ident))
+
+
+def CreateResponse(data, ident):
+  """Create a JSON-RPC response.
+
+  Args:
+    data: The data to return.
+    ident: The response identifier.
+
+  Returns:
+    A valid JSON-RPC response object.
+  """
+  if isinstance(data, Fault):
+    response = {
+        'jsonrpc': '2.0',
+        'error': {
+            'code': data.code,
+            'message': data.message},
+        'id': ident
+        }
+  else:
+    response = {
+        'jsonrpc': '2.0',
+        'response': data,
+        'id': ident
+        }
+
+  return response
+
+
+def CreateResponseString(data, ident):
+  """Create a JSON-RPC response string.
+
+  Args:
+    data: The data to return.
+    ident: The response identifier.
+
+  Returns:
+    A valid JSON-RPC response object.
+  """
+  return json.dumps(CreateResponse(data, ident))
+
+
+def ParseHTTPResponse(response):
+  """Parse an HTTP response object and return the JSON object.
+
+  Args:
+    response: An HTTP response object.
+
+  Returns:
+    The returned JSON-RPC object.
+
+  Raises:
+    ProtocolError: if the object format is not correct.
+    Fault: If a Fault error is returned from the server.
+  """
+  # Check for new http response object, else it is a file object
+  if hasattr(response, 'getheader'):
+    if response.getheader('Content-Encoding', '') == 'gzip':
+      stream = _base.GzipDecodedResponse(response)
+    else:
+      stream = response
+  else:
+    stream = response
+
+  data = ''
+  while 1:
+    chunk = stream.read(1024)
+    if not chunk:
+      break
+    data += chunk
+
+  response = json.loads(data)
+  ValidateBasicJSONRPCData(response)
+
+  if 'response' in response:
+    ValidateResponse(response)
+    return response['response']
+  elif 'error' in response:
+    ValidateError(response)
+    code = response['error']['code']
+    message = response['error']['message']
+    raise Fault(code, message)
+  else:
+    raise ProtocolError('No valid JSON returned')
+
+
+def ValidateRequest(data):
+  """Validate a JSON-RPC request object.
+
+  Args:
+    data: The JSON-RPC object (dict).
+
+  Raises:
+    ProtocolError: if the object format is not correct.
+  """
+  ValidateBasicJSONRPCData(data)
+  if 'method' not in data or 'params' not in data:
+    raise ProtocolError('JSON is not a valid request')
+
+
+def ValidateResponse(data):
+  """Validate a JSON-RPC response object.
+
+  Args:
+    data: The JSON-RPC object (dict).
+
+  Raises:
+    ProtocolError: if the object format is not correct.
+  """
+  ValidateBasicJSONRPCData(data)
+  if 'response' not in data:
+    raise ProtocolError('JSON is not a valid response')
+
+
+def ValidateError(data):
+  """Validate a JSON-RPC error object.
+
+  Args:
+    data: The JSON-RPC object (dict).
+
+  Raises:
+    ProtocolError: if the object format is not correct.
+  """
+  ValidateBasicJSONRPCData(data)
+  if ('error' not in data or
+      'code' not in data['error'] or
+      'message' not in data['error']):
+    raise ProtocolError('JSON is not a valid error response')
+
+
+def ValidateBasicJSONRPCData(data):
+  """Validate a basic JSON-RPC object.
+
+  Args:
+    data: The JSON-RPC object (dict).
+
+  Raises:
+    ProtocolError: if the object format is not correct.
+  """
+  error = None
+  if not isinstance(data, dict):
+    error = 'JSON data is not a dictionary'
+  elif 'jsonrpc' not in data or data['jsonrpc'] != '2.0':
+    error = 'JSON is not a valid JSON RPC 2.0 message'
+  elif 'id' not in data:
+    error = 'JSON data missing required id entry'
+  if error:
+    raise ProtocolError(error)
+
+
+class Transport(_base.Transport):
+  """RPC transport class.
+
+  This class extends the functionality of xmlrpclib.Transport and only
+  overrides the operations needed to change the protocol from XML-RPC to
+  JSON-RPC.
+  """
+
+  user_agent = 'jsonrpclib.py/' + __version__
+
+  def send_content(self, connection, request_body):
+    """Send the request."""
+    connection.putheader('Content-Type','application/json')
+
+    #optionally encode the request
+    if (self.encode_threshold is not None and
+        self.encode_threshold < len(request_body) and
+        gzip):
+      connection.putheader('Content-Encoding', 'gzip')
+      request_body = gzip_encode(request_body)
+
+    connection.putheader('Content-Length', str(len(request_body)))
+    connection.endheaders(request_body)
+
+  def single_request(self, host, handler, request_body, verbose=0):
+    """Issue a single JSON-RPC request."""
+
+    h = self.make_connection(host)
+    if verbose:
+      h.set_debuglevel(1)
+    try:
+      self.send_request(h, handler, request_body)
+      self.send_host(h, host)
+      self.send_user_agent(h)
+      self.send_content(h, request_body)
+
+      response = h.getresponse(buffering=True)
+      if response.status == 200:
+        self.verbose = verbose  #pylint: disable=attribute-defined-outside-init
+
+        return self.parse_response(response)
+
+    except Fault:
+      raise
+    except Exception:
+      # All unexpected errors leave connection in
+      # a strange state, so we clear it.
+      self.close()
+      raise
+
+    # discard any response data and raise exception
+    if response.getheader('content-length', 0):
+      response.read()
+    raise ProtocolError(
+        host + handler,
+        response.status, response.reason,
+        response.msg,
+        )
+
+  def parse_response(self, response):
+    """Parse the HTTP resoponse from the server."""
+    return ParseHTTPResponse(response)
+
+
+class SafeTransport(_base.SafeTransport):
+  """Transport class for HTTPS servers.
+
+  This class extends the functionality of xmlrpclib.SafeTransport and only
+  overrides the operations needed to change the protocol from XML-RPC to
+  JSON-RPC.
+  """
+
+  def parse_response(self, response):
+    return ParseHTTPResponse(response)
+
+
+class ServerProxy(_base.ServerProxy):
+  """Proxy class to the RPC server.
+
+  This class extends the functionality of xmlrpclib.ServerProxy and only
+  overrides the operations needed to change the protocol from XML-RPC to
+  JSON-RPC.
+  """
+
+  def __init__(self, uri, transport=None, encoding=None, verbose=0,
+               allow_none=0, use_datetime=0):
+    urltype, _ = urllib.splittype(uri)
+    if urltype not in ('http', 'https'):
+      raise IOError('unsupported JSON-RPC protocol')
+
+    _base.ServerProxy.__init__(self, uri, transport, encoding, verbose,
+                               allow_none, use_datetime)
+    transport_type, uri = urllib.splittype(uri)
+    if transport is None:
+      if transport_type == 'https':
+        transport = SafeTransport(use_datetime=use_datetime)
+      else:
+        transport = Transport(use_datetime=use_datetime)
+    self.__transport = transport
+
+  def __request(self, methodname, params):
+    """Call a method on the remote server."""
+    request = CreateRequestString(methodname, params)
+
+    response = self.__transport.request(
+        self.__host,
+        self.__handler,
+        request,
+        verbose=self.__verbose
+        )
+
+    return response
+
+
+Server = ServerProxy
diff --git a/testing/legion/legion.isolate b/testing/legion/legion.isolate
index 44bc2e1..63e3fc8 100644
--- a/testing/legion/legion.isolate
+++ b/testing/legion/legion.isolate
@@ -7,6 +7,7 @@
     'files': [
       '__init__.py',
       'common_lib.py',
+      'jsonrpclib.py',
       'legion_test_case.py',
       'legion.isolate',
       'process.py',
@@ -17,6 +18,7 @@
       'task_controller.py',
       'task_registration_server.py',
       'test_controller.py',
+      'SimpleJSONRPCServer.py',
       '../../tools/swarming_client/',
     ],
   },
diff --git a/testing/legion/rpc_server.py b/testing/legion/rpc_server.py
index d9d98cc..24ebe25 100644
--- a/testing/legion/rpc_server.py
+++ b/testing/legion/rpc_server.py
@@ -17,17 +17,16 @@ be achieved in 2 ways:
 import logging
 import threading
 import time
-import xmlrpclib
-import SimpleXMLRPCServer
 import SocketServer
 
 #pylint: disable=relative-import
 import common_lib
 import rpc_methods
 import ssl_util
+import SimpleJSONRPCServer
 
 
-class RequestHandler(SimpleXMLRPCServer.SimpleXMLRPCRequestHandler):
+class RequestHandler(SimpleJSONRPCServer.SimpleJSONRPCRequestHandler):
   """Restricts access to only specified IP address.
 
   This call assumes the server is RPCServer.
@@ -45,7 +44,7 @@ class RequestHandler(SimpleXMLRPCServer.SimpleXMLRPCRequestHandler):
       self.end_headers()
       self.wfile.write(response)
     else:
-      return SimpleXMLRPCServer.SimpleXMLRPCRequestHandler.do_POST(self)
+      return SimpleJSONRPCServer.SimpleJSONRPCRequestHandler.do_POST(self)
 
 
 class RpcServer(ssl_util.SslRpcServer,
@@ -73,7 +72,7 @@ class RpcServer(ssl_util.SslRpcServer,
     idle timeout thread to quit.
     """
     self._shutdown_requested_event.set()
-    SimpleXMLRPCServer.SimpleXMLRPCServer.shutdown(self)
+    SimpleJSONRPCServer.SimpleJSONRPCServer.shutdown(self)
     logging.info('Server shutdown complete')
 
   def serve_forever(self, poll_interval=0.5):
@@ -88,7 +87,7 @@ class RpcServer(ssl_util.SslRpcServer,
     """
     logging.info('RPC server starting')
     self._idle_thread.start()
-    SimpleXMLRPCServer.SimpleXMLRPCServer.serve_forever(self, poll_interval)
+    SimpleJSONRPCServer.SimpleJSONRPCServer.serve_forever(self, poll_interval)
 
   def _dispatch(self, method, params):
     """Dispatch the call to the correct method with the provided params.
@@ -105,7 +104,8 @@ class RpcServer(ssl_util.SslRpcServer,
     """
     logging.debug('Calling %s%s', method, params)
     self._rpc_received_event.set()
-    return SimpleXMLRPCServer.SimpleXMLRPCServer._dispatch(self, method, params)
+    return SimpleJSONRPCServer.SimpleJSONRPCServer._dispatch(
+        self, method, params)
 
   def _CheckForIdleQuit(self):
     """Check for, and exit, if the server is idle for too long.
diff --git a/testing/legion/ssl_util.py b/testing/legion/ssl_util.py
index 1571c77..41e1e75 100644
--- a/testing/legion/ssl_util.py
+++ b/testing/legion/ssl_util.py
@@ -8,11 +8,11 @@ import logging
 import ssl
 import subprocess
 import tempfile
-import xmlrpclib
-import SimpleXMLRPCServer
 
 #pylint: disable=relative-import
 import common_lib
+import jsonrpclib
+import SimpleJSONRPCServer
 
 
 class Error(Exception):
@@ -87,11 +87,11 @@ def _RunCommand(cmd):
   return out
 
 
-class SslRpcServer(SimpleXMLRPCServer.SimpleXMLRPCServer):
+class SslRpcServer(SimpleJSONRPCServer.SimpleJSONRPCServer):
   """Class to add SSL support to the RPC server."""
 
   def __init__(self, *args, **kwargs):
-    SimpleXMLRPCServer.SimpleXMLRPCServer.__init__(self, *args, **kwargs)
+    SimpleJSONRPCServer.SimpleJSONRPCServer.__init__(self, *args, **kwargs)
     self.socket = ssl.wrap_socket(self.socket, certfile=CreatePemFile(),
                                   server_side=True)
 
@@ -100,4 +100,4 @@ class SslRpcServer(SimpleXMLRPCServer.SimpleXMLRPCServer):
     """Creates and returns a connection to an SSL RPC server."""
     addr = 'https://%s:%d' % (server, port)
     logging.debug('Connecting to RPC server at %s', addr)
-    return xmlrpclib.Server(addr, allow_none=True)
+    return jsonrpclib.ServerProxy(addr, allow_none=True)
diff --git a/testing/legion/task_controller.py b/testing/legion/task_controller.py
index 5522a2b..cb7ab15 100644
--- a/testing/legion/task_controller.py
+++ b/testing/legion/task_controller.py
@@ -13,12 +13,12 @@ import subprocess
 import sys
 import tempfile
 import threading
-import xmlrpclib
 
 #pylint: disable=relative-import
 import common_lib
 import process
 import ssl_util
+import jsonrpclib
 
 ISOLATE_PY = os.path.join(common_lib.SWARMING_DIR, 'isolate.py')
 SWARMING_PY = os.path.join(common_lib.SWARMING_DIR, 'swarming.py')
@@ -164,7 +164,7 @@ class TaskController(object):
       logging.info('Releasing %s', self._name)
       try:
         self._rpc.Quit()
-      except (socket.error, xmlrpclib.Fault):
+      except (socket.error, jsonrpclib.Fault):
         logging.error('Unable to connect to %s to call Quit', self.name)
       self._rpc = None
       self._connected = False
diff --git a/testing/legion/task_registration_server.py b/testing/legion/task_registration_server.py
index 85d6eab..1b15952 100644
--- a/testing/legion/task_registration_server.py
+++ b/testing/legion/task_registration_server.py
@@ -11,7 +11,6 @@ is based on an OTP passed to the run_task binary on startup.
 
 import logging
 import threading
-import xmlrpclib
 
 #pylint: disable=relative-import
 import common_lib