You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
267 lines
8.5 KiB
Python
267 lines
8.5 KiB
Python
1 year ago
|
import asyncio
|
||
|
|
import functools
|
||
|
|
|
||
|
|
from loguru import logger
|
||
|
|
|
||
|
|
from .busrt_client import (
|
||
|
|
ERR_IO,
|
||
|
|
ERR_TIMEOUT,
|
||
|
|
OP_MESSAGE,
|
||
|
|
RESPONSE_OK,
|
||
|
|
on_frame_default,
|
||
|
|
)
|
||
|
|
from .msgutils import deserialize, serialize
|
||
|
|
|
||
|
|
RPC_NOTIFICATION_HEADER = b'\x00'
|
||
|
|
RPC_REQUEST_HEADER = b'\x01'
|
||
|
|
RPC_REPLY_HEADER = b'\x11'
|
||
|
|
RPC_ERROR_REPLY_HEADER = b'\x12'
|
||
|
|
|
||
|
|
RPC_NOTIFICATION = 0x00
|
||
|
|
RPC_REQUEST = 0x01
|
||
|
|
RPC_REPLY = 0x11
|
||
|
|
RPC_ERROR = 0x12
|
||
|
|
|
||
|
|
RPC_ERROR_CODE_PARSE = -32700
|
||
|
|
RPC_ERROR_CODE_INVALID_REQUEST = -32600
|
||
|
|
RPC_ERROR_CODE_METHOD_NOT_FOUND = -32601
|
||
|
|
RPC_ERROR_CODE_INVALID_METHOD_PARAMS = -32602
|
||
|
|
RPC_ERROR_CODE_INTERNAL = -32603
|
||
|
|
|
||
|
|
|
||
|
|
async def on_call_default(event):
|
||
|
|
raise RpcException('RPC Engine not initialized',
|
||
|
|
RPC_ERROR_CODE_METHOD_NOT_FOUND)
|
||
|
|
|
||
|
|
|
||
|
|
async def on_notification_default(event):
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
def format_rpc_e_msg(e):
|
||
|
|
if isinstance(e, RpcException):
|
||
|
|
return e.rpc_error_payload
|
||
|
|
else:
|
||
|
|
return str(e)
|
||
|
|
|
||
|
|
|
||
|
|
class RpcException(Exception):
|
||
|
|
|
||
|
|
def __init__(self, msg: str | bytes = '', code=RPC_ERROR_CODE_INTERNAL):
|
||
|
|
self.rpc_error_code = code
|
||
|
|
self.rpc_error_payload = msg
|
||
|
|
super().__init__(msg if isinstance(msg, str) else msg.decode())
|
||
|
|
|
||
|
|
def __str__(self):
|
||
|
|
return super().__str__() + f' (code: {self.rpc_error_code})'
|
||
|
|
|
||
|
|
|
||
|
|
class Rpc:
|
||
|
|
|
||
|
|
def __init__(self, client):
|
||
|
|
self.client = client
|
||
|
|
self.client.on_frame = self._handle_frame
|
||
|
|
self.call_id = 0
|
||
|
|
self.call_lock = asyncio.Lock()
|
||
|
|
self.calls = {}
|
||
|
|
self.on_frame = on_frame_default
|
||
|
|
self.on_call = on_call_default
|
||
|
|
self.on_notification = on_notification_default
|
||
|
|
|
||
|
|
def is_connected(self):
|
||
|
|
return self.client.connected
|
||
|
|
|
||
|
|
def notify(self, target, notification):
|
||
|
|
return self.client.send(target, notification)
|
||
|
|
|
||
|
|
def call0(self, target, request):
|
||
|
|
request.header = RPC_REQUEST_HEADER + b'\x00\x00\x00\x00' + \
|
||
|
|
request.method + b'\x00'
|
||
|
|
return self.client.send(target, request)
|
||
|
|
|
||
|
|
async def call(self, target, request):
|
||
|
|
async with self.call_lock:
|
||
|
|
call_id = self.call_id + 1
|
||
|
|
if call_id == 0xffff_ffff:
|
||
|
|
self.call_id = 0
|
||
|
|
else:
|
||
|
|
self.call_id = call_id
|
||
|
|
call_event = RpcCallEvent()
|
||
|
|
self.calls[call_id] = call_event
|
||
|
|
request.header = RPC_REQUEST_HEADER + call_id.to_bytes(
|
||
|
|
4, 'little') + request.method + b'\x00'
|
||
|
|
try:
|
||
|
|
try:
|
||
|
|
code = await (await self.client.send(
|
||
|
|
target,
|
||
|
|
request)).wait_completed(timeout=self.client.timeout)
|
||
|
|
if code != RESPONSE_OK:
|
||
|
|
try:
|
||
|
|
del self.calls[call_id]
|
||
|
|
except KeyError:
|
||
|
|
pass
|
||
|
|
err_code = -32000 - code
|
||
|
|
call_event.error = RpcException('RPC error', code=err_code)
|
||
|
|
call_event.completed.set()
|
||
|
|
except asyncio.TimeoutError:
|
||
|
|
try:
|
||
|
|
del self.calls[call_id]
|
||
|
|
except KeyError:
|
||
|
|
pass
|
||
|
|
err_code = -32000 - ERR_TIMEOUT
|
||
|
|
call_event.error = RpcException('RPC timeout', code=err_code)
|
||
|
|
call_event.completed.set()
|
||
|
|
except Exception as e:
|
||
|
|
try:
|
||
|
|
del self.calls[call_id]
|
||
|
|
except KeyError:
|
||
|
|
pass
|
||
|
|
call_event.error = RpcException(str(e), code=-32000 - ERR_IO)
|
||
|
|
call_event.completed.set()
|
||
|
|
return call_event
|
||
|
|
|
||
|
|
def __getattr__(self, method):
|
||
|
|
return functools.partial(self.params_call, method=method)
|
||
|
|
|
||
|
|
async def params_call(self, target, method: str, **kwargs):
|
||
|
|
params = kwargs
|
||
|
|
if '_raw' in kwargs:
|
||
|
|
params = kwargs.pop('_raw')
|
||
|
|
params = params or None
|
||
|
|
c0 = False
|
||
|
|
if method.endswith('0'):
|
||
|
|
c0 = True
|
||
|
|
method = method[:-1]
|
||
|
|
request = Request(method, serialize(params))
|
||
|
|
if c0:
|
||
|
|
await self.call0(target, request)
|
||
|
|
return None
|
||
|
|
result = await self.call(target, request)
|
||
|
|
return deserialize((await result.wait_completed()).get_payload())
|
||
|
|
|
||
|
|
async def _handle_frame(self, frame):
|
||
|
|
try:
|
||
|
|
if frame.type == OP_MESSAGE:
|
||
|
|
if frame.payload[0] == RPC_NOTIFICATION:
|
||
|
|
event = Event(RPC_NOTIFICATION, frame, 1)
|
||
|
|
await self.on_notification(event)
|
||
|
|
elif frame.payload[0] == RPC_REQUEST:
|
||
|
|
sender = frame.sender
|
||
|
|
call_id_b = frame.payload[1:5]
|
||
|
|
call_id = int.from_bytes(call_id_b, 'little')
|
||
|
|
method = frame.payload[5:5 +
|
||
|
|
frame.payload[5:].index(b'\x00')]
|
||
|
|
event = Event(RPC_REQUEST, frame, 6 + len(method))
|
||
|
|
event.call_id = call_id
|
||
|
|
event.method = method
|
||
|
|
if call_id == 0:
|
||
|
|
await self.on_call(event)
|
||
|
|
else:
|
||
|
|
reply = Reply()
|
||
|
|
try:
|
||
|
|
reply.payload = await self.on_call(event)
|
||
|
|
if reply.payload is None:
|
||
|
|
reply.payload = b''
|
||
|
|
reply.header = RPC_REPLY_HEADER + call_id_b
|
||
|
|
except Exception as e:
|
||
|
|
code = getattr(e, 'rpc_error_code', RPC_ERROR_CODE_INTERNAL)
|
||
|
|
|
||
|
|
reply.header = (
|
||
|
|
RPC_ERROR_REPLY_HEADER + call_id_b +
|
||
|
|
code.to_bytes(2, 'little', signed=True))
|
||
|
|
reply.payload = format_rpc_e_msg(e)
|
||
|
|
await self.client.send(sender, reply)
|
||
|
|
elif frame.payload[0] == RPC_REPLY or frame.payload[
|
||
|
|
0] == RPC_ERROR:
|
||
|
|
call_id = int.from_bytes(frame.payload[1:5], 'little')
|
||
|
|
try:
|
||
|
|
call_event = self.calls.pop(call_id)
|
||
|
|
call_event.frame = frame
|
||
|
|
if frame.payload[0] == RPC_ERROR:
|
||
|
|
err_code = int.from_bytes(frame.payload[5:7],
|
||
|
|
'little',
|
||
|
|
signed=True)
|
||
|
|
call_event.error = RpcException(frame.payload[7:],
|
||
|
|
code=err_code)
|
||
|
|
call_event.completed.set()
|
||
|
|
except KeyError:
|
||
|
|
logger.warning(f'orphaned RPC response: {call_id}')
|
||
|
|
else:
|
||
|
|
logger.error(f'Invalid RPC frame code: {frame.payload[0]}')
|
||
|
|
else:
|
||
|
|
await self.on_frame(frame)
|
||
|
|
except Exception:
|
||
|
|
import traceback
|
||
|
|
logger.error(traceback.format_exc())
|
||
|
|
|
||
|
|
|
||
|
|
class RpcCallEvent:
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
self.frame = None
|
||
|
|
self.error = None
|
||
|
|
self.completed = asyncio.Event()
|
||
|
|
|
||
|
|
def is_completed(self):
|
||
|
|
return self.completed.is_set()
|
||
|
|
|
||
|
|
async def wait_completed(self, timeout=None):
|
||
|
|
if timeout:
|
||
|
|
if not await asyncio.wait_for(self.completed.wait(),
|
||
|
|
timeout=timeout):
|
||
|
|
raise TimeoutError
|
||
|
|
else:
|
||
|
|
await self.completed.wait()
|
||
|
|
if self.error is None:
|
||
|
|
return self
|
||
|
|
else:
|
||
|
|
raise self.error
|
||
|
|
|
||
|
|
def get_payload(self):
|
||
|
|
return self.frame.payload[5:]
|
||
|
|
|
||
|
|
def is_empty(self):
|
||
|
|
return len(self.frame.payload) <= 5
|
||
|
|
|
||
|
|
|
||
|
|
class Event:
|
||
|
|
|
||
|
|
def __init__(self, tp, frame, payload_pos):
|
||
|
|
self.tp = tp
|
||
|
|
self.frame = frame
|
||
|
|
self._payload_pos = payload_pos
|
||
|
|
|
||
|
|
def get_payload(self):
|
||
|
|
return self.frame.payload[self._payload_pos:]
|
||
|
|
|
||
|
|
|
||
|
|
class Notification:
|
||
|
|
|
||
|
|
def __init__(self, payload=b''):
|
||
|
|
self.payload = payload
|
||
|
|
self.type = OP_MESSAGE
|
||
|
|
self.qos = 1
|
||
|
|
self.header = RPC_NOTIFICATION_HEADER
|
||
|
|
|
||
|
|
|
||
|
|
class Request:
|
||
|
|
|
||
|
|
def __init__(self, method, params=b''):
|
||
|
|
self.payload = b'' if params is None else params
|
||
|
|
self.type = OP_MESSAGE
|
||
|
|
self.qos = 1
|
||
|
|
self.method = method.encode() if isinstance(method, str) else method
|
||
|
|
self.header = None
|
||
|
|
|
||
|
|
|
||
|
|
class Reply:
|
||
|
|
|
||
|
|
def __init__(self, result=b''):
|
||
|
|
self.payload = result
|
||
|
|
self.type = OP_MESSAGE
|
||
|
|
self.qos = 1
|
||
|
|
self.header = None
|
||
|
|
|
||
|
|
|
||
|
|
__all__ = ["Rpc", "Request", "Reply", "RpcCallEvent", "Event", "RpcException"]
|