# wireprotov2peer.py - client side code for wire protocol version 2
#
# Copyright 2018 Gregory Szorc <gregory.szorc@gmail.com>
#
# This software may be used and distributed according to the terms of the
# GNU General Public License version 2 or any later version.
from __future__ import absolute_import
from .i18n import _
from .thirdparty import (
cbor,
)
from . import (
encoding,
error,
util,
wireprotoframing,
)
def formatrichmessage(atoms):
"""Format an encoded message from the framing protocol."""
chunks = []
for atom in atoms:
msg = _(atom[b'msg'])
if b'args' in atom:
msg = msg % atom[b'args']
chunks.append(msg)
return b''.join(chunks)
class commandresponse(object):
"""Represents the response to a command request."""
def __init__(self, requestid, command):
self.requestid = requestid
self.command = command
self.b = util.bytesio()
def cborobjects(self):
"""Obtain decoded CBOR objects from this response."""
size = self.b.tell()
self.b.seek(0)
decoder = cbor.CBORDecoder(self.b)
while self.b.tell() < size:
yield decoder.decode()
class clienthandler(object):
"""Object to handle higher-level client activities.
The ``clientreactor`` is used to hold low-level state about the frame-based
protocol, such as which requests and streams are active. This type is used
for higher-level operations, such as reading frames from a socket, exposing
and managing a higher-level primitive for representing command responses,
etc. This class is what peers should probably use to bridge wire activity
with the higher-level peer API.
"""
def __init__(self, ui, clientreactor):
self._ui = ui
self._reactor = clientreactor
self._requests = {}
self._futures = {}
self._responses = {}
def callcommand(self, command, args, f):
"""Register a request to call a command.
Returns an iterable of frames that should be sent over the wire.
"""
request, action, meta = self._reactor.callcommand(command, args)
if action != 'noop':
raise error.ProgrammingError('%s not yet supported' % action)
rid = request.requestid
self._requests[rid] = request
self._futures[rid] = f
self._responses[rid] = commandresponse(rid, command)
return iter(())
def flushcommands(self):
"""Flush all queued commands.
Returns an iterable of frames that should be sent over the wire.
"""
action, meta = self._reactor.flushcommands()
if action != 'sendframes':
raise error.ProgrammingError('%s not yet supported' % action)
return meta['framegen']
def readframe(self, fh):
"""Attempt to read and process a frame.
Returns None if no frame was read. Presumably this means EOF.
"""
frame = wireprotoframing.readframe(fh)
if frame is None:
# TODO tell reactor?
return
self._ui.note(_('received %r\n') % frame)
self._processframe(frame)
return True
def _processframe(self, frame):
"""Process a single read frame."""
action, meta = self._reactor.onframerecv(frame)
if action == 'error':
e = error.RepoError(meta['message'])
if frame.requestid in self._futures:
self._futures[frame.requestid].set_exception(e)
else:
raise e
if frame.requestid not in self._requests:
raise error.ProgrammingError(
'received frame for unknown request; this is either a bug in '
'the clientreactor not screening for this or this instance was '
'never told about this request: %r' % frame)
response = self._responses[frame.requestid]
if action == 'responsedata':
response.b.write(meta['data'])
if meta['eos']:
# If the command has a decoder, resolve the future to the
# decoded value. Otherwise resolve to the rich response object.
decoder = COMMAND_DECODERS.get(response.command)
# TODO consider always resolving the overall status map.
if decoder:
objs = response.cborobjects()
overall = next(objs)
if overall['status'] == 'ok':
self._futures[frame.requestid].set_result(decoder(objs))
else:
e = error.RepoError(
formatrichmessage(overall['error']['message']))
self._futures[frame.requestid].set_exception(e)
else:
self._futures[frame.requestid].set_result(response)
del self._requests[frame.requestid]
del self._futures[frame.requestid]
else:
raise error.ProgrammingError(
'unhandled action from clientreactor: %s' % action)
def decodebranchmap(objs):
# Response should be a single CBOR map of branch name to array of nodes.
bm = next(objs)
return {encoding.tolocal(k): v for k, v in bm.items()}
def decodeheads(objs):
# Array of node bytestrings.
return next(objs)
def decodeknown(objs):
# Bytestring where each byte is a 0 or 1.
raw = next(objs)
return [True if c == '1' else False for c in raw]
def decodelistkeys(objs):
# Map with bytestring keys and values.
return next(objs)
def decodelookup(objs):
return next(objs)
def decodepushkey(objs):
return next(objs)
COMMAND_DECODERS = {
'branchmap': decodebranchmap,
'heads': decodeheads,
'known': decodeknown,
'listkeys': decodelistkeys,
'lookup': decodelookup,
'pushkey': decodepushkey,
}