def test_read_next_chunk_resize(self):
# Setup: Create a byte array for test input
test_bytes = bytearray(b'1234567890')
with io.BytesIO(test_bytes) as stream:
# If:
# ... I create a reader with an artificially low initial buffer size
# and prefill the buffer
reader = JSONRPCReader(stream, logger=utils.get_mock_logger())
reader._buffer = bytearray(5)
reader._buffer_end_offset = 4
# ... and I read a chunk from the stream
result = reader._read_next_chunk()
# Then:
# ... The read should have succeeded
self.assertTrue(result, True)
# ... The size of the buffer should have doubled
self.assertEqual(len(reader._buffer), 10)
# ... The buffer end offset should be the size of the buffer
self.assertEqual(reader._buffer_end_offset, 10)
# ... The buffer should contain the first 6 elements of the test data
expected = test_bytes[:6]
actual = reader._buffer[4:]
self.assertEqual(actual, expected)
def test_read_next_chunk_eof(self):
with io.BytesIO() as stream:
# If:
# ... I create a reader with a stream that has no bytes
reader = JSONRPCReader(stream, logger=utils.get_mock_logger())
# ... and I read a chunk from the stream
# Then: I should get an exception
with self.assertRaises(EOFError):
reader._read_next_chunk()
def test_closes(self):
with io.BytesIO(b'') as stream:
# If:
# ... I create a JSON RPC reader with an opened stream
reader = JSONRPCReader(stream, logger=utils.get_mock_logger())
# ... and I close the reader
reader.close()
# Then: The stream should be closed
self.assertTrue(reader.stream.closed)
def test_read_next_chunk_closed(self):
# Setup: Create a stream that has already been closed
stream = io.BytesIO()
stream.close()
# If:
# ... I create a reader with a closed stream
reader = JSONRPCReader(stream, logger=utils.get_mock_logger())
# ... and I read a chunk from the stream
# Then: I should get an exception
with self.assertRaises(ValueError):
reader._read_next_chunk()
def test_read_message_invalid_json(self):
# Setup: Reader with a stream that has an invalid message
test_bytes = bytearray(b'Content-Length: 10\r\n\r\nabcdefghij')
with io.BytesIO(test_bytes) as stream:
reader = JSONRPCReader(stream, logger=utils.get_mock_logger())
reader._buffer = bytearray(100)
# If: I read a message
# Then:
# ... It should throw an exception
with self.assertRaises(ValueError):
reader.read_message()
# ... The buffer should be trashed
self.assertEqual(len(reader._buffer), reader.DEFAULT_BUFFER_SIZE)
def test_closes_exception():
# Setup: Patch the stream to have a custom close handler
stream = io.BytesIO(b'')
close_orig = stream.close
stream.close = mock.MagicMock(side_effect=AttributeError)
# If: Close a reader and it throws an exception
logger = utils.get_mock_logger()
reader = JSONRPCReader(stream, logger=logger)
reader.close()
# Then: There should not have been an exception throws
logger.exception.assert_called_once()
# Cleanup: Close the stream
close_orig()
def test_read_headers_not_found(self):
# Setup: Create a reader with a loaded buffer that does not contain the \r\n\r\n control
reader = JSONRPCReader(None, logger=utils.get_mock_logger())
reader._buffer = bytearray(b'1234567890')
reader._buffer_end_offset = len(reader._buffer)
# If: I look for a header block in the buffer
result = reader._try_read_headers()
# Then:
# ... I should not have found any
self.assertFalse(result)
# ... The current reading position of the buffer should not have moved
self.assertEqual(reader._read_offset, 0)
self.assertEqual(reader._read_state, JSONRPCReader.ReadState.Header)
def test_read_next_chunk_success(self):
# Setup: Create a byte array for test input
test_bytes = bytearray(b'123')
with io.BytesIO(test_bytes) as stream:
# If: I attempt to read a chunk from the stream
reader = JSONRPCReader(stream, logger=utils.get_mock_logger())
result = reader._read_next_chunk()
# Then:
# ... The result should be true
self.assertTrue(result)
# ... The buffer should contain 3 byte and should not have been read yet
self.assertEqual(reader._buffer_end_offset, len(test_bytes))
self.assertEqual(reader._read_offset, 0)
# ... The buffer should now contain the bytes from the stream
buffer_contents = reader._buffer[:len(test_bytes)]
self.assertEqual(buffer_contents, test_bytes)
def test_read_headers_success(self):
# Setup: Create a reader with a loaded buffer that contains a a complete header
reader = JSONRPCReader(None, logger=utils.get_mock_logger())
reader._buffer = bytearray(b'Content-Length: 56\r\n\r\n')
reader._buffer_end_offset = len(reader._buffer)
# If: I look for a header block in the buffer
result = reader._try_read_headers()
# Then:
# ... I should have found a header
self.assertTrue(result)
# ... The current reading position should have moved to the end of the buffer
self.assertEqual(reader._read_offset, len(reader._buffer))
self.assertEqual(reader._read_state, JSONRPCReader.ReadState.Content)
# ... The headers should have been stored
self.assertEqual(reader._expected_content_length, 56)
self.assertDictEqual(reader._headers, {'content-length': '56'})
def test_create_standard_encoding(self):
with io.BytesIO(b'') as stream:
# If: I create a JSON RPC reader with a stream without specifying the encoding
reader = JSONRPCReader(stream, logger=utils.get_mock_logger())
# Then: The stream and encoding should be set appropriately
self.assertIsNotNone(reader)
self.assertIs(reader.stream, stream)
self.assertEqual(reader.encoding, 'UTF-8')
self.assertEqual(reader._read_state,
JSONRPCReader.ReadState.Header)
def test_read_multiple_messages(self):
test_string = b'Content-Length: 32\r\n\r\n{"method":"test", "params":null}'
test_bytes = bytearray(test_string + test_string)
with io.BytesIO(test_bytes) as stream:
reader = JSONRPCReader(stream, logger=utils.get_mock_logger())
reader._buffer = bytearray(100)
# If:
# ... I read a message
msg1 = reader.read_message()
# ... And I read another message
msg2 = reader.read_message()
# Then:
# ... The messages should be real
self.assertIsNotNone(msg1)
self.assertIsNotNone(msg2)
# ... The buffer should have been trashed
self.assertEqual(len(reader._buffer), reader.DEFAULT_BUFFER_SIZE)
def test_read_message_multi_read_content(self):
# Setup: Reader with a stream that has an entire message read
test_bytes = bytearray(
b'Content-Length: 32\r\n\r\n{"method":"test", "params":null}')
with io.BytesIO(test_bytes) as stream:
reader = JSONRPCReader(stream, logger=utils.get_mock_logger())
reader._buffer = bytearray(25)
# If: I read a message with the reader
message = reader.read_message()
# Then:
# ... I should have a successful message
self.assertIsNotNone(message)
# ... The reader should be back in header mode
self.assertEqual(reader._read_state,
JSONRPCReader.ReadState.Header)
# ... The buffer should have been trimmed
self.assertEqual(len(reader._buffer),
JSONRPCReader.DEFAULT_BUFFER_SIZE)
def test_create_nonstandard_encoding(self):
with io.BytesIO(b'') as stream:
# If: I create a JSON RPC reader with a non-standard encoding
reader = JSONRPCReader(stream,
encoding="ascii",
logger=utils.get_mock_logger())
# Then: The stream and encoding should be set appropriately
self.assertIsNotNone(reader)
self.assertIs(reader.stream, stream)
self.assertEqual(reader.encoding, 'ascii')
self.assertEqual(reader._read_state,
JSONRPCReader.ReadState.Header)
def __init__(self, in_stream, out_stream, logger=None, version='0'):
"""
Initializes internal state of the server and sets up a few useful built-in request handlers
:param in_stream: Input stream that will provide messages from the client
:param out_stream: Output stream that will send message to the client
:param logger: Optional logger
:param version: Protocol version. Defaults to 0
"""
self.writer = JSONRPCWriter(out_stream, logger=logger)
self.reader = JSONRPCReader(in_stream, logger=logger)
self._logger = logger
self._version = version
self._stop_requested = False
self._output_queue = Queue()
self._request_handlers = {}
self._notification_handlers = {}
self._shutdown_handlers = []
self._output_consumer = None
self._input_consumer = None
# Register built-in handlers
# 1) Echo
echo_config = IncomingMessageConfiguration('echo', None)
self.set_request_handler(echo_config, self._handle_echo_request)
# 2) Protocol version
version_config = IncomingMessageConfiguration('version', None)
self.set_request_handler(version_config, self._handle_version_request)
# 3) Shutdown/exit
shutdown_config = IncomingMessageConfiguration('shutdown', None)
self.set_request_handler(shutdown_config,
self._handle_shutdown_request)
exit_config = IncomingMessageConfiguration('exit', None)
self.set_request_handler(exit_config, self._handle_shutdown_request)
def test_read_content_success(self):
# Setup: Create a reader that has read in headers and has all of a message buffered
test_buffer = bytearray(b"message")
reader = JSONRPCReader(None, logger=utils.get_mock_logger())
reader._buffer = test_buffer
reader._buffer_end_offset = len(reader._buffer)
reader._read_offset = 0
reader._read_state = JSONRPCReader.ReadState.Content
reader._expected_content_length = 5
# If: I read a message from the buffer
output = ['']
result = reader._try_read_content(output)
# Then:
# ... The message should be successfully read
self.assertTrue(result)
self.assertEqual(output[0], 'messa')
# ... The state of the reader should have been updated
self.assertEqual(reader._read_state, JSONRPCReader.ReadState.Header)
self.assertEqual(reader._read_offset, 5)
self.assertEqual(reader._buffer_end_offset, len(reader._buffer))
def test_read_headers_no_colon(self):
# Setup: Create a reader with a buffer that contains the control sequence but does not
# match the header format
test_buffer = bytearray(b'1234567890\r\n\r\n')
reader = JSONRPCReader(None, logger=utils.get_mock_logger())
reader._buffer = test_buffer
reader._buffer_end_offset = len(reader._buffer)
reader._read_offset = 1
# If: I look for a header block in the buffer
# Then:
# ... I should get an exception b/c of the malformed header
with self.assertRaises(KeyError):
reader._try_read_headers()
# ... The current reading position of the buffer should be reset to 0
self.assertEqual(reader._read_offset, 0)
# ... The buffer should have been trashed
self.assertIsNot(reader._buffer, test_buffer)
def test_read_content_not_enough_buffer(self):
# Setup: Create a reader that has read in headers and has part of a message buffered
test_buffer = bytearray(b'message')
reader = JSONRPCReader(None, logger=utils.get_mock_logger())
reader._buffer = test_buffer
reader._buffer_end_offset = len(reader._buffer)
reader._read_offset = 0
reader._read_state = JSONRPCReader.ReadState.Content
reader._expected_content_length = 15
# If: I read a message from the buffer
output = ['']
result = reader._try_read_content(output)
# Then:
# ... The result should be false and the output should be empty
self.assertFalse(result)
self.assertEqual(output[0], '')
# ... The state of the reader should stay the same
self.assertEqual(reader._read_state, JSONRPCReader.ReadState.Content)
self.assertEqual(reader._read_offset, 0)
self.assertEqual(reader._expected_content_length, 15)
self.assertEqual(reader._buffer_end_offset, len(reader._buffer))
def test_read_recover_from_content_message(self):
test_string = b'Content-Length: 10\r\n\r\nabcdefghij' + \
b'Content-Length: 32\r\n\r\n{"method":"test", "params":null}'
test_bytes = bytearray(test_string)
with io.BytesIO(test_bytes) as stream:
reader = JSONRPCReader(stream, logger=utils.get_mock_logger())
reader._buffer = bytearray(100)
# If: I read a message with invalid content
# Then: I should get an exception
with self.assertRaises(ValueError):
reader.read_message()
# If: I read another valid message
msg = reader.read_message()
# Then: I should have a valid message
self.assertIsNotNone(msg)
def test_read_headers_bad_format(self):
# Setup: Create a reader with a header block that contains invalid content-length
test_buffer = bytearray(b'Content-Length: abc\r\n\r\n')
reader = JSONRPCReader(None, logger=utils.get_mock_logger())
reader._buffer = test_buffer
reader._buffer_end_offset = len(reader._buffer)
reader._read_offset = 1
# If: I look for a header block in the buffer
# Then:
# ... I should get an exception from there not being a content-length header
with self.assertRaises(LookupError):
reader._try_read_headers()
# ... The current reading position of the buffer should be reset to 0
self.assertEqual(reader._read_offset, 0)
# ... The buffer should have been trashed
self.assertIsNot(reader._buffer, test_buffer)
# ... The headers should have been trashed
self.assertEqual(len(reader._headers), 0)
class JSONRPCServer:
"""
Handles async requests, async notifications, and async responses
"""
# CONSTANTS ############################################################
OUTPUT_THREAD_NAME = u"JSON_RPC_Output_Thread"
INPUT_THREAD_NAME = u"JSON_RPC_Input_Thread"
class Handler:
def __init__(self, class_, handler):
self.class_ = class_
self.handler = handler
def __init__(self, in_stream, out_stream, logger=None, version='0'):
"""
Initializes internal state of the server and sets up a few useful built-in request handlers
:param in_stream: Input stream that will provide messages from the client
:param out_stream: Output stream that will send message to the client
:param logger: Optional logger
:param version: Protocol version. Defaults to 0
"""
self.writer = JSONRPCWriter(out_stream, logger=logger)
self.reader = JSONRPCReader(in_stream, logger=logger)
self._logger = logger
self._version = version
self._stop_requested = False
self._output_queue = Queue()
self._request_handlers = {}
self._notification_handlers = {}
self._shutdown_handlers = []
self._output_consumer = None
self._input_consumer = None
# Register built-in handlers
# 1) Echo
echo_config = IncomingMessageConfiguration('echo', None)
self.set_request_handler(echo_config, self._handle_echo_request)
# 2) Protocol version
version_config = IncomingMessageConfiguration('version', None)
self.set_request_handler(version_config, self._handle_version_request)
# 3) Shutdown/exit
shutdown_config = IncomingMessageConfiguration('shutdown', None)
self.set_request_handler(shutdown_config,
self._handle_shutdown_request)
exit_config = IncomingMessageConfiguration('exit', None)
self.set_request_handler(exit_config, self._handle_shutdown_request)
# METHODS ##############################################################
def add_shutdown_handler(self, handler):
"""
Adds the provided shutdown handler to the list of shutdown handlers
:param handler: The callable handler to call when shutdown occurs
"""
self._shutdown_handlers.append(handler)
def count_shutdown_handlers(self) -> int:
"""
Returns the number of shutdown handlers registered
"""
return len(self._shutdown_handlers)
def start(self):
"""
Starts the background threads to listen for responses and requests from the underlying
streams. Encapsulated into its own method for future async extensions without threads
"""
if self._logger is not None:
self._logger.info("JSON RPC server starting...")
self._output_consumer = threading.Thread(target=self._consume_output,
name=self.OUTPUT_THREAD_NAME)
self._output_consumer.daemon = True
self._output_consumer.start()
self._input_consumer = threading.Thread(target=self._consume_input,
name=self.INPUT_THREAD_NAME)
self._input_consumer.daemon = True
self._input_consumer.start()
def stop(self):
"""
Signal input and output threads to halt asap
"""
self._stop_requested = True
# Enqueue None to optimistically unblock output thread so it can check for the cancellation flag
self._output_queue.put(None)
if self._logger is not None:
self._logger.info('JSON RPC server stopping...')
def send_request(self, method, params):
"""
Add a new request to the output queue
:param method: Method string of the request
:param params: Data to send along with the request
"""
message_id = str(uuid.uuid4())
# Create the message
message = JSONRPCMessage.create_request(message_id, method, params)
# TODO: Add support for handlers for the responses
# Add the message to the output queue
self._output_queue.put(message)
def send_notification(self, method, params):
"""
Sends a notification, independent of any request
:param method: String name of the method for the notification
:param params: Data to send with the notification
"""
# Create the message
message = JSONRPCMessage.create_notification(method, params)
# TODO: Add support for handlers for the responses
# Add the message to the output queue
self._output_queue.put(message)
def set_request_handler(self, config, handler):
"""
Sets the handler for a request with a given configuration
:param config: Configuration of the request to listen for
:param handler: Handler to call when the server receives a request that matches the config
"""
self._request_handlers[config.method] = self.Handler(
config.parameter_class, handler)
def set_notification_handler(self, config, handler):
"""
Sets the handler for a notification with a given configuration
:param config: Configuration of the notification to listen for
:param handler: Handler to call when the server receives a notification that matches the config
"""
self._notification_handlers[config.method] = self.Handler(
config.parameter_class, handler)
def wait_for_exit(self):
"""
Blocks until both input and output threads return, ie, until the server stops.
"""
self._input_consumer.join()
self._output_consumer.join()
if self._logger is not None:
self._logger.info('Input and output threads have completed')
# Close the reader/writer here instead of in the stop method in order to allow "softer"
# shutdowns that will read or write the last message before halting
self.reader.close()
self.writer.close()
# BUILT-IN HANDLERS ####################################################
@staticmethod
def _handle_echo_request(request_context, params):
request_context.send_response(params)
def _handle_version_request(self, request_context, params):
request_context.send_response(self._version)
def _handle_shutdown_request(self, request_context, params):
# Signal that the threads should stop
if self._logger is not None:
self._logger.info('Received shutdown request')
self._stop_requested = True
# Execute the shutdown request handlers
for handler in self._shutdown_handlers:
handler()
self.stop()
# IMPLEMENTATION DETAILS ###############################################
def _consume_input(self):
"""
Listen for messages from the input stream and dispatch them to the registered listeners
:raises ValueError: The stream was closed. Exit the thread immediately.
:raises LookupError: No void header with content-length was found
:raises EOFError: The stream may not contain any bytes yet, so retry.
"""
if self._logger is not None:
self._logger.info('Input thread started')
while not self._stop_requested:
try:
message = self.reader.read_message()
self._dispatch_message(message)
except EOFError as error:
# Thread fails once we read EOF. Halt the input thread
self._log_exception(error, self.INPUT_THREAD_NAME)
self.stop()
break
except (LookupError, ValueError) as error:
# LookupError: Content-Length header was not found
# ValueError: JSON deserialization failed
self._log_exception(error, self.INPUT_THREAD_NAME)
# Do not halt the input thread
except Exception as error:
# Catch generic exceptions
self._log_exception(error, self.INPUT_THREAD_NAME)
# Do not halt the input thread
def _consume_output(self):
"""
Send output over the output stream
"""
if self._logger is not None:
self._logger.info('Output thread started')
while not self._stop_requested:
try:
# Block until queue contains a message to send
message = self._output_queue.get()
if message is not None:
# It is necessary to check for None here b/c unblock the queue get by adding
# None when we want to stop the service
self.writer.send_message(message)
except ValueError as error:
# Stream is closed, break out of the loop
self._log_exception(error, self.OUTPUT_THREAD_NAME)
break
except Exception as error:
# Catch generic exceptions without breaking out of loop
self._log_exception(error, self.OUTPUT_THREAD_NAME)
def _dispatch_message(self, message):
"""
Dispatches a message that was received to the necessary handler
:param message: The message that was received
"""
if message.message_type in [
JSONRPCMessageType.ResponseSuccess,
JSONRPCMessageType.ResponseError
]:
# Responses need to be routed to the handler that requested them
# TODO: Route to the handler or send error message
return
# Figure out which handler will execute the request/notification
if message.message_type is JSONRPCMessageType.Request:
if self._logger is not None:
self._logger.info('Received request id=%s method=%s',
message.message_id, message.message_method)
handler = self._request_handlers.get(message.message_method)
request_context = RequestContext(message, self._output_queue)
# Make sure we got a handler for the request
if handler is None:
# TODO: Localize?
request_context.send_error(
f'Requested method is unsupported: {message.message_method}'
)
if self._logger is not None:
self._logger.warn('Requested method is unsupported: %s',
message.message_method)
return
# Call the handler with a request context and the deserialized parameter object
if handler.class_ is None:
# Don't attempt to do complex deserialization
deserialized_object = message.message_params
else:
# Use the complex deserializer
deserialized_object = handler.class_.from_dict(
message.message_params)
try:
handler.handler(request_context, deserialized_object)
except Exception as e:
error_message = f'Unhandled exception while handling request method {message.message_method}: "{e}"' # TODO: Localize
if self._logger is not None:
self._logger.exception(error_message)
request_context.send_error(error_message, code=-32603)
elif message.message_type is JSONRPCMessageType.Notification:
if self._logger is not None:
self._logger.info('Received notification method=%s',
message.message_method)
handler = self._notification_handlers.get(message.message_method)
if handler is None:
# Ignore the notification
if self._logger is not None:
self._logger.warn('Notification method %s is unsupported',
message.message_method)
return
# Call the handler with a notification context
notification_context = NotificationContext(self._output_queue)
deserialized_object = None
if handler.class_ is None:
# Don't attempt to do complex deserialization
deserialized_object = message.message_params
else:
# Use the complex deserializer
deserialized_object = handler.class_.from_dict(
message.message_params)
try:
handler.handler(notification_context, deserialized_object)
except Exception:
error_message = f'Unhandled exception while handling notification method {message.message_method}'
if self._logger is not None:
self._logger.exception(error_message)
else:
# If this happens we have a serious issue with the JSON RPC reader
if self._logger is not None:
self._logger.warn('Received unsupported message type %s',
message.message_type)
return
def _log_exception(self, ex, thread_name):
"""
Logs an exception if the logger is defined
:param ex: Exception to log
:param thread_name: Name of the thread that encountered the exception
"""
if self._logger is not None:
self._logger.exception('Thread %s encountered exception %s',
thread_name, ex)