wireproto.py
421 lines
| 14.1 KiB
| text/x-python
|
PythonLexer
/ mercurial / wireproto.py
|
r11581 | # wireproto.py - generic wire protocol support functions | ||
# | ||||
# Copyright 2005-2010 Matt Mackall <mpm@selenic.com> | ||||
# | ||||
# This software may be used and distributed according to the terms of the | ||||
# GNU General Public License version 2 or any later version. | ||||
|
r11879 | import urllib, tempfile, os, sys | ||
|
r11581 | from i18n import _ | ||
from node import bin, hex | ||||
|
r11593 | import changegroup as changegroupmod | ||
|
r11627 | import repo, error, encoding, util, store | ||
|
r12085 | import pushkey as pushkeymod | ||
|
r11581 | |||
|
r11597 | # list of nodes encoding / decoding | ||
def decodelist(l, sep=' '): | ||||
|
r13722 | if l: | ||
return map(bin, l.split(sep)) | ||||
return [] | ||||
|
r11597 | |||
def encodelist(l, sep=' '): | ||||
return sep.join(map(hex, l)) | ||||
|
r11586 | # client side | ||
class wirerepository(repo.repository): | ||||
def lookup(self, key): | ||||
self.requirecap('lookup', _('look up remote revision')) | ||||
|
r13049 | d = self._call("lookup", key=encoding.fromlocal(key)) | ||
|
r11586 | success, data = d[:-1].split(" ", 1) | ||
if int(success): | ||||
return bin(data) | ||||
self._abort(error.RepoError(data)) | ||||
def heads(self): | ||||
d = self._call("heads") | ||||
try: | ||||
|
r11597 | return decodelist(d[:-1]) | ||
|
r13726 | except ValueError: | ||
|
r11879 | self._abort(error.ResponseError(_("unexpected response:"), d)) | ||
|
r11586 | |||
|
r13723 | def known(self, nodes): | ||
n = encodelist(nodes) | ||||
d = self._call("known", nodes=n) | ||||
try: | ||||
return [bool(int(f)) for f in d] | ||||
|
r13726 | except ValueError: | ||
|
r13723 | self._abort(error.ResponseError(_("unexpected response:"), d)) | ||
|
r11586 | def branchmap(self): | ||
d = self._call("branchmap") | ||||
try: | ||||
branchmap = {} | ||||
for branchpart in d.splitlines(): | ||||
|
r11597 | branchname, branchheads = branchpart.split(' ', 1) | ||
|
r13047 | branchname = encoding.tolocal(urllib.unquote(branchname)) | ||
|
r11597 | branchheads = decodelist(branchheads) | ||
|
r11586 | branchmap[branchname] = branchheads | ||
return branchmap | ||||
except TypeError: | ||||
self._abort(error.ResponseError(_("unexpected response:"), d)) | ||||
def branches(self, nodes): | ||||
|
r11597 | n = encodelist(nodes) | ||
|
r11586 | d = self._call("branches", nodes=n) | ||
try: | ||||
|
r11597 | br = [tuple(decodelist(b)) for b in d.splitlines()] | ||
|
r11586 | return br | ||
|
r13726 | except ValueError: | ||
|
r11586 | self._abort(error.ResponseError(_("unexpected response:"), d)) | ||
def between(self, pairs): | ||||
|
r11587 | batch = 8 # avoid giant requests | ||
r = [] | ||||
for i in xrange(0, len(pairs), batch): | ||||
|
r11597 | n = " ".join([encodelist(p, '-') for p in pairs[i:i + batch]]) | ||
|
r11587 | d = self._call("between", pairs=n) | ||
try: | ||||
|
r11597 | r.extend(l and decodelist(l) or [] for l in d.splitlines()) | ||
|
r13726 | except ValueError: | ||
|
r11587 | self._abort(error.ResponseError(_("unexpected response:"), d)) | ||
return r | ||||
|
r11586 | |||
def pushkey(self, namespace, key, old, new): | ||||
if not self.capable('pushkey'): | ||||
return False | ||||
d = self._call("pushkey", | ||||
|
r13050 | namespace=encoding.fromlocal(namespace), | ||
key=encoding.fromlocal(key), | ||||
old=encoding.fromlocal(old), | ||||
new=encoding.fromlocal(new)) | ||||
|
r13450 | try: | ||
d = bool(int(d)) | ||||
except ValueError: | ||||
raise error.ResponseError( | ||||
_('push failed (unexpected response):'), d) | ||||
return d | ||||
|
r11586 | |||
def listkeys(self, namespace): | ||||
if not self.capable('pushkey'): | ||||
return {} | ||||
|
r13050 | d = self._call("listkeys", namespace=encoding.fromlocal(namespace)) | ||
|
r11586 | r = {} | ||
for l in d.splitlines(): | ||||
k, v = l.split('\t') | ||||
|
r13050 | r[encoding.tolocal(k)] = encoding.tolocal(v) | ||
|
r11586 | return r | ||
|
r11588 | def stream_out(self): | ||
return self._callstream('stream_out') | ||||
|
r11591 | def changegroup(self, nodes, kind): | ||
|
r11597 | n = encodelist(nodes) | ||
|
r11591 | f = self._callstream("changegroup", roots=n) | ||
|
r12337 | return changegroupmod.unbundle10(self._decompress(f), 'UN') | ||
|
r11591 | |||
def changegroupsubset(self, bases, heads, kind): | ||||
self.requirecap('changegroupsubset', _('look up remote changes')) | ||||
|
r11597 | bases = encodelist(bases) | ||
heads = encodelist(heads) | ||||
|
r12337 | f = self._callstream("changegroupsubset", | ||
bases=bases, heads=heads) | ||||
return changegroupmod.unbundle10(self._decompress(f), 'UN') | ||||
|
r11591 | |||
|
r13741 | def getbundle(self, source, heads=None, common=None): | ||
self.requirecap('getbundle', _('look up remote changes')) | ||||
opts = {} | ||||
if heads is not None: | ||||
opts['heads'] = encodelist(heads) | ||||
if common is not None: | ||||
opts['common'] = encodelist(common) | ||||
f = self._callstream("getbundle", **opts) | ||||
return changegroupmod.unbundle10(self._decompress(f), 'UN') | ||||
|
r11592 | def unbundle(self, cg, heads, source): | ||
'''Send cg (a readable file-like object representing the | ||||
changegroup to push, typically a chunkbuffer object) to the | ||||
remote server as a bundle. Return an integer indicating the | ||||
result of the push (see localrepository.addchangegroup()).''' | ||||
|
r13942 | if self.capable('unbundlehash'): | ||
heads = encodelist(['hashed', | ||||
util.sha1(''.join(sorted(heads))).digest()]) | ||||
else: | ||||
heads = encodelist(heads) | ||||
ret, output = self._callpush("unbundle", cg, heads=heads) | ||||
|
r11592 | if ret == "": | ||
raise error.ResponseError( | ||||
_('push failed:'), output) | ||||
try: | ||||
ret = int(ret) | ||||
|
r12063 | except ValueError: | ||
|
r11592 | raise error.ResponseError( | ||
_('push failed (unexpected response):'), ret) | ||||
for l in output.splitlines(True): | ||||
self.ui.status(_('remote: '), l) | ||||
return ret | ||||
|
r14048 | def debugwireargs(self, one, two, three=None, four=None, five=None): | ||
|
r13720 | # don't pass optional arguments left at their default value | ||
opts = {} | ||||
if three is not None: | ||||
opts['three'] = three | ||||
if four is not None: | ||||
opts['four'] = four | ||||
return self._call('debugwireargs', one=one, two=two, **opts) | ||||
|
r11586 | # server side | ||
|
r11625 | class streamres(object): | ||
def __init__(self, gen): | ||||
self.gen = gen | ||||
class pushres(object): | ||||
def __init__(self, res): | ||||
self.res = res | ||||
|
r12703 | class pusherr(object): | ||
def __init__(self, res): | ||||
self.res = res | ||||
|
r11581 | def dispatch(repo, proto, command): | ||
func, spec = commands[command] | ||||
args = proto.getargs(spec) | ||||
|
r11625 | return func(repo, proto, *args) | ||
|
r11581 | |||
|
r13721 | def options(cmd, keys, others): | ||
opts = {} | ||||
for k in keys: | ||||
if k in others: | ||||
opts[k] = others[k] | ||||
del others[k] | ||||
if others: | ||||
sys.stderr.write("abort: %s got unexpected arguments %s\n" | ||||
% (cmd, ",".join(others))) | ||||
return opts | ||||
|
r11583 | def between(repo, proto, pairs): | ||
|
r11597 | pairs = [decodelist(p, '-') for p in pairs.split(" ")] | ||
|
r11581 | r = [] | ||
for b in repo.between(pairs): | ||||
|
r11597 | r.append(encodelist(b) + "\n") | ||
|
r11581 | return "".join(r) | ||
|
r11583 | def branchmap(repo, proto): | ||
|
r11581 | branchmap = repo.branchmap() | ||
heads = [] | ||||
for branch, nodes in branchmap.iteritems(): | ||||
|
r13047 | branchname = urllib.quote(encoding.fromlocal(branch)) | ||
|
r11597 | branchnodes = encodelist(nodes) | ||
heads.append('%s %s' % (branchname, branchnodes)) | ||||
|
r11581 | return '\n'.join(heads) | ||
|
r11583 | def branches(repo, proto, nodes): | ||
|
r11597 | nodes = decodelist(nodes) | ||
|
r11581 | r = [] | ||
for b in repo.branches(nodes): | ||||
|
r11597 | r.append(encodelist(b) + "\n") | ||
|
r11581 | return "".join(r) | ||
|
r11594 | def capabilities(repo, proto): | ||
|
r13942 | caps = ('lookup changegroupsubset branchmap pushkey known getbundle ' | ||
'unbundlehash').split() | ||||
|
r11627 | if _allowstream(repo.ui): | ||
|
r12296 | requiredformats = repo.requirements & repo.supportedformats | ||
# if our local revlogs are just revlogv1, add 'stream' cap | ||||
if not requiredformats - set(('revlogv1',)): | ||||
caps.append('stream') | ||||
# otherwise, add 'streamreqs' detailing our local revlog format | ||||
else: | ||||
caps.append('streamreqs=%s' % ','.join(requiredformats)) | ||||
|
r11594 | caps.append('unbundle=%s' % ','.join(changegroupmod.bundlepriority)) | ||
|
r14093 | caps.append('httpheader=1024') | ||
|
r11594 | return ' '.join(caps) | ||
|
r11584 | def changegroup(repo, proto, roots): | ||
|
r11597 | nodes = decodelist(roots) | ||
|
r11584 | cg = repo.changegroup(nodes, 'serve') | ||
|
r11625 | return streamres(proto.groupchunks(cg)) | ||
|
r11584 | |||
def changegroupsubset(repo, proto, bases, heads): | ||||
|
r11597 | bases = decodelist(bases) | ||
heads = decodelist(heads) | ||||
|
r11584 | cg = repo.changegroupsubset(bases, heads, 'serve') | ||
|
r11625 | return streamres(proto.groupchunks(cg)) | ||
|
r11584 | |||
|
r13721 | def debugwireargs(repo, proto, one, two, others): | ||
# only accept optional args from the known set | ||||
opts = options('debugwireargs', ['three', 'four'], others) | ||||
return repo.debugwireargs(one, two, **opts) | ||||
|
r13720 | |||
|
r13741 | def getbundle(repo, proto, others): | ||
opts = options('getbundle', ['heads', 'common'], others) | ||||
for k, v in opts.iteritems(): | ||||
opts[k] = decodelist(v) | ||||
cg = repo.getbundle('serve', **opts) | ||||
return streamres(proto.groupchunks(cg)) | ||||
|
r11583 | def heads(repo, proto): | ||
|
r11581 | h = repo.heads() | ||
|
r11597 | return encodelist(h) + "\n" | ||
|
r11581 | |||
|
r11594 | def hello(repo, proto): | ||
'''the hello command returns a set of lines describing various | ||||
interesting things about the server, in an RFC822-like format. | ||||
Currently the only one defined is "capabilities", which | ||||
consists of a line in the form: | ||||
capabilities: space separated list of tokens | ||||
''' | ||||
return "capabilities: %s\n" % (capabilities(repo, proto)) | ||||
|
r11583 | def listkeys(repo, proto, namespace): | ||
|
r13050 | d = pushkeymod.list(repo, encoding.tolocal(namespace)).items() | ||
t = '\n'.join(['%s\t%s' % (encoding.fromlocal(k), encoding.fromlocal(v)) | ||||
for k, v in d]) | ||||
|
r11581 | return t | ||
|
r11583 | def lookup(repo, proto, key): | ||
|
r11581 | try: | ||
|
r13049 | r = hex(repo.lookup(encoding.tolocal(key))) | ||
|
r11581 | success = 1 | ||
except Exception, inst: | ||||
r = str(inst) | ||||
success = 0 | ||||
return "%s %s\n" % (success, r) | ||||
|
r13723 | def known(repo, proto, nodes): | ||
return ''.join(b and "1" or "0" for b in repo.known(decodelist(nodes))) | ||||
|
r11583 | def pushkey(repo, proto, namespace, key, old, new): | ||
|
r13050 | # compatibility with pre-1.8 clients which were accidentally | ||
# sending raw binary nodes rather than utf-8-encoded hex | ||||
if len(new) == 20 and new.encode('string-escape') != new: | ||||
# looks like it could be a binary node | ||||
try: | ||||
|
r14064 | new.decode('utf-8') | ||
|
r13050 | new = encoding.tolocal(new) # but cleanly decodes as UTF-8 | ||
except UnicodeDecodeError: | ||||
pass # binary, leave unmodified | ||||
else: | ||||
new = encoding.tolocal(new) # normal path | ||||
r = pushkeymod.push(repo, | ||||
encoding.tolocal(namespace), encoding.tolocal(key), | ||||
encoding.tolocal(old), new) | ||||
|
r11581 | return '%s\n' % int(r) | ||
|
r11627 | def _allowstream(ui): | ||
return ui.configbool('server', 'uncompressed', True, untrusted=True) | ||||
|
r11585 | def stream(repo, proto): | ||
|
r11627 | '''If the server supports streaming clone, it advertises the "stream" | ||
capability with a value representing the version and flags of the repo | ||||
it is serving. Client checks to see if it understands the format. | ||||
The format is simple: the server writes out a line with the amount | ||||
of files, then the total amount of bytes to be transfered (separated | ||||
by a space). Then, for each file, the server first writes the filename | ||||
and filesize (separated by the null character), then the file contents. | ||||
''' | ||||
if not _allowstream(repo.ui): | ||||
return '1\n' | ||||
entries = [] | ||||
total_bytes = 0 | ||||
try: | ||||
# get consistent snapshot of repo, lock during scan | ||||
lock = repo.lock() | ||||
try: | ||||
repo.ui.debug('scanning\n') | ||||
for name, ename, size in repo.store.walk(): | ||||
entries.append((name, size)) | ||||
total_bytes += size | ||||
finally: | ||||
lock.release() | ||||
except error.LockError: | ||||
return '2\n' # error: 2 | ||||
def streamer(repo, entries, total): | ||||
'''stream out all metadata files in repository.''' | ||||
yield '0\n' # success | ||||
repo.ui.debug('%d files, %d bytes to transfer\n' % | ||||
(len(entries), total_bytes)) | ||||
yield '%d %d\n' % (len(entries), total_bytes) | ||||
for name, size in entries: | ||||
repo.ui.debug('sending %s (%d bytes)\n' % (name, size)) | ||||
# partially encode name over the wire for backwards compat | ||||
yield '%s\0%d\n' % (store.encodedir(name), size) | ||||
for chunk in util.filechunkiter(repo.sopener(name), limit=size): | ||||
yield chunk | ||||
return streamres(streamer(repo, entries, total_bytes)) | ||||
|
r11585 | |||
|
r11593 | def unbundle(repo, proto, heads): | ||
|
r11597 | their_heads = decodelist(heads) | ||
|
r11593 | |||
def check_heads(): | ||||
|
r11597 | heads = repo.heads() | ||
|
r13942 | heads_hash = util.sha1(''.join(sorted(heads))).digest() | ||
return (their_heads == ['force'] or their_heads == heads or | ||||
their_heads == ['hashed', heads_hash]) | ||||
|
r11593 | |||
|
r12702 | proto.redirect() | ||
|
r11593 | # fail early if possible | ||
if not check_heads(): | ||||
|
r12703 | return pusherr('unsynced changes') | ||
|
r11593 | |||
# write bundle data to temporary file because it can be big | ||||
fd, tempname = tempfile.mkstemp(prefix='hg-unbundle-') | ||||
fp = os.fdopen(fd, 'wb+') | ||||
r = 0 | ||||
try: | ||||
proto.getfile(fp) | ||||
lock = repo.lock() | ||||
try: | ||||
if not check_heads(): | ||||
# someone else committed/pushed/unbundled while we | ||||
# were transferring data | ||||
|
r12703 | return pusherr('unsynced changes') | ||
|
r11593 | |||
# push can proceed | ||||
fp.seek(0) | ||||
|
r12042 | gen = changegroupmod.readbundle(fp, None) | ||
|
r11593 | |||
try: | ||||
r = repo.addchangegroup(gen, 'serve', proto._client(), | ||||
lock=lock) | ||||
except util.Abort, inst: | ||||
sys.stderr.write("abort: %s\n" % inst) | ||||
finally: | ||||
lock.release() | ||||
|
r12701 | return pushres(r) | ||
|
r11593 | |||
finally: | ||||
fp.close() | ||||
os.unlink(tempname) | ||||
|
r11581 | commands = { | ||
'between': (between, 'pairs'), | ||||
'branchmap': (branchmap, ''), | ||||
'branches': (branches, 'nodes'), | ||||
|
r11594 | 'capabilities': (capabilities, ''), | ||
|
r11584 | 'changegroup': (changegroup, 'roots'), | ||
'changegroupsubset': (changegroupsubset, 'bases heads'), | ||||
|
r13721 | 'debugwireargs': (debugwireargs, 'one two *'), | ||
|
r13741 | 'getbundle': (getbundle, '*'), | ||
|
r11581 | 'heads': (heads, ''), | ||
|
r11594 | 'hello': (hello, ''), | ||
|
r13723 | 'known': (known, 'nodes'), | ||
|
r11581 | 'listkeys': (listkeys, 'namespace'), | ||
'lookup': (lookup, 'key'), | ||||
'pushkey': (pushkey, 'namespace key old new'), | ||||
|
r11585 | 'stream_out': (stream, ''), | ||
|
r11593 | 'unbundle': (unbundle, 'heads'), | ||
|
r11581 | } | ||