#!/usr/bin/env python3 # -*- mode: python -*- # -*- coding: utf-8 -*- # Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """RPC/IPC support.""" import abc import http.client import http.server import io import logging import os import socketserver from avro import io as avro_io from avro import protocol from avro import schema logger = logging.getLogger(__name__) # ------------------------------------------------------------------------------ # Constants def LoadResource(name): dir_path = os.path.dirname(__file__) rsrc_path = os.path.join(dir_path, name) with open(rsrc_path, 'r') as f: return f.read() # Handshake schema is pulled in during build HANDSHAKE_REQUEST_SCHEMA_JSON = LoadResource('HandshakeRequest.avsc') HANDSHAKE_RESPONSE_SCHEMA_JSON = LoadResource('HandshakeResponse.avsc') HANDSHAKE_REQUEST_SCHEMA = schema.Parse(HANDSHAKE_REQUEST_SCHEMA_JSON) HANDSHAKE_RESPONSE_SCHEMA = schema.Parse(HANDSHAKE_RESPONSE_SCHEMA_JSON) HANDSHAKE_REQUESTOR_WRITER = avro_io.DatumWriter(HANDSHAKE_REQUEST_SCHEMA) HANDSHAKE_REQUESTOR_READER = avro_io.DatumReader(HANDSHAKE_RESPONSE_SCHEMA) HANDSHAKE_RESPONDER_WRITER = avro_io.DatumWriter(HANDSHAKE_RESPONSE_SCHEMA) HANDSHAKE_RESPONDER_READER = avro_io.DatumReader(HANDSHAKE_REQUEST_SCHEMA) META_SCHEMA = schema.Parse('{"type": "map", "values": "bytes"}') META_WRITER = avro_io.DatumWriter(META_SCHEMA) META_READER = avro_io.DatumReader(META_SCHEMA) SYSTEM_ERROR_SCHEMA = schema.Parse('["string"]') AVRO_RPC_MIME = 'avro/binary' # protocol cache # Map: remote name -> remote MD5 hash _REMOTE_HASHES = {} # Decoder/encoder for a 32 bits big-endian integer. UINT32_BE = avro_io.STRUCT_INT # Default size of the buffers use to frame messages: BUFFER_SIZE = 8192 # ------------------------------------------------------------------------------ # Exceptions class AvroRemoteException(schema.AvroException): """ Raised when an error message is sent by an Avro requestor or responder. """ def __init__(self, fail_msg=None): schema.AvroException.__init__(self, fail_msg) class ConnectionClosedException(schema.AvroException): pass # ------------------------------------------------------------------------------ # Base IPC Classes (Requestor/Responder) class BaseRequestor(object, metaclass=abc.ABCMeta): """Base class for the client side of a protocol interaction.""" def __init__(self, local_protocol, transceiver): """Initializes a new requestor object. Args: local_protocol: Avro Protocol describing the messages sent and received. transceiver: Transceiver instance to channel messages through. """ self._local_protocol = local_protocol self._transceiver = transceiver self._remote_protocol = None self._remote_hash = None self._send_protocol = None @property def local_protocol(self): """Returns: the Avro Protocol describing the messages sent and received.""" return self._local_protocol @property def transceiver(self): """Returns: the underlying channel used by this requestor.""" return self._transceiver @abc.abstractmethod def _IssueRequest(self, call_request, message_name, request_datum): """TODO: Document this method. Args: call_request: ??? message_name: Name of the message. request_datum: ??? Returns: ??? """ raise Error('Abstract method') def Request(self, message_name, request_datum): """Writes a request message and reads a response or error message. Args: message_name: Name of the IPC method. request_datum: IPC request. Returns: The IPC response. """ # build handshake and call request buffer_writer = io.BytesIO() buffer_encoder = avro_io.BinaryEncoder(buffer_writer) self._WriteHandshakeRequest(buffer_encoder) self._WriteCallRequest(message_name, request_datum, buffer_encoder) # send the handshake and call request; block until call response call_request = buffer_writer.getvalue() return self._IssueRequest(call_request, message_name, request_datum) def _WriteHandshakeRequest(self, encoder): """Emits the handshake request. Args: encoder: Encoder to write the handshake request into. """ local_hash = self._local_protocol.md5 # if self._remote_hash is None: # remote_name = self.transceiver.remote_name # self._remote_hash = _REMOTE_HASHES.get(remote_name) if self._remote_hash is None: self._remote_hash = local_hash self._remote_protocol = self._local_protocol request_datum = { 'clientHash': local_hash, 'serverHash': self._remote_hash, } if self._send_protocol: request_datum['clientProtocol'] = str(self._local_protocol) logger.info('Sending handshake request: %s', request_datum) HANDSHAKE_REQUESTOR_WRITER.write(request_datum, encoder) def _WriteCallRequest(self, message_name, request_datum, encoder): """ The format of a call request is: * request metadata, a map with values of type bytes * the message name, an Avro string, followed by * the message parameters. Parameters are serialized according to the message's request declaration. """ # request metadata (not yet implemented) request_metadata = {} META_WRITER.write(request_metadata, encoder) # Identify message to send: message = self.local_protocol.message_map.get(message_name) if message is None: raise schema.AvroException('Unknown message: %s' % message_name) encoder.write_utf8(message.name) # message parameters self._WriteRequest(message.request, request_datum, encoder) def _WriteRequest(self, request_schema, request_datum, encoder): logger.info('writing request: %s', request_datum) datum_writer = avro_io.DatumWriter(request_schema) datum_writer.write(request_datum, encoder) def _ReadHandshakeResponse(self, decoder): """Reads and processes the handshake response message. Args: decoder: Decoder to read messages from. Returns: call-response exists (boolean) ??? Raises: schema.AvroException on ??? """ handshake_response = HANDSHAKE_REQUESTOR_READER.read(decoder) logger.info('Processing handshake response: %s', handshake_response) match = handshake_response['match'] if match == 'BOTH': # Both client and server protocol hashes match: self._send_protocol = False return True elif match == 'CLIENT': # Client's side hash mismatch: self._remote_protocol = \ protocol.Parse(handshake_response['serverProtocol']) self._remote_hash = handshake_response['serverHash'] self._send_protocol = False return True elif match == 'NONE': # Neither client nor server match: self._remote_protocol = \ protocol.Parse(handshake_response['serverProtocol']) self._remote_hash = handshake_response['serverHash'] self._send_protocol = True return False else: raise schema.AvroException('handshake_response.match=%r' % match) def _ReadCallResponse(self, message_name, decoder): """Reads and processes a method call response. The format of a call response is: - response metadata, a map with values of type bytes - a one-byte error flag boolean, followed by either: - if the error flag is false, the message response, serialized per the message's response schema. - if the error flag is true, the error, serialized per the message's error union schema. Args: message_name: decoder: Returns: ??? Raises: schema.AvroException on ??? """ # response metadata response_metadata = META_READER.read(decoder) # remote response schema remote_message_schema = self._remote_protocol.message_map.get(message_name) if remote_message_schema is None: raise schema.AvroException('Unknown remote message: %s' % message_name) # local response schema local_message_schema = self._local_protocol.message_map.get(message_name) if local_message_schema is None: raise schema.AvroException('Unknown local message: %s' % message_name) # error flag if not decoder.read_boolean(): writer_schema = remote_message_schema.response reader_schema = local_message_schema.response return self._ReadResponse(writer_schema, reader_schema, decoder) else: writer_schema = remote_message_schema.errors reader_schema = local_message_schema.errors raise self._ReadError(writer_schema, reader_schema, decoder) def _ReadResponse(self, writer_schema, reader_schema, decoder): datum_reader = avro_io.DatumReader(writer_schema, reader_schema) result = datum_reader.read(decoder) return result def _ReadError(self, writer_schema, reader_schema, decoder): datum_reader = avro_io.DatumReader(writer_schema, reader_schema) return AvroRemoteException(datum_reader.read(decoder)) class Requestor(BaseRequestor): """Concrete requestor implementation.""" def _IssueRequest(self, call_request, message_name, request_datum): call_response = self.transceiver.Transceive(call_request) # process the handshake and call response buffer_decoder = avro_io.BinaryDecoder(io.BytesIO(call_response)) call_response_exists = self._ReadHandshakeResponse(buffer_decoder) if call_response_exists: return self._ReadCallResponse(message_name, buffer_decoder) else: return self.Request(message_name, request_datum) # ------------------------------------------------------------------------------ class Responder(object, metaclass=abc.ABCMeta): """Base class for the server side of a protocol interaction.""" def __init__(self, local_protocol): self._local_protocol = local_protocol self._local_hash = self._local_protocol.md5 self._protocol_cache = {} self.set_protocol_cache(self._local_hash, self._local_protocol) @property def local_protocol(self): return self._local_protocol # utility functions to manipulate protocol cache def get_protocol_cache(self, hash): return self._protocol_cache.get(hash) def set_protocol_cache(self, hash, protocol): self._protocol_cache[hash] = protocol def Respond(self, call_request): """Entry point to process one procedure call. Args: call_request: Serialized procedure call request. Returns: Serialized procedure call response. Raises: ??? """ buffer_reader = io.BytesIO(call_request) buffer_decoder = avro_io.BinaryDecoder(buffer_reader) buffer_writer = io.BytesIO() buffer_encoder = avro_io.BinaryEncoder(buffer_writer) error = None response_metadata = {} try: remote_protocol = self._ProcessHandshake(buffer_decoder, buffer_encoder) # handshake failure if remote_protocol is None: return buffer_writer.getvalue() # read request using remote protocol request_metadata = META_READER.read(buffer_decoder) remote_message_name = buffer_decoder.read_utf8() # get remote and local request schemas so we can do # schema resolution (one fine day) remote_message = remote_protocol.message_map.get(remote_message_name) if remote_message is None: fail_msg = 'Unknown remote message: %s' % remote_message_name raise schema.AvroException(fail_msg) local_message = self.local_protocol.message_map.get(remote_message_name) if local_message is None: fail_msg = 'Unknown local message: %s' % remote_message_name raise schema.AvroException(fail_msg) writer_schema = remote_message.request reader_schema = local_message.request request = self._ReadRequest(writer_schema, reader_schema, buffer_decoder) logger.info('Processing request: %r', request) # perform server logic try: response = self.Invoke(local_message, request) except AvroRemoteException as exn: error = exn except Exception as exn: error = AvroRemoteException(str(exn)) # write response using local protocol META_WRITER.write(response_metadata, buffer_encoder) buffer_encoder.write_boolean(error is not None) if error is None: writer_schema = local_message.response self._WriteResponse(writer_schema, response, buffer_encoder) else: writer_schema = local_message.errors self._WriteError(writer_schema, error, buffer_encoder) except schema.AvroException as exn: error = AvroRemoteException(str(exn)) buffer_encoder = avro_io.BinaryEncoder(io.StringIO()) META_WRITER.write(response_metadata, buffer_encoder) buffer_encoder.write_boolean(True) self._WriteError(SYSTEM_ERROR_SCHEMA, error, buffer_encoder) return buffer_writer.getvalue() def _ProcessHandshake(self, decoder, encoder): """Processes an RPC handshake. Args: decoder: Where to read from. encoder: Where to write to. Returns: The requested Protocol. """ handshake_request = HANDSHAKE_RESPONDER_READER.read(decoder) logger.info('Processing handshake request: %s', handshake_request) # determine the remote protocol client_hash = handshake_request.get('clientHash') client_protocol = handshake_request.get('clientProtocol') remote_protocol = self.get_protocol_cache(client_hash) if remote_protocol is None and client_protocol is not None: remote_protocol = protocol.Parse(client_protocol) self.set_protocol_cache(client_hash, remote_protocol) # evaluate remote's guess of the local protocol server_hash = handshake_request.get('serverHash') handshake_response = {} if self._local_hash == server_hash: if remote_protocol is None: handshake_response['match'] = 'NONE' else: handshake_response['match'] = 'BOTH' else: if remote_protocol is None: handshake_response['match'] = 'NONE' else: handshake_response['match'] = 'CLIENT' if handshake_response['match'] != 'BOTH': handshake_response['serverProtocol'] = str(self.local_protocol) handshake_response['serverHash'] = self._local_hash logger.info('Handshake response: %s', handshake_response) HANDSHAKE_RESPONDER_WRITER.write(handshake_response, encoder) return remote_protocol @abc.abstractmethod def Invoke(self, local_message, request): """Processes one procedure call. Args: local_message: Avro message specification. request: Call request. Returns: Call response. Raises: ??? """ raise Error('abtract method') def _ReadRequest(self, writer_schema, reader_schema, decoder): datum_reader = avro_io.DatumReader(writer_schema, reader_schema) return datum_reader.read(decoder) def _WriteResponse(self, writer_schema, response_datum, encoder): datum_writer = avro_io.DatumWriter(writer_schema) datum_writer.write(response_datum, encoder) def _WriteError(self, writer_schema, error_exception, encoder): datum_writer = avro_io.DatumWriter(writer_schema) datum_writer.write(str(error_exception), encoder) # ------------------------------------------------------------------------------ # Framed message class FramedReader(object): """Wrapper around a file-like object to read framed data.""" def __init__(self, reader): self._reader = reader def Read(self): """Reads one message from the configured reader. Returns: The message, as bytes. """ message = io.BytesIO() # Read and append frames until we encounter a 0-size frame: while self._ReadFrame(message) > 0: pass return message.getvalue() def _ReadFrame(self, message): """Reads and appends one frame into the given message bytes. Args: message: Message to append the frame to. Returns: Size of the frame that was read. The empty frame (size 0) indicates the end of a message. """ frame_size = self._ReadInt32() remaining = frame_size while remaining > 0: data_bytes = self._reader.read(remaining) if len(data_bytes) == 0: raise ConnectionClosedException( 'FramedReader: expecting %d more bytes in frame of size %d, got 0.' % (remaining, frame_size)) message.write(data_bytes) remaining -= len(data_bytes) return frame_size def _ReadInt32(self): encoded = self._reader.read(UINT32_BE.size) if len(encoded) != UINT32_BE.size: raise ConnectionClosedException('Invalid header: %r' % encoded) return UINT32_BE.unpack(encoded)[0] class FramedWriter(object): """Wrapper around a file-like object to write framed data.""" def __init__(self, writer): self._writer = writer def Write(self, message): """Writes a message. Message is chunked into sequences of frames terminated by an empty frame. Args: message: Message to write, as bytes. """ while len(message) > 0: chunk_size = max(BUFFER_SIZE, len(message)) chunk = message[:chunk_size] self._WriteBuffer(chunk) message = message[chunk_size:] # A message is always terminated by a zero-length buffer. self._WriteUnsignedInt32(0) def _WriteBuffer(self, chunk): self._WriteUnsignedInt32(len(chunk)) self._writer.write(chunk) def _WriteUnsignedInt32(self, uint32): self._writer.write(UINT32_BE.pack(uint32)) # ------------------------------------------------------------------------------ # Transceiver (send/receive channel) class Transceiver(object, metaclass=abc.ABCMeta): @abc.abstractproperty def remote_name(self): pass @abc.abstractmethod def ReadMessage(self): """Reads a single message from the channel. Blocks until a message can be read. Returns: The message read from the channel. """ pass @abc.abstractmethod def WriteMessage(self, message): """Writes a message into the channel. Blocks until the message has been written. Args: message: Message to write. """ pass def Transceive(self, request): """Processes a single request-reply interaction. Synchronous request-reply interaction. Args: request: Request message. Returns: The reply message. """ self.WriteMessage(request) result = self.ReadMessage() return result def Close(self): """Closes this transceiver.""" pass class HTTPTransceiver(Transceiver): """HTTP-based transceiver implementation.""" def __init__(self, host, port, req_resource='/'): """Initializes a new HTTP transceiver. Args: host: Name or IP address of the remote host to interact with. port: Port the remote server is listening on. req_resource: Optional HTTP resource path to use, '/' by default. """ self._req_resource = req_resource self._conn = http.client.HTTPConnection(host, port) self._conn.connect() self._remote_name = self._conn.sock.getsockname() @property def remote_name(self): return self._remote_name def ReadMessage(self): response = self._conn.getresponse() response_reader = FramedReader(response) framed_message = response_reader.Read() response.read() # ensure we're ready for subsequent requests return framed_message def WriteMessage(self, message): req_method = 'POST' req_headers = {'Content-Type': AVRO_RPC_MIME} bio = io.BytesIO() req_body_buffer = FramedWriter(bio) req_body_buffer.Write(message) req_body = bio.getvalue() self._conn.request(req_method, self._req_resource, req_body, req_headers) def Close(self): self._conn.close() self._conn = None # ------------------------------------------------------------------------------ # Server Implementations def _MakeHandlerClass(responder): class AvroHTTPRequestHandler(http.server.BaseHTTPRequestHandler): def do_POST(self): reader = FramedReader(self.rfile) call_request = reader.Read() logger.info('Serialized request: %r', call_request) call_response = responder.Respond(call_request) logger.info('Serialized response: %r', call_response) self.send_response(200) self.send_header('Content-type', AVRO_RPC_MIME) self.end_headers() framed_writer = FramedWriter(self.wfile) framed_writer.Write(call_response) self.wfile.flush() logger.info('Response sent') return AvroHTTPRequestHandler class MultiThreadedHTTPServer( socketserver.ThreadingMixIn, http.server.HTTPServer, ): """Multi-threaded HTTP server.""" pass class AvroIpcHttpServer(MultiThreadedHTTPServer): """Avro IPC server implemented on top of an HTTP server.""" def __init__(self, interface, port, responder): """Initializes a new Avro IPC server. Args: interface: Interface the server listens on, eg. 'localhost' or '0.0.0.0'. port: TCP port the server listens on, eg. 8000. responder: Responder implementation to handle RPCs. """ super(AvroIpcHttpServer, self).__init__( server_address=(interface, port), RequestHandlerClass=_MakeHandlerClass(responder), ) if __name__ == '__main__': raise Exception('Not a standalone module')