wireprototypes.py
453 lines
| 13.3 KiB
| text/x-python
|
PythonLexer
/ mercurial / wireprototypes.py
|
r36090 | # 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 | ||||
|
r37630 | from .node import ( | ||
bin, | ||||
hex, | ||||
) | ||||
|
r37801 | from .i18n import _ | ||
|
r43346 | from .thirdparty import attr | ||
|
r37801 | from . import ( | ||
error, | ||||
util, | ||||
) | ||||
|
r43346 | from .interfaces import util as interfaceutil | ||
from .utils import compression | ||||
|
r36389 | |||
|
r36553 | # Names of the SSH protocol implementations. | ||
|
r43347 | SSHV1 = b'ssh-v1' | ||
|
r37064 | # These are advertised over the wire. Increment the counters at the end | ||
|
r36553 | # to reflect BC breakages. | ||
|
r43347 | SSHV2 = b'exp-ssh-v2-0003' | ||
HTTP_WIREPROTO_V2 = b'exp-http-v2-0003' | ||||
|
r36553 | |||
|
r43347 | NARROWCAP = b'exp-narrow-1' | ||
ELLIPSESCAP1 = b'exp-ellipses-1' | ||||
ELLIPSESCAP = b'exp-ellipses-2' | ||||
|
r42605 | SUPPORTED_ELLIPSESCAP = (ELLIPSESCAP1, ELLIPSESCAP) | ||
|
r40110 | |||
|
r36627 | # All available wire protocol transports. | ||
TRANSPORTS = { | ||||
|
r43347 | SSHV1: {b'transport': b'ssh', b'version': 1,}, | ||
|
r36627 | SSHV2: { | ||
|
r43347 | b'transport': b'ssh', | ||
|
r37310 | # TODO mark as version 2 once all commands are implemented. | ||
|
r43347 | b'version': 1, | ||
|
r36627 | }, | ||
|
r43347 | b'http-v1': {b'transport': b'http', b'version': 1,}, | ||
HTTP_WIREPROTO_V2: {b'transport': b'http', b'version': 2,}, | ||||
|
r36627 | } | ||
|
r43346 | |||
|
r36091 | class bytesresponse(object): | ||
"""A wire protocol response consisting of raw bytes.""" | ||||
|
r43346 | |||
|
r36091 | def __init__(self, data): | ||
self.data = data | ||||
|
r43346 | |||
|
r36090 | class ooberror(object): | ||
"""wireproto reply: failure of a batch of operation | ||||
Something failed during a batch call. The error message is stored in | ||||
`self.message`. | ||||
""" | ||||
|
r43346 | |||
|
r36090 | def __init__(self, message): | ||
self.message = message | ||||
|
r43346 | |||
|
r36090 | class pushres(object): | ||
"""wireproto reply: success with simple integer return | ||||
The call was successful and returned an integer contained in `self.res`. | ||||
""" | ||||
|
r43346 | |||
|
r36090 | def __init__(self, res, output): | ||
self.res = res | ||||
self.output = output | ||||
|
r43346 | |||
|
r36090 | class pusherr(object): | ||
"""wireproto reply: failure | ||||
The call failed. The `self.res` attribute contains the error message. | ||||
""" | ||||
|
r43346 | |||
|
r36090 | def __init__(self, res, output): | ||
self.res = res | ||||
self.output = output | ||||
|
r43346 | |||
|
r36090 | class streamres(object): | ||
"""wireproto reply: binary stream | ||||
The call was successful and the result is a stream. | ||||
Accepts a generator containing chunks of data to be sent to the client. | ||||
``prefer_uncompressed`` indicates that the data is expected to be | ||||
uncompressable and that the stream should therefore use the ``none`` | ||||
engine. | ||||
""" | ||||
|
r43346 | |||
|
r36090 | def __init__(self, gen=None, prefer_uncompressed=False): | ||
self.gen = gen | ||||
self.prefer_uncompressed = prefer_uncompressed | ||||
|
r43346 | |||
|
r36090 | class streamreslegacy(object): | ||
"""wireproto reply: uncompressed binary stream | ||||
The call was successful and the result is a stream. | ||||
Accepts a generator containing chunks of data to be sent to the client. | ||||
Like ``streamres``, but sends an uncompressed data for "version 1" clients | ||||
using the application/mercurial-0.1 media type. | ||||
""" | ||||
|
r43346 | |||
|
r36090 | def __init__(self, gen=None): | ||
self.gen = gen | ||||
|
r36389 | |||
|
r43346 | |||
|
r37630 | # list of nodes encoding / decoding | ||
|
r43347 | def decodelist(l, sep=b' '): | ||
|
r37630 | if l: | ||
|
r43346 | return [bin(v) for v in l.split(sep)] | ||
|
r37630 | return [] | ||
|
r43346 | |||
|
r43347 | def encodelist(l, sep=b' '): | ||
|
r37630 | try: | ||
return sep.join(map(hex, l)) | ||||
except TypeError: | ||||
raise | ||||
|
r43346 | |||
|
r37630 | # batched call argument encoding | ||
|
r43346 | |||
|
r37630 | def escapebatcharg(plain): | ||
|
r43346 | return ( | ||
|
r43347 | plain.replace(b':', b':c') | ||
.replace(b',', b':o') | ||||
.replace(b';', b':s') | ||||
.replace(b'=', b':e') | ||||
|
r43346 | ) | ||
|
r37630 | |||
def unescapebatcharg(escaped): | ||||
|
r43346 | return ( | ||
|
r43347 | escaped.replace(b':e', b'=') | ||
.replace(b':s', b';') | ||||
.replace(b':o', b',') | ||||
.replace(b':c', b':') | ||||
|
r43346 | ) | ||
|
r37630 | |||
|
r37631 | # mapping of options accepted by getbundle and their types | ||
# | ||||
|
r43225 | # Meant to be extended by extensions. It is the extension's responsibility to | ||
# ensure such options are properly processed in exchange.getbundle. | ||||
|
r37631 | # | ||
# supported types are: | ||||
# | ||||
|
r43225 | # :nodes: list of binary nodes, transmitted as space-separated hex nodes | ||
# :csv: list of values, transmitted as comma-separated values | ||||
# :scsv: set of values, transmitted as comma-separated values | ||||
|
r37631 | # :plain: string with no transformation needed. | ||
GETBUNDLE_ARGUMENTS = { | ||||
|
r43347 | b'heads': b'nodes', | ||
b'bookmarks': b'boolean', | ||||
b'common': b'nodes', | ||||
b'obsmarkers': b'boolean', | ||||
b'phases': b'boolean', | ||||
b'bundlecaps': b'scsv', | ||||
b'listkeys': b'csv', | ||||
b'cg': b'boolean', | ||||
b'cbattempted': b'boolean', | ||||
b'stream': b'boolean', | ||||
b'includepats': b'csv', | ||||
b'excludepats': b'csv', | ||||
|
r37631 | } | ||
|
r43346 | |||
|
r37828 | class baseprotocolhandler(interfaceutil.Interface): | ||
|
r36389 | """Abstract base class for wire protocol handlers. | ||
A wire protocol handler serves as an interface between protocol command | ||||
handlers and the wire protocol transport layer. Protocol handlers provide | ||||
methods to read command arguments, redirect stdio for the duration of | ||||
the request, handle response types, etc. | ||||
""" | ||||
|
r37828 | name = interfaceutil.Attribute( | ||
|
r36389 | """The name of the protocol implementation. | ||
Used for uniquely identifying the transport type. | ||||
|
r43346 | """ | ||
) | ||||
|
r36389 | |||
|
r37312 | def getargs(args): | ||
|
r36389 | """return the value for arguments in <args> | ||
|
r37503 | For version 1 transports, returns a list of values in the same | ||
order they appear in ``args``. For version 2 transports, returns | ||||
a dict mapping argument name to value. | ||||
""" | ||||
|
r36389 | |||
|
r37411 | def getprotocaps(): | ||
"""Returns the list of protocol-level capabilities of client | ||||
Returns a list of capabilities as declared by the client for | ||||
the current request (or connection for stateful protocol handlers).""" | ||||
|
r37432 | def getpayload(): | ||
"""Provide a generator for the raw payload. | ||||
|
r36389 | |||
|
r37432 | The caller is responsible for ensuring that the full payload is | ||
processed. | ||||
|
r36389 | """ | ||
|
r37312 | def mayberedirectstdio(): | ||
|
r36389 | """Context manager to possibly redirect stdio. | ||
The context manager yields a file-object like object that receives | ||||
stdout and stderr output when the context manager is active. Or it | ||||
yields ``None`` if no I/O redirection occurs. | ||||
The intent of this context manager is to capture stdio output | ||||
so it may be sent in the response. Some transports support streaming | ||||
stdio to the client in real time. For these transports, stdio output | ||||
won't be captured. | ||||
""" | ||||
|
r37312 | def client(): | ||
|
r36389 | """Returns a string representation of this client (as bytes).""" | ||
|
r36631 | |||
|
r37312 | def addcapabilities(repo, caps): | ||
|
r36631 | """Adds advertised capabilities specific to this protocol. | ||
Receives the list of capabilities collected so far. | ||||
Returns a list of capabilities. The passed in argument can be returned. | ||||
""" | ||||
|
r36819 | |||
|
r37312 | def checkperm(perm): | ||
|
r36819 | """Validate that the client has permissions to perform a request. | ||
The argument is the permission required to proceed. If the client | ||||
doesn't have that permission, the exception should raise or abort | ||||
in a protocol specific manner. | ||||
""" | ||||
|
r37799 | |||
|
r43346 | |||
|
r37799 | class commandentry(object): | ||
"""Represents a declared wire protocol command.""" | ||||
|
r43346 | |||
def __init__( | ||||
self, | ||||
func, | ||||
|
r43347 | args=b'', | ||
|
r43346 | transports=None, | ||
|
r43347 | permission=b'push', | ||
|
r43346 | cachekeyfn=None, | ||
extracapabilitiesfn=None, | ||||
): | ||||
|
r37799 | self.func = func | ||
self.args = args | ||||
self.transports = transports or set() | ||||
self.permission = permission | ||||
|
r40057 | self.cachekeyfn = cachekeyfn | ||
|
r40208 | self.extracapabilitiesfn = extracapabilitiesfn | ||
|
r37799 | |||
def _merge(self, func, args): | ||||
"""Merge this instance with an incoming 2-tuple. | ||||
This is called when a caller using the old 2-tuple API attempts | ||||
to replace an instance. The incoming values are merged with | ||||
data not captured by the 2-tuple and a new instance containing | ||||
the union of the two objects is returned. | ||||
""" | ||||
|
r43346 | return commandentry( | ||
func, | ||||
args=args, | ||||
transports=set(self.transports), | ||||
permission=self.permission, | ||||
) | ||||
|
r37799 | |||
# Old code treats instances as 2-tuples. So expose that interface. | ||||
def __iter__(self): | ||||
yield self.func | ||||
yield self.args | ||||
def __getitem__(self, i): | ||||
if i == 0: | ||||
return self.func | ||||
elif i == 1: | ||||
return self.args | ||||
else: | ||||
|
r43347 | raise IndexError(b'can only access elements 0 and 1') | ||
|
r37799 | |||
|
r43346 | |||
|
r37799 | class commanddict(dict): | ||
"""Container for registered wire protocol commands. | ||||
It behaves like a dict. But __setitem__ is overwritten to allow silent | ||||
coercion of values from 2-tuples for API compatibility. | ||||
""" | ||||
|
r43346 | |||
|
r37799 | def __setitem__(self, k, v): | ||
if isinstance(v, commandentry): | ||||
pass | ||||
# Cast 2-tuples to commandentry instances. | ||||
elif isinstance(v, tuple): | ||||
if len(v) != 2: | ||||
|
r43347 | raise ValueError(b'command tuples must have exactly 2 elements') | ||
|
r37799 | |||
# It is common for extensions to wrap wire protocol commands via | ||||
# e.g. ``wireproto.commands[x] = (newfn, args)``. Because callers | ||||
# doing this aren't aware of the new API that uses objects to store | ||||
# command entries, we automatically merge old state with new. | ||||
if k in self: | ||||
v = self[k]._merge(v[0], v[1]) | ||||
else: | ||||
# Use default values from @wireprotocommand. | ||||
|
r43346 | v = commandentry( | ||
v[0], | ||||
args=v[1], | ||||
transports=set(TRANSPORTS), | ||||
|
r43347 | permission=b'push', | ||
|
r43346 | ) | ||
|
r37799 | else: | ||
|
r43346 | raise ValueError( | ||
|
r43347 | b'command entries must be commandentry instances ' | ||
b'or 2-tuples' | ||||
|
r43346 | ) | ||
|
r37799 | |||
return super(commanddict, self).__setitem__(k, v) | ||||
def commandavailable(self, command, proto): | ||||
"""Determine if a command is available for the requested protocol.""" | ||||
assert proto.name in TRANSPORTS | ||||
entry = self.get(command) | ||||
if not entry: | ||||
return False | ||||
if proto.name not in entry.transports: | ||||
return False | ||||
return True | ||||
|
r37801 | |||
|
r43346 | |||
|
r37801 | def supportedcompengines(ui, role): | ||
"""Obtain the list of supported compression engines for a request.""" | ||||
r42208 | assert role in (compression.CLIENTROLE, compression.SERVERROLE) | |||
|
r37801 | |||
r42208 | compengines = compression.compengines.supportedwireengines(role) | |||
|
r37801 | |||
# Allow config to override default list and ordering. | ||||
r42208 | if role == compression.SERVERROLE: | |||
|
r43347 | configengines = ui.configlist(b'server', b'compressionengines') | ||
config = b'server.compressionengines' | ||||
|
r37801 | else: | ||
# This is currently implemented mainly to facilitate testing. In most | ||||
# cases, the server should be in charge of choosing a compression engine | ||||
# because a server has the most to lose from a sub-optimal choice. (e.g. | ||||
# CPU DoS due to an expensive engine or a network DoS due to poor | ||||
# compression ratio). | ||||
|
r43346 | configengines = ui.configlist( | ||
|
r43347 | b'experimental', b'clientcompressionengines' | ||
|
r43346 | ) | ||
|
r43347 | config = b'experimental.clientcompressionengines' | ||
|
r37801 | |||
# No explicit config. Filter out the ones that aren't supposed to be | ||||
# advertised and return default ordering. | ||||
if not configengines: | ||||
|
r43347 | attr = ( | ||
b'serverpriority' if role == util.SERVERROLE else b'clientpriority' | ||||
) | ||||
|
r43346 | return [ | ||
e for e in compengines if getattr(e.wireprotosupport(), attr) > 0 | ||||
] | ||||
|
r37801 | |||
# If compression engines are listed in the config, assume there is a good | ||||
# reason for it (like server operators wanting to achieve specific | ||||
# performance characteristics). So fail fast if the config references | ||||
# unusable compression engines. | ||||
validnames = set(e.name() for e in compengines) | ||||
invalidnames = set(e for e in configengines if e not in validnames) | ||||
if invalidnames: | ||||
|
r43346 | raise error.Abort( | ||
|
r43347 | _(b'invalid compression engine defined in %s: %s') | ||
% (config, b', '.join(sorted(invalidnames))) | ||||
|
r43346 | ) | ||
|
r37801 | |||
compengines = [e for e in compengines if e.name() in configengines] | ||||
|
r43346 | compengines = sorted( | ||
compengines, key=lambda e: configengines.index(e.name()) | ||||
) | ||||
|
r37801 | |||
if not compengines: | ||||
|
r43346 | raise error.Abort( | ||
_( | ||||
|
r43347 | b'%s config option does not specify any known ' | ||
b'compression engines' | ||||
|
r43346 | ) | ||
% config, | ||||
|
r43347 | hint=_(b'usable compression engines: %s') | ||
% b', '.sorted(validnames), | ||||
|
r43346 | ) | ||
|
r37801 | |||
return compengines | ||||
|
r40056 | |||
|
r43346 | |||
|
r40056 | @attr.s | ||
class encodedresponse(object): | ||||
"""Represents response data that is already content encoded. | ||||
Wire protocol version 2 only. | ||||
Commands typically emit Python objects that are encoded and sent over the | ||||
wire. If commands emit an object of this type, the encoding step is bypassed | ||||
and the content from this object is used instead. | ||||
""" | ||||
|
r43346 | |||
|
r40056 | data = attr.ib() | ||
|
r40061 | |||
|
r43346 | |||
|
r40061 | @attr.s | ||
class alternatelocationresponse(object): | ||||
"""Represents a response available at an alternate location. | ||||
Instances are sent in place of actual response objects when the server | ||||
is sending a "content redirect" response. | ||||
Only compatible with wire protocol version 2. | ||||
""" | ||||
|
r43346 | |||
|
r40061 | url = attr.ib() | ||
mediatype = attr.ib() | ||||
size = attr.ib(default=None) | ||||
fullhashes = attr.ib(default=None) | ||||
fullhashseed = attr.ib(default=None) | ||||
serverdercerts = attr.ib(default=None) | ||||
servercadercerts = attr.ib(default=None) | ||||
|
r40364 | |||
|
r43346 | |||
|
r40364 | @attr.s | ||
class indefinitebytestringresponse(object): | ||||
"""Represents an object to be encoded to an indefinite length bytestring. | ||||
Instances are initialized from an iterable of chunks, with each chunk being | ||||
a bytes instance. | ||||
""" | ||||
|
r43346 | |||
|
r40364 | chunks = attr.ib() | ||