##// END OF EJS Templates
sshserver: use `iter(callable, sentinel)` instead of while True...
Augie Fackler -
r29728:1a29db79 default
parent child Browse files
Show More
@@ -1,135 +1,131 b''
1 # sshserver.py - ssh protocol server support for mercurial
1 # sshserver.py - ssh protocol server support for mercurial
2 #
2 #
3 # Copyright 2005-2007 Matt Mackall <mpm@selenic.com>
3 # Copyright 2005-2007 Matt Mackall <mpm@selenic.com>
4 # Copyright 2006 Vadim Gelfer <vadim.gelfer@gmail.com>
4 # Copyright 2006 Vadim Gelfer <vadim.gelfer@gmail.com>
5 #
5 #
6 # This software may be used and distributed according to the terms of the
6 # This software may be used and distributed according to the terms of the
7 # GNU General Public License version 2 or any later version.
7 # GNU General Public License version 2 or any later version.
8
8
9 from __future__ import absolute_import
9 from __future__ import absolute_import
10
10
11 import os
11 import os
12 import sys
12 import sys
13
13
14 from .i18n import _
14 from .i18n import _
15 from . import (
15 from . import (
16 error,
16 error,
17 hook,
17 hook,
18 util,
18 util,
19 wireproto,
19 wireproto,
20 )
20 )
21
21
22 class sshserver(wireproto.abstractserverproto):
22 class sshserver(wireproto.abstractserverproto):
23 def __init__(self, ui, repo):
23 def __init__(self, ui, repo):
24 self.ui = ui
24 self.ui = ui
25 self.repo = repo
25 self.repo = repo
26 self.lock = None
26 self.lock = None
27 self.fin = ui.fin
27 self.fin = ui.fin
28 self.fout = ui.fout
28 self.fout = ui.fout
29
29
30 hook.redirect(True)
30 hook.redirect(True)
31 ui.fout = repo.ui.fout = ui.ferr
31 ui.fout = repo.ui.fout = ui.ferr
32
32
33 # Prevent insertion/deletion of CRs
33 # Prevent insertion/deletion of CRs
34 util.setbinary(self.fin)
34 util.setbinary(self.fin)
35 util.setbinary(self.fout)
35 util.setbinary(self.fout)
36
36
37 def getargs(self, args):
37 def getargs(self, args):
38 data = {}
38 data = {}
39 keys = args.split()
39 keys = args.split()
40 for n in xrange(len(keys)):
40 for n in xrange(len(keys)):
41 argline = self.fin.readline()[:-1]
41 argline = self.fin.readline()[:-1]
42 arg, l = argline.split()
42 arg, l = argline.split()
43 if arg not in keys:
43 if arg not in keys:
44 raise error.Abort(_("unexpected parameter %r") % arg)
44 raise error.Abort(_("unexpected parameter %r") % arg)
45 if arg == '*':
45 if arg == '*':
46 star = {}
46 star = {}
47 for k in xrange(int(l)):
47 for k in xrange(int(l)):
48 argline = self.fin.readline()[:-1]
48 argline = self.fin.readline()[:-1]
49 arg, l = argline.split()
49 arg, l = argline.split()
50 val = self.fin.read(int(l))
50 val = self.fin.read(int(l))
51 star[arg] = val
51 star[arg] = val
52 data['*'] = star
52 data['*'] = star
53 else:
53 else:
54 val = self.fin.read(int(l))
54 val = self.fin.read(int(l))
55 data[arg] = val
55 data[arg] = val
56 return [data[k] for k in keys]
56 return [data[k] for k in keys]
57
57
58 def getarg(self, name):
58 def getarg(self, name):
59 return self.getargs(name)[0]
59 return self.getargs(name)[0]
60
60
61 def getfile(self, fpout):
61 def getfile(self, fpout):
62 self.sendresponse('')
62 self.sendresponse('')
63 count = int(self.fin.readline())
63 count = int(self.fin.readline())
64 while count:
64 while count:
65 fpout.write(self.fin.read(count))
65 fpout.write(self.fin.read(count))
66 count = int(self.fin.readline())
66 count = int(self.fin.readline())
67
67
68 def redirect(self):
68 def redirect(self):
69 pass
69 pass
70
70
71 def groupchunks(self, changegroup):
71 def groupchunks(self, changegroup):
72 while True:
72 return iter(lambda: changegroup.read(4096), '')
73 d = changegroup.read(4096)
74 if not d:
75 break
76 yield d
77
73
78 def sendresponse(self, v):
74 def sendresponse(self, v):
79 self.fout.write("%d\n" % len(v))
75 self.fout.write("%d\n" % len(v))
80 self.fout.write(v)
76 self.fout.write(v)
81 self.fout.flush()
77 self.fout.flush()
82
78
83 def sendstream(self, source):
79 def sendstream(self, source):
84 write = self.fout.write
80 write = self.fout.write
85 for chunk in source.gen:
81 for chunk in source.gen:
86 write(chunk)
82 write(chunk)
87 self.fout.flush()
83 self.fout.flush()
88
84
89 def sendpushresponse(self, rsp):
85 def sendpushresponse(self, rsp):
90 self.sendresponse('')
86 self.sendresponse('')
91 self.sendresponse(str(rsp.res))
87 self.sendresponse(str(rsp.res))
92
88
93 def sendpusherror(self, rsp):
89 def sendpusherror(self, rsp):
94 self.sendresponse(rsp.res)
90 self.sendresponse(rsp.res)
95
91
96 def sendooberror(self, rsp):
92 def sendooberror(self, rsp):
97 self.ui.ferr.write('%s\n-\n' % rsp.message)
93 self.ui.ferr.write('%s\n-\n' % rsp.message)
98 self.ui.ferr.flush()
94 self.ui.ferr.flush()
99 self.fout.write('\n')
95 self.fout.write('\n')
100 self.fout.flush()
96 self.fout.flush()
101
97
102 def serve_forever(self):
98 def serve_forever(self):
103 try:
99 try:
104 while self.serve_one():
100 while self.serve_one():
105 pass
101 pass
106 finally:
102 finally:
107 if self.lock is not None:
103 if self.lock is not None:
108 self.lock.release()
104 self.lock.release()
109 sys.exit(0)
105 sys.exit(0)
110
106
111 handlers = {
107 handlers = {
112 str: sendresponse,
108 str: sendresponse,
113 wireproto.streamres: sendstream,
109 wireproto.streamres: sendstream,
114 wireproto.pushres: sendpushresponse,
110 wireproto.pushres: sendpushresponse,
115 wireproto.pusherr: sendpusherror,
111 wireproto.pusherr: sendpusherror,
116 wireproto.ooberror: sendooberror,
112 wireproto.ooberror: sendooberror,
117 }
113 }
118
114
119 def serve_one(self):
115 def serve_one(self):
120 cmd = self.fin.readline()[:-1]
116 cmd = self.fin.readline()[:-1]
121 if cmd and cmd in wireproto.commands:
117 if cmd and cmd in wireproto.commands:
122 rsp = wireproto.dispatch(self.repo, self, cmd)
118 rsp = wireproto.dispatch(self.repo, self, cmd)
123 self.handlers[rsp.__class__](self, rsp)
119 self.handlers[rsp.__class__](self, rsp)
124 elif cmd:
120 elif cmd:
125 impl = getattr(self, 'do_' + cmd, None)
121 impl = getattr(self, 'do_' + cmd, None)
126 if impl:
122 if impl:
127 r = impl()
123 r = impl()
128 if r is not None:
124 if r is not None:
129 self.sendresponse(r)
125 self.sendresponse(r)
130 else: self.sendresponse("")
126 else: self.sendresponse("")
131 return cmd != ''
127 return cmd != ''
132
128
133 def _client(self):
129 def _client(self):
134 client = os.environ.get('SSH_CLIENT', '').split(' ', 1)[0]
130 client = os.environ.get('SSH_CLIENT', '').split(' ', 1)[0]
135 return 'remote:ssh:' + client
131 return 'remote:ssh:' + client
General Comments 0
You need to be logged in to leave comments. Login now