|
1 # commandserver.py - communicate with Mercurial's API over a pipe |
|
2 # |
|
3 # Copyright Matt Mackall <mpm@selenic.com> |
|
4 # |
|
5 # This software may be used and distributed according to the terms of the |
|
6 # GNU General Public License version 2 or any later version. |
|
7 |
|
8 from i18n import _ |
|
9 import struct |
|
10 import sys |
|
11 import dispatch, encoding, util |
|
12 |
|
13 logfile = None |
|
14 |
|
15 def log(*args): |
|
16 if not logfile: |
|
17 return |
|
18 |
|
19 for a in args: |
|
20 logfile.write(str(a)) |
|
21 |
|
22 logfile.flush() |
|
23 |
|
24 class channeledoutput(object): |
|
25 """ |
|
26 Write data from in_ to out in the following format: |
|
27 |
|
28 data length (unsigned int), |
|
29 data |
|
30 """ |
|
31 def __init__(self, in_, out, channel): |
|
32 self.in_ = in_ |
|
33 self.out = out |
|
34 self.channel = channel |
|
35 |
|
36 def write(self, data): |
|
37 if not data: |
|
38 return |
|
39 self.out.write(struct.pack('>cI', self.channel, len(data))) |
|
40 self.out.write(data) |
|
41 self.out.flush() |
|
42 |
|
43 def __getattr__(self, attr): |
|
44 if attr in ('isatty', 'fileno'): |
|
45 raise AttributeError, attr |
|
46 return getattr(self.in_, attr) |
|
47 |
|
48 class channeledinput(object): |
|
49 """ |
|
50 Read data from in_. |
|
51 |
|
52 Requests for input are written to out in the following format: |
|
53 channel identifier - 'I' for plain input, 'L' line based (1 byte) |
|
54 how many bytes to send at most (unsigned int), |
|
55 |
|
56 The client replies with: |
|
57 data length (unsigned int), 0 meaning EOF |
|
58 data |
|
59 """ |
|
60 |
|
61 maxchunksize = 4 * 1024 |
|
62 |
|
63 def __init__(self, in_, out, channel): |
|
64 self.in_ = in_ |
|
65 self.out = out |
|
66 self.channel = channel |
|
67 |
|
68 def read(self, size=-1): |
|
69 if size < 0: |
|
70 # if we need to consume all the clients input, ask for 4k chunks |
|
71 # so the pipe doesn't fill up risking a deadlock |
|
72 size = self.maxchunksize |
|
73 s = self._read(size, self.channel) |
|
74 buf = s |
|
75 while s: |
|
76 buf += s |
|
77 s = self._read(size, self.channel) |
|
78 |
|
79 return buf |
|
80 else: |
|
81 return self._read(size, self.channel) |
|
82 |
|
83 def _read(self, size, channel): |
|
84 if not size: |
|
85 return '' |
|
86 assert size > 0 |
|
87 |
|
88 # tell the client we need at most size bytes |
|
89 self.out.write(struct.pack('>cI', channel, size)) |
|
90 self.out.flush() |
|
91 |
|
92 length = self.in_.read(4) |
|
93 length = struct.unpack('>I', length)[0] |
|
94 if not length: |
|
95 return '' |
|
96 else: |
|
97 return self.in_.read(length) |
|
98 |
|
99 def readline(self, size=-1): |
|
100 if size < 0: |
|
101 size = self.maxchunksize |
|
102 s = self._read(size, 'L') |
|
103 buf = s |
|
104 # keep asking for more until there's either no more or |
|
105 # we got a full line |
|
106 while s and s[-1] != '\n': |
|
107 buf += s |
|
108 s = self._read(size, 'L') |
|
109 |
|
110 return buf |
|
111 else: |
|
112 return self._read(size, 'L') |
|
113 |
|
114 def __iter__(self): |
|
115 return self |
|
116 |
|
117 def next(self): |
|
118 l = self.readline() |
|
119 if not l: |
|
120 raise StopIteration |
|
121 return l |
|
122 |
|
123 def __getattr__(self, attr): |
|
124 if attr in ('isatty', 'fileno'): |
|
125 raise AttributeError, attr |
|
126 return getattr(self.in_, attr) |
|
127 |
|
128 class server(object): |
|
129 """ |
|
130 Listens for commands on stdin, runs them and writes the output on a channel |
|
131 based stream to stdout. |
|
132 """ |
|
133 def __init__(self, ui, repo, mode): |
|
134 self.ui = ui |
|
135 |
|
136 logpath = ui.config("cmdserver", "log", None) |
|
137 if logpath: |
|
138 global logfile |
|
139 if logpath == '-': |
|
140 # write log on a special 'd'ebug channel |
|
141 logfile = channeledoutput(sys.stdout, sys.stdout, 'd') |
|
142 else: |
|
143 logfile = open(logpath, 'a') |
|
144 |
|
145 self.repo = repo |
|
146 |
|
147 if mode == 'pipe': |
|
148 self.cerr = channeledoutput(sys.stderr, sys.stdout, 'e') |
|
149 self.cout = channeledoutput(sys.stdout, sys.stdout, 'o') |
|
150 self.cin = channeledinput(sys.stdin, sys.stdout, 'I') |
|
151 self.cresult = channeledoutput(sys.stdout, sys.stdout, 'r') |
|
152 |
|
153 self.client = sys.stdin |
|
154 else: |
|
155 raise util.Abort(_('unknown mode %s') % mode) |
|
156 |
|
157 def _read(self, size): |
|
158 data = self.client.read(size) |
|
159 |
|
160 # is the other end closed? |
|
161 if not data: |
|
162 raise EOFError() |
|
163 |
|
164 return data |
|
165 |
|
166 def runcommand(self): |
|
167 """ reads a list of \0 terminated arguments, executes |
|
168 and writes the return code to the result channel """ |
|
169 |
|
170 length = struct.unpack('>I', self._read(4))[0] |
|
171 args = self._read(length).split('\0') |
|
172 |
|
173 # copy the ui so changes to it don't persist between requests |
|
174 req = dispatch.request(args, self.ui.copy(), self.repo, self.cin, |
|
175 self.cout, self.cerr) |
|
176 |
|
177 ret = dispatch.dispatch(req) or 0 # might return None |
|
178 |
|
179 self.cresult.write(struct.pack('>i', int(ret))) |
|
180 |
|
181 def getencoding(self): |
|
182 """ writes the current encoding to the result channel """ |
|
183 self.cresult.write(encoding.encoding) |
|
184 |
|
185 def serveone(self): |
|
186 cmd = self.client.readline()[:-1] |
|
187 if cmd: |
|
188 handler = self.capabilities.get(cmd) |
|
189 if handler: |
|
190 handler(self) |
|
191 else: |
|
192 # clients are expected to check what commands are supported by |
|
193 # looking at the servers capabilities |
|
194 raise util.Abort(_('unknown command %s') % cmd) |
|
195 |
|
196 return cmd != '' |
|
197 |
|
198 capabilities = {'runcommand' : runcommand, |
|
199 'getencoding' : getencoding} |
|
200 |
|
201 def serve(self): |
|
202 self.cout.write('capabilities: %s' % ' '.join(self.capabilities.keys())) |
|
203 self.cout.write('encoding: %s' % encoding.encoding) |
|
204 |
|
205 try: |
|
206 while self.serveone(): |
|
207 pass |
|
208 except EOFError: |
|
209 # we'll get here if the client disconnected while we were reading |
|
210 # its request |
|
211 return 1 |
|
212 |
|
213 return 0 |