mirror of
https://github.com/facebook/sapling.git
synced 2024-10-09 08:18:15 +03:00
2112fb0fd2
Currently, the "streamres" response type is populated with a generator of chunks with compression possibly already applied. This puts the onus on commands to perform chunking and compression. Architecturally, I think this is the wrong place to perform this work. I think commands should say "here is the data" and the protocol layer should take care of encoding the final bytes to put on the wire. Additionally, upcoming commits will improve wire protocol support for compression. Having a central place for performing compression in the protocol transport layer will be easier than having to deal with compression at the commands layer. This commit refactors the "streamres" response type to accept either a generator or an object with "read." Additionally, the type now accepts a flag indicating whether the response is a "version 1 compressible" response. This basically identifies all commands currently performing compression. I could have used a special type for this, but a flag works just as well. The argument name foreshadows the introduction of wire protocol changes, hence the "v1." The code for chunking and compressing has been moved to the output generation function for each protocol transport. Some code has been inlined, resulting in the deletion of now unused methods.
135 lines
3.6 KiB
Python
135 lines
3.6 KiB
Python
# sshserver.py - ssh protocol server support for mercurial
|
|
#
|
|
# Copyright 2005-2007 Matt Mackall <mpm@selenic.com>
|
|
# Copyright 2006 Vadim Gelfer <vadim.gelfer@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
|
|
|
|
import os
|
|
import sys
|
|
|
|
from .i18n import _
|
|
from . import (
|
|
error,
|
|
hook,
|
|
util,
|
|
wireproto,
|
|
)
|
|
|
|
class sshserver(wireproto.abstractserverproto):
|
|
def __init__(self, ui, repo):
|
|
self.ui = ui
|
|
self.repo = repo
|
|
self.lock = None
|
|
self.fin = ui.fin
|
|
self.fout = ui.fout
|
|
|
|
hook.redirect(True)
|
|
ui.fout = repo.ui.fout = ui.ferr
|
|
|
|
# Prevent insertion/deletion of CRs
|
|
util.setbinary(self.fin)
|
|
util.setbinary(self.fout)
|
|
|
|
def getargs(self, args):
|
|
data = {}
|
|
keys = args.split()
|
|
for n in xrange(len(keys)):
|
|
argline = self.fin.readline()[:-1]
|
|
arg, l = argline.split()
|
|
if arg not in keys:
|
|
raise error.Abort(_("unexpected parameter %r") % arg)
|
|
if arg == '*':
|
|
star = {}
|
|
for k in xrange(int(l)):
|
|
argline = self.fin.readline()[:-1]
|
|
arg, l = argline.split()
|
|
val = self.fin.read(int(l))
|
|
star[arg] = val
|
|
data['*'] = star
|
|
else:
|
|
val = self.fin.read(int(l))
|
|
data[arg] = val
|
|
return [data[k] for k in keys]
|
|
|
|
def getarg(self, name):
|
|
return self.getargs(name)[0]
|
|
|
|
def getfile(self, fpout):
|
|
self.sendresponse('')
|
|
count = int(self.fin.readline())
|
|
while count:
|
|
fpout.write(self.fin.read(count))
|
|
count = int(self.fin.readline())
|
|
|
|
def redirect(self):
|
|
pass
|
|
|
|
def sendresponse(self, v):
|
|
self.fout.write("%d\n" % len(v))
|
|
self.fout.write(v)
|
|
self.fout.flush()
|
|
|
|
def sendstream(self, source):
|
|
write = self.fout.write
|
|
|
|
if source.reader:
|
|
gen = iter(lambda: source.reader.read(4096), '')
|
|
else:
|
|
gen = source.gen
|
|
|
|
for chunk in gen:
|
|
write(chunk)
|
|
self.fout.flush()
|
|
|
|
def sendpushresponse(self, rsp):
|
|
self.sendresponse('')
|
|
self.sendresponse(str(rsp.res))
|
|
|
|
def sendpusherror(self, rsp):
|
|
self.sendresponse(rsp.res)
|
|
|
|
def sendooberror(self, rsp):
|
|
self.ui.ferr.write('%s\n-\n' % rsp.message)
|
|
self.ui.ferr.flush()
|
|
self.fout.write('\n')
|
|
self.fout.flush()
|
|
|
|
def serve_forever(self):
|
|
try:
|
|
while self.serve_one():
|
|
pass
|
|
finally:
|
|
if self.lock is not None:
|
|
self.lock.release()
|
|
sys.exit(0)
|
|
|
|
handlers = {
|
|
str: sendresponse,
|
|
wireproto.streamres: sendstream,
|
|
wireproto.pushres: sendpushresponse,
|
|
wireproto.pusherr: sendpusherror,
|
|
wireproto.ooberror: sendooberror,
|
|
}
|
|
|
|
def serve_one(self):
|
|
cmd = self.fin.readline()[:-1]
|
|
if cmd and cmd in wireproto.commands:
|
|
rsp = wireproto.dispatch(self.repo, self, cmd)
|
|
self.handlers[rsp.__class__](self, rsp)
|
|
elif cmd:
|
|
impl = getattr(self, 'do_' + cmd, None)
|
|
if impl:
|
|
r = impl()
|
|
if r is not None:
|
|
self.sendresponse(r)
|
|
else: self.sendresponse("")
|
|
return cmd != ''
|
|
|
|
def _client(self):
|
|
client = os.environ.get('SSH_CLIENT', '').split(' ', 1)[0]
|
|
return 'remote:ssh:' + client
|