# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2.1 of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
#
# StreamRpcChannel.py
# Copyright (C) 2005 Simon Newton
"""A RpcChannel that works over a TCP socket."""
__author__ = 'nomis52@gmail.com (Simon Newton)'
import logging
import struct
from google.protobuf import service
from ola.rpc import Rpc_pb2
from ola.rpc.SimpleRpcController import SimpleRpcController
class OutstandingRequest(object):
"""These represent requests on the server side that haven't completed yet."""
def __init__(self, request_id, controller):
self.controller = controller
self.id = request_id
class OutstandingResponse(object):
"""These represent requests on the client side that haven't completed yet."""
def __init__(self, response_id, controller, callback, reply_class):
self.callback = callback
self.controller = controller
self.id = response_id
self.reply = None
self.reply_class = reply_class
class StreamRpcChannel(service.RpcChannel):
"""Implements a RpcChannel over a TCP socket."""
PROTOCOL_VERSION = 1
VERSION_MASK = 0xf0000000
SIZE_MASK = 0x0fffffff
RECEIVE_BUFFER_SIZE = 8192
def __init__(self, socket, service_impl, close_callback = None):
"""Create a new StreamRpcChannel.
Args:
socket: the socket to communicate on
service: the service to route incoming requests to
"""
self._service = service_impl
self._socket = socket
self._sequence = 0
self._outstanding_requests = {}
self._outstanding_responses = {}
self._buffer = [] # the received data
self._expected_size = None # the size of the message we're receiving
self._skip_message = False # skip the current message
self._close_callback = close_callback
def SocketReady(self):
"""Read data from the socket and handle when we get a full message.
Returns:
True if the socket remains connected, False if it was closed.
"""
data = self._socket.recv(self.RECEIVE_BUFFER_SIZE)
if data == '':
logging.info('OLAD Server Socket closed')
if self._close_callback is not None:
self._close_callback()
return False
self._buffer.append(data)
self._ProcessIncomingData()
return True
def CallMethod(self, method, controller, request, response_pb, done):
"""Call a method.
Don't use this directly, use the Stub object.
Args:
method: The MethodDescriptor to call
controller: An RpcController object
request: The request message
response: The response class
done: A closure to call once complete.
"""
message = Rpc_pb2.RpcMessage()
message.type = Rpc_pb2.REQUEST
message.id = self._sequence
message.name = method.name
message.buffer = request.SerializeToString()
self._SendMessage(message)
self._sequence += 1
if message.id in self._outstanding_responses:
# fail any outstanding response with the same id, not the best approach
# but it'll do for now.
logging.warning('Response %d already pending, failing now', message.id)
response = self._outstanding_responses[message.id]
response.controller.SetFailed('Duplicate request found')
self._InvokeCallback(response)
response = OutstandingResponse(message.id, controller, done, response_pb)
self._outstanding_responses[message.id] = response
def RequestComplete(self, request, response):
"""This is called on the server side when a request has completed.
Args:
request: the OutstandingRequest object that has completed.
response: the response to send to the client
"""
message = Rpc_pb2.RpcMessage()
if request.controller.Failed():
self._SendRequestFailed(request)
return
message.type = Rpc_pb2.RESPONSE
message.id = request.id
message.buffer = response.SerializeToString()
self._SendMessage(message)
del self._outstanding_requests[request.id]
def _EncodeHeader(self, size):
"""Encode a version and size into a header.
Args:
size: the size of the rpc message
Returns:
The header field
"""
header = (self.PROTOCOL_VERSION << 28) & self.VERSION_MASK
header |= size & self.SIZE_MASK
return struct.pack('=L', header)
def _DecodeHeader(self, header):
"""Decode a header into the version and size.
Args:
header: the received header
Return:
A tuple in the form (version, size)
"""
return ((header & self.VERSION_MASK) >> 28, header & self.SIZE_MASK)
def _SendMessage(self, message):
"""Send an RpcMessage.
Args:
message: An RpcMessage object.
Returns:
True if the send succeeded, False otherwise.
"""
data = message.SerializeToString()
# combine into one buffer to send so we avoid sending two packets
data = self._EncodeHeader(len(data)) + data
sent_bytes = self._socket.send(data)
if sent_bytes != len(data):
logging.warning('Failed to send full datagram')
return False
return True
def _SendRequestFailed(self, request):
"""Send a response indiciating this request failed.
Args:
request: An OutstandingRequest object.
"""
message = Rpc_pb2.RpcMessage()
message.type = Rpc_pb2.RESPONSE_FAILED
message.id = request.id
message.buffer = request.controller.ErrorText()
self._SendMessage(message)
del self._outstanding_requests[request.id]
def _SendNotImplemented(self, message_id):
"""Send a not-implemented response.
Args:
message_id: The message number.
"""
message = Rpc_pb2.RpcMessage
message.type = Rpc_pb2.RESPONSE_NOT_IMPLEMENTED
message.id = message_id
self._SendMessage(message)
def _GrabData(self, size):
"""Fetch the next N bytes of data from the buffer.
Args:
size: the amount of data to fetch.
Returns:
The next size bytes of data, or None if there isn't enough.
"""
data_size = sum([len(s) for s in self._buffer])
if data_size < size:
return None
data = []
size_left = size
while size_left > 0:
chunk = self._buffer.pop(0)
if len(chunk) > size_left:
# we only want part of it
data.append(chunk[0:size_left])
self._buffer.insert(0, chunk[size_left:])
size_left = 0
else:
data.append(chunk)
size_left -= len(chunk)
return b''.join(data)
def _ProcessIncomingData(self):
"""Process the received data."""
while True:
if not self._expected_size:
# this is a new msg
raw_header = self._GrabData(4)
if not raw_header:
# not enough data yet
return
header = struct.unpack('<L', raw_header)[0]
version, size = self._DecodeHeader(header)
if version != self.PROTOCOL_VERSION:
logging.warning('Protocol mismatch: %d != %d', version,
self.PROTOCOL_VERSION)
self._skip_message = True
self._expected_size = size
data = self._GrabData(self._expected_size)
if not data:
# not enough data yet
return
if not self._skip_message:
self._HandleNewMessage(data)
self._expected_size = 0
self._skip_message = 0
def _HandleNewMessage(self, data):
"""Handle a new Message.
Args:
data: A chunk of data representing a RpcMessage
"""
message = Rpc_pb2.RpcMessage()
message.ParseFromString(data)
if message.type in self.MESSAGE_HANDLERS:
self.MESSAGE_HANDLERS[message.type](self, message)
else:
logging.warning('Not sure of message type %d', message.type())
def _HandleRequest(self, message):
"""Handle a Request message.
Args:
message: The RpcMessage object.
"""
if not self._service:
logging.warning('No service registered')
return
descriptor = self._service.GetDescriptor()
method = descriptor.FindMethodByName(message.name)
if not method:
logging.warning('Failed to get method descriptor for %s', message.name)
self._SendNotImplemented(message.id)
return
request_pb = self._service.GetRequestClass(method)()
request_pb.ParseFromString(message.buffer)
controller = SimpleRpcController()
request = OutstandingRequest(message.id, controller)
if message.id in self._outstanding_requests:
logging.warning('Duplicate request for %d', message.id)
self._SendRequestFailed(message.id)
self._outstanding_requests[message.id] = request
callback = lambda x: self.RequestComplete(request, x)
self._service.CallMethod(method, request.controller, request_pb, callback)
def _HandleResponse(self, message):
"""Handle a Response message.
Args:
message: The RpcMessage object.
"""
response = self._outstanding_responses.get(message.id, None)
if response:
response.reply = response.reply_class()
response.reply.ParseFromString(message.buffer)
self._InvokeCallback(response)
def _HandleCanceledResponse(self, message):
"""Handle a Not Implemented message.
Args:
message: The RpcMessage object.
"""
logging.warning('Received a canceled response')
response = self._outstanding_responses.get(message.id, None)
if response:
response.controller.SetFailed(message.buffer)
self._InvokeCallback(response)
def _HandleFailedReponse(self, message):
"""Handle a Failed Response message.
Args:
message: The RpcMessage object.
"""
response = self._outstanding_responses.get(message.id, None)
if response:
response.controller.SetFailed(message.buffer)
self._InvokeCallback(response)
def _HandleNotImplemented(self, message):
"""Handle a Not Implemented message.
Args:
message: The RpcMessage object.
"""
logging.warning('Received a non-implemented response')
response = self._outstanding_responses.get(message.id, None)
if response:
response.controller.SetFailed('Not Implemented')
self._InvokeCallback(response)
def _InvokeCallback(self, response):
"""Run the callback and delete the outstanding response.
Args:
response: the resoponse to run.
"""
response.callback(response.controller, response.reply)
del self._outstanding_responses[response.id]
MESSAGE_HANDLERS = {
Rpc_pb2.REQUEST: _HandleRequest,
Rpc_pb2.RESPONSE: _HandleResponse,
Rpc_pb2.RESPONSE_CANCEL: _HandleCanceledResponse,
Rpc_pb2.RESPONSE_FAILED: _HandleFailedReponse,
Rpc_pb2.RESPONSE_NOT_IMPLEMENTED: _HandleNotImplemented,
}