sshpeer.py
386 lines
| 12.0 KiB
| text/x-python
|
PythonLexer
/ mercurial / sshpeer.py
Peter Arrenbrecht
|
r17192 | # sshpeer.py - ssh repository proxy class for mercurial | ||
# | ||||
# Copyright 2005, 2006 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. | ||||
Gregory Szorc
|
r25975 | from __future__ import absolute_import | ||
Peter Arrenbrecht
|
r17192 | import re | ||
Gregory Szorc
|
r25975 | |||
from .i18n import _ | ||||
from . import ( | ||||
error, | ||||
Pulkit Goyal
|
r33100 | pycompat, | ||
Gregory Szorc
|
r25975 | util, | ||
wireproto, | ||||
) | ||||
Peter Arrenbrecht
|
r17192 | |||
def _serverquote(s): | ||||
Yuya Nishihara
|
r35475 | """quote a string for the remote shell ... which we assume is sh""" | ||
Matt Mackall
|
r23671 | if not s: | ||
return s | ||||
Peter Arrenbrecht
|
r17192 | if re.match('[a-zA-Z0-9@%_+=:,./-]*$', s): | ||
return s | ||||
return "'%s'" % s.replace("'", "'\\''") | ||||
Pierre-Yves David
|
r25244 | def _forwardoutput(ui, pipe): | ||
"""display all data currently available on pipe as remote output. | ||||
This is non blocking.""" | ||||
s = util.readpipe(pipe) | ||||
if s: | ||||
for l in s.splitlines(): | ||||
ui.status(_("remote: "), l, '\n') | ||||
Pierre-Yves David
|
r25421 | class doublepipe(object): | ||
"""Operate a side-channel pipe in addition of a main one | ||||
The side-channel pipe contains server output to be forwarded to the user | ||||
input. The double pipe will behave as the "main" pipe, but will ensure the | ||||
content of the "side" pipe is properly processed while we wait for blocking | ||||
call on the "main" pipe. | ||||
If large amounts of data are read from "main", the forward will cease after | ||||
the first bytes start to appear. This simplifies the implementation | ||||
without affecting actual output of sshpeer too much as we rarely issue | ||||
large read for data not yet emitted by the server. | ||||
The main pipe is expected to be a 'bufferedinputpipe' from the util module | ||||
Augie Fackler
|
r31953 | that handle all the os specific bits. This class lives in this module | ||
Mads Kiilerich
|
r26781 | because it focus on behavior specific to the ssh protocol.""" | ||
Pierre-Yves David
|
r25421 | |||
def __init__(self, ui, main, side): | ||||
self._ui = ui | ||||
self._main = main | ||||
self._side = side | ||||
def _wait(self): | ||||
"""wait until some data are available on main or side | ||||
return a pair of boolean (ismainready, issideready) | ||||
(This will only wait for data if the setup is supported by `util.poll`) | ||||
""" | ||||
Pierre-Yves David
|
r25457 | if getattr(self._main, 'hasbuffer', False): # getattr for classic pipe | ||
Pierre-Yves David
|
r25421 | return (True, True) # main has data, assume side is worth poking at. | ||
fds = [self._main.fileno(), self._side.fileno()] | ||||
try: | ||||
act = util.poll(fds) | ||||
except NotImplementedError: | ||||
# non supported yet case, assume all have data. | ||||
act = fds | ||||
return (self._main.fileno() in act, self._side.fileno() in act) | ||||
Pierre-Yves David
|
r25456 | def write(self, data): | ||
return self._call('write', data) | ||||
Pierre-Yves David
|
r25421 | def read(self, size): | ||
Augie Fackler
|
r32062 | r = self._call('read', size) | ||
if size != 0 and not r: | ||||
# We've observed a condition that indicates the | ||||
# stdout closed unexpectedly. Check stderr one | ||||
# more time and snag anything that's there before | ||||
# letting anyone know the main part of the pipe | ||||
# closed prematurely. | ||||
_forwardoutput(self._ui, self._side) | ||||
return r | ||||
Pierre-Yves David
|
r25421 | |||
def readline(self): | ||||
return self._call('readline') | ||||
Pierre-Yves David
|
r25455 | def _call(self, methname, data=None): | ||
Pierre-Yves David
|
r25421 | """call <methname> on "main", forward output of "side" while blocking | ||
""" | ||||
Pierre-Yves David
|
r25455 | # data can be '' or 0 | ||
if (data is not None and not data) or self._main.closed: | ||||
Pierre-Yves David
|
r25421 | _forwardoutput(self._ui, self._side) | ||
return '' | ||||
while True: | ||||
mainready, sideready = self._wait() | ||||
if sideready: | ||||
_forwardoutput(self._ui, self._side) | ||||
if mainready: | ||||
meth = getattr(self._main, methname) | ||||
Pierre-Yves David
|
r25455 | if data is None: | ||
Pierre-Yves David
|
r25421 | return meth() | ||
else: | ||||
Pierre-Yves David
|
r25455 | return meth(data) | ||
Pierre-Yves David
|
r25421 | |||
def close(self): | ||||
return self._main.close() | ||||
Pierre-Yves David
|
r25456 | def flush(self): | ||
return self._main.flush() | ||||
Gregory Szorc
|
r35951 | def _cleanuppipes(ui, pipei, pipeo, pipee): | ||
"""Clean up pipes used by an SSH connection.""" | ||||
if pipeo: | ||||
pipeo.close() | ||||
if pipei: | ||||
pipei.close() | ||||
if pipee: | ||||
# Try to read from the err descriptor until EOF. | ||||
try: | ||||
for l in pipee: | ||||
ui.status(_('remote: '), l) | ||||
except (IOError, ValueError): | ||||
pass | ||||
pipee.close() | ||||
Peter Arrenbrecht
|
r17192 | class sshpeer(wireproto.wirepeer): | ||
Gregory Szorc
|
r35950 | def __init__(self, ui, path, create=False, sshstate=None): | ||
Peter Arrenbrecht
|
r17192 | self._url = path | ||
Gregory Szorc
|
r33803 | self._ui = ui | ||
Gregory Szorc
|
r33763 | self._pipeo = self._pipei = self._pipee = None | ||
Peter Arrenbrecht
|
r17192 | |||
u = util.url(path, parsequery=False, parsefragment=False) | ||||
Gregory Szorc
|
r33763 | self._path = u.path or '.' | ||
Peter Arrenbrecht
|
r17192 | |||
Gregory Szorc
|
r35950 | self._validaterepo(*sshstate) | ||
Peter Arrenbrecht
|
r17192 | |||
Gregory Szorc
|
r33803 | # Begin of _basepeer interface. | ||
@util.propertycache | ||||
def ui(self): | ||||
return self._ui | ||||
Peter Arrenbrecht
|
r17192 | def url(self): | ||
return self._url | ||||
Gregory Szorc
|
r33803 | def local(self): | ||
return None | ||||
def peer(self): | ||||
return self | ||||
def canpush(self): | ||||
return True | ||||
def close(self): | ||||
pass | ||||
# End of _basepeer interface. | ||||
# Begin of _basewirecommands interface. | ||||
def capabilities(self): | ||||
return self._caps | ||||
# End of _basewirecommands interface. | ||||
Kostia Balytskyi
|
r35452 | def _validaterepo(self, sshcmd, args, remotecmd, sshenv=None): | ||
Peter Arrenbrecht
|
r17192 | # cleanup up previous run | ||
Gregory Szorc
|
r33763 | self._cleanup() | ||
Peter Arrenbrecht
|
r17192 | |||
cmd = '%s %s %s' % (sshcmd, args, | ||||
util.shellquote("%s -R %s serve --stdio" % | ||||
Gregory Szorc
|
r33763 | (_serverquote(remotecmd), _serverquote(self._path)))) | ||
Siddharth Agarwal
|
r20794 | self.ui.debug('running %s\n' % cmd) | ||
Peter Arrenbrecht
|
r17192 | cmd = util.quotecommand(cmd) | ||
Durham Goode
|
r18759 | |||
Gregory Szorc
|
r33763 | # while self._subprocess isn't used, having it allows the subprocess to | ||
Durham Goode
|
r18759 | # to clean up correctly later | ||
Pierre-Yves David
|
r25408 | # | ||
# no buffer allow the use of 'select' | ||||
# feel free to remove buffering and select usage when we ultimately | ||||
# move to threading. | ||||
Kostia Balytskyi
|
r35452 | sub = util.popen4(cmd, bufsize=0, env=sshenv) | ||
Gregory Szorc
|
r33763 | self._pipeo, self._pipei, self._pipee, self._subprocess = sub | ||
Peter Arrenbrecht
|
r17192 | |||
Gregory Szorc
|
r33763 | self._pipei = util.bufferedinputpipe(self._pipei) | ||
self._pipei = doublepipe(self.ui, self._pipei, self._pipee) | ||||
self._pipeo = doublepipe(self.ui, self._pipeo, self._pipee) | ||||
Pierre-Yves David
|
r25407 | |||
Durham Goode
|
r34194 | def badresponse(): | ||
Zuzanna Mroczek
|
r35107 | msg = _("no suitable response from remote hg") | ||
hint = self.ui.config("ui", "ssherrorhint") | ||||
self._abort(error.RepoError(msg, hint=hint)) | ||||
Durham Goode
|
r34194 | |||
try: | ||||
# skip any noise generated by remote shell | ||||
self._callstream("hello") | ||||
r = self._callstream("between", pairs=("%s-%s" % ("0"*40, "0"*40))) | ||||
except IOError: | ||||
badresponse() | ||||
Peter Arrenbrecht
|
r17192 | lines = ["", "dummy"] | ||
max_noise = 500 | ||||
while lines[-1] and max_noise: | ||||
Durham Goode
|
r34194 | try: | ||
l = r.readline() | ||||
Durham Goode
|
r34107 | self._readerr() | ||
Durham Goode
|
r34194 | if lines[-1] == "1\n" and l == "\n": | ||
break | ||||
if l: | ||||
self.ui.debug("remote: ", l) | ||||
lines.append(l) | ||||
max_noise -= 1 | ||||
except IOError: | ||||
badresponse() | ||||
Peter Arrenbrecht
|
r17192 | else: | ||
Durham Goode
|
r34194 | badresponse() | ||
Peter Arrenbrecht
|
r17192 | |||
self._caps = set() | ||||
for l in reversed(lines): | ||||
if l.startswith("capabilities:"): | ||||
self._caps.update(l[:-1].split(":")[1].split()) | ||||
break | ||||
Gregory Szorc
|
r33763 | def _readerr(self): | ||
_forwardoutput(self.ui, self._pipee) | ||||
Peter Arrenbrecht
|
r17192 | |||
def _abort(self, exception): | ||||
Gregory Szorc
|
r33763 | self._cleanup() | ||
Peter Arrenbrecht
|
r17192 | raise exception | ||
Gregory Szorc
|
r33763 | def _cleanup(self): | ||
Gregory Szorc
|
r35951 | _cleanuppipes(self.ui, self._pipei, self._pipeo, self._pipee) | ||
Peter Arrenbrecht
|
r17192 | |||
Gregory Szorc
|
r33763 | __del__ = _cleanup | ||
Peter Arrenbrecht
|
r17192 | |||
Augie Fackler
|
r28438 | def _submitbatch(self, req): | ||
Gregory Szorc
|
r29733 | rsp = self._callstream("batch", cmds=wireproto.encodebatchcmds(req)) | ||
Augie Fackler
|
r28438 | available = self._getamount() | ||
# TODO this response parsing is probably suboptimal for large | ||||
# batches with large responses. | ||||
toread = min(available, 1024) | ||||
work = rsp.read(toread) | ||||
available -= toread | ||||
chunk = work | ||||
while chunk: | ||||
while ';' in work: | ||||
one, work = work.split(';', 1) | ||||
yield wireproto.unescapearg(one) | ||||
toread = min(available, 1024) | ||||
chunk = rsp.read(toread) | ||||
available -= toread | ||||
work += chunk | ||||
yield wireproto.unescapearg(work) | ||||
Peter Arrenbrecht
|
r17192 | def _callstream(self, cmd, **args): | ||
Pulkit Goyal
|
r33100 | args = pycompat.byteskwargs(args) | ||
Boris Feld
|
r35717 | if (self.ui.debugflag | ||
and self.ui.configbool('devel', 'debug.peer-request')): | ||||
dbg = self.ui.debug | ||||
line = 'devel-peer-request: %s\n' | ||||
dbg(line % cmd) | ||||
for key, value in sorted(args.items()): | ||||
if not isinstance(value, dict): | ||||
dbg(line % ' %s: %d bytes' % (key, len(value))) | ||||
else: | ||||
for dk, dv in sorted(value.items()): | ||||
dbg(line % ' %s-%s: %d' % (key, dk, len(dv))) | ||||
Peter Arrenbrecht
|
r17192 | self.ui.debug("sending %s command\n" % cmd) | ||
Gregory Szorc
|
r33763 | self._pipeo.write("%s\n" % cmd) | ||
Peter Arrenbrecht
|
r17192 | _func, names = wireproto.commands[cmd] | ||
keys = names.split() | ||||
wireargs = {} | ||||
for k in keys: | ||||
if k == '*': | ||||
wireargs['*'] = args | ||||
break | ||||
else: | ||||
wireargs[k] = args[k] | ||||
del args[k] | ||||
for k, v in sorted(wireargs.iteritems()): | ||||
Gregory Szorc
|
r33763 | self._pipeo.write("%s %d\n" % (k, len(v))) | ||
Peter Arrenbrecht
|
r17192 | if isinstance(v, dict): | ||
for dk, dv in v.iteritems(): | ||||
Gregory Szorc
|
r33763 | self._pipeo.write("%s %d\n" % (dk, len(dv))) | ||
self._pipeo.write(dv) | ||||
Peter Arrenbrecht
|
r17192 | else: | ||
Gregory Szorc
|
r33763 | self._pipeo.write(v) | ||
self._pipeo.flush() | ||||
Peter Arrenbrecht
|
r17192 | |||
Gregory Szorc
|
r33763 | return self._pipei | ||
Peter Arrenbrecht
|
r17192 | |||
Pierre-Yves David
|
r20905 | def _callcompressable(self, cmd, **args): | ||
return self._callstream(cmd, **args) | ||||
Peter Arrenbrecht
|
r17192 | def _call(self, cmd, **args): | ||
self._callstream(cmd, **args) | ||||
return self._recv() | ||||
def _callpush(self, cmd, fp, **args): | ||||
r = self._call(cmd, **args) | ||||
if r: | ||||
return '', r | ||||
Augie Fackler
|
r29727 | for d in iter(lambda: fp.read(4096), ''): | ||
Peter Arrenbrecht
|
r17192 | self._send(d) | ||
self._send("", flush=True) | ||||
r = self._recv() | ||||
if r: | ||||
return '', r | ||||
return self._recv(), '' | ||||
Pierre-Yves David
|
r21073 | def _calltwowaystream(self, cmd, fp, **args): | ||
r = self._call(cmd, **args) | ||||
if r: | ||||
# XXX needs to be made better | ||||
liscju
|
r29389 | raise error.Abort(_('unexpected remote reply: %s') % r) | ||
Augie Fackler
|
r29727 | for d in iter(lambda: fp.read(4096), ''): | ||
Pierre-Yves David
|
r21073 | self._send(d) | ||
self._send("", flush=True) | ||||
Gregory Szorc
|
r33763 | return self._pipei | ||
Peter Arrenbrecht
|
r17192 | |||
Augie Fackler
|
r28438 | def _getamount(self): | ||
Gregory Szorc
|
r33763 | l = self._pipei.readline() | ||
Peter Arrenbrecht
|
r17192 | if l == '\n': | ||
Gregory Szorc
|
r33763 | self._readerr() | ||
Pierre-Yves David
|
r25243 | msg = _('check previous remote output') | ||
self._abort(error.OutOfBandError(hint=msg)) | ||||
Gregory Szorc
|
r33763 | self._readerr() | ||
Peter Arrenbrecht
|
r17192 | try: | ||
Augie Fackler
|
r28438 | return int(l) | ||
Peter Arrenbrecht
|
r17192 | except ValueError: | ||
self._abort(error.ResponseError(_("unexpected response:"), l)) | ||||
Augie Fackler
|
r28438 | |||
def _recv(self): | ||||
Gregory Szorc
|
r33763 | return self._pipei.read(self._getamount()) | ||
Peter Arrenbrecht
|
r17192 | |||
def _send(self, data, flush=False): | ||||
Gregory Szorc
|
r33763 | self._pipeo.write("%d\n" % len(data)) | ||
Peter Arrenbrecht
|
r17192 | if data: | ||
Gregory Szorc
|
r33763 | self._pipeo.write(data) | ||
Peter Arrenbrecht
|
r17192 | if flush: | ||
Gregory Szorc
|
r33763 | self._pipeo.flush() | ||
self._readerr() | ||||
Peter Arrenbrecht
|
r17192 | |||
Gregory Szorc
|
r35946 | def instance(ui, path, create): | ||
Gregory Szorc
|
r35949 | """Create an SSH peer. | ||
The returned object conforms to the ``wireproto.wirepeer`` interface. | ||||
""" | ||||
u = util.url(path, parsequery=False, parsefragment=False) | ||||
if u.scheme != 'ssh' or not u.host or u.path is None: | ||||
raise error.RepoError(_("couldn't parse location %s") % path) | ||||
util.checksafessh(path) | ||||
if u.passwd is not None: | ||||
raise error.RepoError(_('password in URL not supported')) | ||||
Gregory Szorc
|
r35950 | sshcmd = ui.config('ui', 'ssh') | ||
remotecmd = ui.config('ui', 'remotecmd') | ||||
sshaddenv = dict(ui.configitems('sshenv')) | ||||
sshenv = util.shellenviron(sshaddenv) | ||||
remotepath = u.path or '.' | ||||
args = util.sshargs(sshcmd, u.host, u.user, u.port) | ||||
if create: | ||||
cmd = '%s %s %s' % (sshcmd, args, | ||||
util.shellquote('%s init %s' % | ||||
(_serverquote(remotecmd), _serverquote(remotepath)))) | ||||
ui.debug('running %s\n' % cmd) | ||||
res = ui.system(cmd, blockedtag='sshpeer', environ=sshenv) | ||||
if res != 0: | ||||
raise error.RepoError(_('could not create remote repo')) | ||||
sshstate = (sshcmd, args, remotecmd, sshenv) | ||||
return sshpeer(ui, path, create=create, sshstate=sshstate) | ||||