wireproto.py
332 lines
| 10.9 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 | ||
|
r11581 | import pushkey as pushkey_ | ||
|
r11597 | # list of nodes encoding / decoding | ||
def decodelist(l, sep=' '): | ||||
return map(bin, l.split(sep)) | ||||
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')) | ||||
d = self._call("lookup", key=key) | ||||
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]) | ||
|
r11586 | except: | ||
|
r11879 | 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) | ||
branchname = urllib.unquote(branchname) | ||||
|
r11586 | # Earlier servers (1.3.x) send branch names in (their) local | ||
# charset. The best we can do is assume it's identical to our | ||||
# own local charset, in case it's not utf-8. | ||||
try: | ||||
branchname.decode('utf-8') | ||||
except UnicodeDecodeError: | ||||
branchname = encoding.fromlocal(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 | ||
except: | ||||
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()) | ||
|
r11587 | except: | ||
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", | ||||
namespace=namespace, key=key, old=old, new=new) | ||||
return bool(int(d)) | ||||
def listkeys(self, namespace): | ||||
if not self.capable('pushkey'): | ||||
return {} | ||||
d = self._call("listkeys", namespace=namespace) | ||||
r = {} | ||||
for l in d.splitlines(): | ||||
k, v = l.split('\t') | ||||
r[k.decode('string-escape')] = v.decode('string-escape') | ||||
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) | ||
return self._decompress(f) | ||||
def changegroupsubset(self, bases, heads, kind): | ||||
self.requirecap('changegroupsubset', _('look up remote changes')) | ||||
|
r11597 | bases = encodelist(bases) | ||
heads = encodelist(heads) | ||||
|
r11591 | return self._decompress(self._callstream("changegroupsubset", | ||
bases=bases, heads=heads)) | ||||
|
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()).''' | ||||
|
r11597 | ret, output = self._callpush("unbundle", cg, heads=encodelist(heads)) | ||
|
r11592 | if ret == "": | ||
raise error.ResponseError( | ||||
_('push failed:'), output) | ||||
try: | ||||
ret = int(ret) | ||||
except ValueError, err: | ||||
raise error.ResponseError( | ||||
_('push failed (unexpected response):'), ret) | ||||
for l in output.splitlines(True): | ||||
self.ui.status(_('remote: '), l) | ||||
return ret | ||||
|
r11586 | # server side | ||
|
r11625 | class streamres(object): | ||
def __init__(self, gen): | ||||
self.gen = gen | ||||
class pushres(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 | |||
|
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(): | ||||
branchname = urllib.quote(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): | ||
caps = 'lookup changegroupsubset branchmap pushkey'.split() | ||||
|
r11627 | if _allowstream(repo.ui): | ||
|
r11594 | caps.append('stream=%d' % repo.changelog.version) | ||
caps.append('unbundle=%s' % ','.join(changegroupmod.bundlepriority)) | ||||
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 | |||
|
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): | ||
|
r11581 | d = pushkey_.list(repo, namespace).items() | ||
t = '\n'.join(['%s\t%s' % (k.encode('string-escape'), | ||||
v.encode('string-escape')) for k, v in d]) | ||||
return t | ||||
|
r11583 | def lookup(repo, proto, key): | ||
|
r11581 | try: | ||
r = hex(repo.lookup(key)) | ||||
success = 1 | ||||
except Exception, inst: | ||||
r = str(inst) | ||||
success = 0 | ||||
return "%s %s\n" % (success, r) | ||||
|
r11583 | def pushkey(repo, proto, namespace, key, old, new): | ||
|
r11581 | r = pushkey_.push(repo, namespace, key, old, new) | ||
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() | ||
return their_heads == ['force'] or their_heads == heads | ||||
|
r11593 | |||
# fail early if possible | ||||
if not check_heads(): | ||||
|
r11620 | return '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 | ||||
proto.redirect() | ||||
try: | ||||
proto.getfile(fp) | ||||
lock = repo.lock() | ||||
try: | ||||
if not check_heads(): | ||||
# someone else committed/pushed/unbundled while we | ||||
# were transferring data | ||||
|
r11620 | return 'unsynced changes' | ||
|
r11593 | |||
# push can proceed | ||||
fp.seek(0) | ||||
header = fp.read(6) | ||||
if header.startswith('HG'): | ||||
if not header.startswith('HG10'): | ||||
raise ValueError('unknown bundle version') | ||||
elif header not in changegroupmod.bundletypes: | ||||
raise ValueError('unknown bundle compression type') | ||||
gen = changegroupmod.unbundle(header, fp) | ||||
try: | ||||
r = repo.addchangegroup(gen, 'serve', proto._client(), | ||||
lock=lock) | ||||
except util.Abort, inst: | ||||
sys.stderr.write("abort: %s\n" % inst) | ||||
finally: | ||||
lock.release() | ||||
|
r11625 | 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'), | ||||
|
r11581 | 'heads': (heads, ''), | ||
|
r11594 | 'hello': (hello, ''), | ||
|
r11581 | 'listkeys': (listkeys, 'namespace'), | ||
'lookup': (lookup, 'key'), | ||||
'pushkey': (pushkey, 'namespace key old new'), | ||||
|
r11585 | 'stream_out': (stream, ''), | ||
|
r11593 | 'unbundle': (unbundle, 'heads'), | ||
|
r11581 | } | ||