##// END OF EJS Templates
chgserver: use _readlist and _readstr...
Jun Wu -
r28158:7cc57a53 default
parent child Browse files
Show More
@@ -1,400 +1,392 b''
1 1 # chgserver.py - command server extension for cHg
2 2 #
3 3 # Copyright 2011 Yuya Nishihara <yuya@tcha.org>
4 4 #
5 5 # This software may be used and distributed according to the terms of the
6 6 # GNU General Public License version 2 or any later version.
7 7
8 8 """command server extension for cHg (EXPERIMENTAL)
9 9
10 10 'S' channel (read/write)
11 11 propagate ui.system() request to client
12 12
13 13 'attachio' command
14 14 attach client's stdio passed by sendmsg()
15 15
16 16 'chdir' command
17 17 change current directory
18 18
19 19 'getpager' command
20 20 checks if pager is enabled and which pager should be executed
21 21
22 22 'setenv' command
23 23 replace os.environ completely
24 24
25 25 'SIGHUP' signal
26 26 reload configuration files
27 27 """
28 28
29 29 from __future__ import absolute_import
30 30
31 31 import SocketServer
32 32 import errno
33 33 import os
34 34 import re
35 35 import signal
36 36 import struct
37 37 import traceback
38 38
39 39 from mercurial.i18n import _
40 40
41 41 from mercurial import (
42 42 cmdutil,
43 43 commands,
44 44 commandserver,
45 45 dispatch,
46 46 error,
47 47 osutil,
48 48 util,
49 49 )
50 50
51 51 # Note for extension authors: ONLY specify testedwith = 'internal' for
52 52 # extensions which SHIP WITH MERCURIAL. Non-mainline extensions should
53 53 # be specifying the version(s) of Mercurial they are tested with, or
54 54 # leave the attribute unspecified.
55 55 testedwith = 'internal'
56 56
57 57 _log = commandserver.log
58 58
59 59 # copied from hgext/pager.py:uisetup()
60 60 def _setuppagercmd(ui, options, cmd):
61 61 if not ui.formatted():
62 62 return
63 63
64 64 p = ui.config("pager", "pager", os.environ.get("PAGER"))
65 65 usepager = False
66 66 always = util.parsebool(options['pager'])
67 67 auto = options['pager'] == 'auto'
68 68
69 69 if not p:
70 70 pass
71 71 elif always:
72 72 usepager = True
73 73 elif not auto:
74 74 usepager = False
75 75 else:
76 76 attended = ['annotate', 'cat', 'diff', 'export', 'glog', 'log', 'qdiff']
77 77 attend = ui.configlist('pager', 'attend', attended)
78 78 ignore = ui.configlist('pager', 'ignore')
79 79 cmds, _ = cmdutil.findcmd(cmd, commands.table)
80 80
81 81 for cmd in cmds:
82 82 var = 'attend-%s' % cmd
83 83 if ui.config('pager', var):
84 84 usepager = ui.configbool('pager', var)
85 85 break
86 86 if (cmd in attend or
87 87 (cmd not in ignore and not attend)):
88 88 usepager = True
89 89 break
90 90
91 91 if usepager:
92 92 ui.setconfig('ui', 'formatted', ui.formatted(), 'pager')
93 93 ui.setconfig('ui', 'interactive', False, 'pager')
94 94 return p
95 95
96 96 _envvarre = re.compile(r'\$[a-zA-Z_]+')
97 97
98 98 def _clearenvaliases(cmdtable):
99 99 """Remove stale command aliases referencing env vars; variable expansion
100 100 is done at dispatch.addaliases()"""
101 101 for name, tab in cmdtable.items():
102 102 cmddef = tab[0]
103 103 if (isinstance(cmddef, dispatch.cmdalias) and
104 104 not cmddef.definition.startswith('!') and # shell alias
105 105 _envvarre.search(cmddef.definition)):
106 106 del cmdtable[name]
107 107
108 108 def _newchgui(srcui, csystem):
109 109 class chgui(srcui.__class__):
110 110 def __init__(self, src=None):
111 111 super(chgui, self).__init__(src)
112 112 if src:
113 113 self._csystem = getattr(src, '_csystem', csystem)
114 114 else:
115 115 self._csystem = csystem
116 116
117 117 def system(self, cmd, environ=None, cwd=None, onerr=None,
118 118 errprefix=None):
119 119 # copied from mercurial/util.py:system()
120 120 self.flush()
121 121 def py2shell(val):
122 122 if val is None or val is False:
123 123 return '0'
124 124 if val is True:
125 125 return '1'
126 126 return str(val)
127 127 env = os.environ.copy()
128 128 if environ:
129 129 env.update((k, py2shell(v)) for k, v in environ.iteritems())
130 130 env['HG'] = util.hgexecutable()
131 131 rc = self._csystem(cmd, env, cwd)
132 132 if rc and onerr:
133 133 errmsg = '%s %s' % (os.path.basename(cmd.split(None, 1)[0]),
134 134 util.explainexit(rc)[0])
135 135 if errprefix:
136 136 errmsg = '%s: %s' % (errprefix, errmsg)
137 137 raise onerr(errmsg)
138 138 return rc
139 139
140 140 return chgui(srcui)
141 141
142 142 def _renewui(srcui):
143 143 newui = srcui.__class__()
144 144 for a in ['fin', 'fout', 'ferr', 'environ']:
145 145 setattr(newui, a, getattr(srcui, a))
146 146 if util.safehasattr(srcui, '_csystem'):
147 147 newui._csystem = srcui._csystem
148 148 # stolen from tortoisehg.util.copydynamicconfig()
149 149 for section, name, value in srcui.walkconfig():
150 150 source = srcui.configsource(section, name)
151 151 if ':' in source:
152 152 # path:line
153 153 continue
154 154 if source == 'none':
155 155 # ui.configsource returns 'none' by default
156 156 source = ''
157 157 newui.setconfig(section, name, value, source)
158 158 return newui
159 159
160 160 class channeledsystem(object):
161 161 """Propagate ui.system() request in the following format:
162 162
163 163 payload length (unsigned int),
164 164 cmd, '\0',
165 165 cwd, '\0',
166 166 envkey, '=', val, '\0',
167 167 ...
168 168 envkey, '=', val
169 169
170 170 and waits:
171 171
172 172 exitcode length (unsigned int),
173 173 exitcode (int)
174 174 """
175 175 def __init__(self, in_, out, channel):
176 176 self.in_ = in_
177 177 self.out = out
178 178 self.channel = channel
179 179
180 180 def __call__(self, cmd, environ, cwd):
181 181 args = [util.quotecommand(cmd), cwd or '.']
182 182 args.extend('%s=%s' % (k, v) for k, v in environ.iteritems())
183 183 data = '\0'.join(args)
184 184 self.out.write(struct.pack('>cI', self.channel, len(data)))
185 185 self.out.write(data)
186 186 self.out.flush()
187 187
188 188 length = self.in_.read(4)
189 189 length, = struct.unpack('>I', length)
190 190 if length != 4:
191 191 raise error.Abort(_('invalid response'))
192 192 rc, = struct.unpack('>i', self.in_.read(4))
193 193 return rc
194 194
195 195 _iochannels = [
196 196 # server.ch, ui.fp, mode
197 197 ('cin', 'fin', 'rb'),
198 198 ('cout', 'fout', 'wb'),
199 199 ('cerr', 'ferr', 'wb'),
200 200 ]
201 201
202 202 class chgcmdserver(commandserver.server):
203 203 def __init__(self, ui, repo, fin, fout, sock):
204 204 super(chgcmdserver, self).__init__(
205 205 _newchgui(ui, channeledsystem(fin, fout, 'S')), repo, fin, fout)
206 206 self.clientsock = sock
207 207 self._oldios = [] # original (self.ch, ui.fp, fd) before "attachio"
208 208
209 209 def cleanup(self):
210 210 # dispatch._runcatch() does not flush outputs if exception is not
211 211 # handled by dispatch._dispatch()
212 212 self.ui.flush()
213 213 self._restoreio()
214 214
215 215 def attachio(self):
216 216 """Attach to client's stdio passed via unix domain socket; all
217 217 channels except cresult will no longer be used
218 218 """
219 219 # tell client to sendmsg() with 1-byte payload, which makes it
220 220 # distinctive from "attachio\n" command consumed by client.read()
221 221 self.clientsock.sendall(struct.pack('>cI', 'I', 1))
222 222 clientfds = osutil.recvfds(self.clientsock.fileno())
223 223 _log('received fds: %r\n' % clientfds)
224 224
225 225 ui = self.ui
226 226 ui.flush()
227 227 first = self._saveio()
228 228 for fd, (cn, fn, mode) in zip(clientfds, _iochannels):
229 229 assert fd > 0
230 230 fp = getattr(ui, fn)
231 231 os.dup2(fd, fp.fileno())
232 232 os.close(fd)
233 233 if not first:
234 234 continue
235 235 # reset buffering mode when client is first attached. as we want
236 236 # to see output immediately on pager, the mode stays unchanged
237 237 # when client re-attached. ferr is unchanged because it should
238 238 # be unbuffered no matter if it is a tty or not.
239 239 if fn == 'ferr':
240 240 newfp = fp
241 241 else:
242 242 # make it line buffered explicitly because the default is
243 243 # decided on first write(), where fout could be a pager.
244 244 if fp.isatty():
245 245 bufsize = 1 # line buffered
246 246 else:
247 247 bufsize = -1 # system default
248 248 newfp = os.fdopen(fp.fileno(), mode, bufsize)
249 249 setattr(ui, fn, newfp)
250 250 setattr(self, cn, newfp)
251 251
252 252 self.cresult.write(struct.pack('>i', len(clientfds)))
253 253
254 254 def _saveio(self):
255 255 if self._oldios:
256 256 return False
257 257 ui = self.ui
258 258 for cn, fn, _mode in _iochannels:
259 259 ch = getattr(self, cn)
260 260 fp = getattr(ui, fn)
261 261 fd = os.dup(fp.fileno())
262 262 self._oldios.append((ch, fp, fd))
263 263 return True
264 264
265 265 def _restoreio(self):
266 266 ui = self.ui
267 267 for (ch, fp, fd), (cn, fn, _mode) in zip(self._oldios, _iochannels):
268 268 newfp = getattr(ui, fn)
269 269 # close newfp while it's associated with client; otherwise it
270 270 # would be closed when newfp is deleted
271 271 if newfp is not fp:
272 272 newfp.close()
273 273 # restore original fd: fp is open again
274 274 os.dup2(fd, fp.fileno())
275 275 os.close(fd)
276 276 setattr(self, cn, ch)
277 277 setattr(ui, fn, fp)
278 278 del self._oldios[:]
279 279
280 280 def chdir(self):
281 281 """Change current directory
282 282
283 283 Note that the behavior of --cwd option is bit different from this.
284 284 It does not affect --config parameter.
285 285 """
286 length = struct.unpack('>I', self._read(4))[0]
287 if not length:
286 path = self._readstr()
287 if not path:
288 288 return
289 path = self._read(length)
290 289 _log('chdir to %r\n' % path)
291 290 os.chdir(path)
292 291
293 292 def getpager(self):
294 293 """Read cmdargs and write pager command to r-channel if enabled
295 294
296 295 If pager isn't enabled, this writes '\0' because channeledoutput
297 296 does not allow to write empty data.
298 297 """
299 length = struct.unpack('>I', self._read(4))[0]
300 if not length:
301 args = []
302 else:
303 args = self._read(length).split('\0')
298 args = self._readlist()
304 299 try:
305 300 cmd, _func, args, options, _cmdoptions = dispatch._parse(self.ui,
306 301 args)
307 302 except (error.Abort, error.AmbiguousCommand, error.CommandError,
308 303 error.UnknownCommand):
309 304 cmd = None
310 305 options = {}
311 306 if not cmd or 'pager' not in options:
312 307 self.cresult.write('\0')
313 308 return
314 309
315 310 pagercmd = _setuppagercmd(self.ui, options, cmd)
316 311 if pagercmd:
317 312 self.cresult.write(pagercmd)
318 313 else:
319 314 self.cresult.write('\0')
320 315
321 316 def setenv(self):
322 317 """Clear and update os.environ
323 318
324 319 Note that not all variables can make an effect on the running process.
325 320 """
326 length = struct.unpack('>I', self._read(4))[0]
327 if not length:
328 return
329 s = self._read(length)
321 l = self._readlist()
330 322 try:
331 newenv = dict(l.split('=', 1) for l in s.split('\0'))
323 newenv = dict(s.split('=', 1) for s in l)
332 324 except ValueError:
333 325 raise ValueError('unexpected value in setenv request')
334 326
335 327 diffkeys = set(k for k in set(os.environ.keys() + newenv.keys())
336 328 if os.environ.get(k) != newenv.get(k))
337 329 _log('change env: %r\n' % sorted(diffkeys))
338 330
339 331 os.environ.clear()
340 332 os.environ.update(newenv)
341 333
342 334 if set(['HGPLAIN', 'HGPLAINEXCEPT']) & diffkeys:
343 335 # reload config so that ui.plain() takes effect
344 336 self.ui = _renewui(self.ui)
345 337
346 338 _clearenvaliases(commands.table)
347 339
348 340 capabilities = commandserver.server.capabilities.copy()
349 341 capabilities.update({'attachio': attachio,
350 342 'chdir': chdir,
351 343 'getpager': getpager,
352 344 'setenv': setenv})
353 345
354 346 # copied from mercurial/commandserver.py
355 347 class _requesthandler(SocketServer.StreamRequestHandler):
356 348 def handle(self):
357 349 # use a different process group from the master process, making this
358 350 # process pass kernel "is_current_pgrp_orphaned" check so signals like
359 351 # SIGTSTP, SIGTTIN, SIGTTOU are not ignored.
360 352 os.setpgid(0, 0)
361 353 ui = self.server.ui
362 354 repo = self.server.repo
363 355 sv = chgcmdserver(ui, repo, self.rfile, self.wfile, self.connection)
364 356 try:
365 357 try:
366 358 sv.serve()
367 359 # handle exceptions that may be raised by command server. most of
368 360 # known exceptions are caught by dispatch.
369 361 except error.Abort as inst:
370 362 ui.warn(_('abort: %s\n') % inst)
371 363 except IOError as inst:
372 364 if inst.errno != errno.EPIPE:
373 365 raise
374 366 except KeyboardInterrupt:
375 367 pass
376 368 finally:
377 369 sv.cleanup()
378 370 except: # re-raises
379 371 # also write traceback to error channel. otherwise client cannot
380 372 # see it because it is written to server's stderr by default.
381 373 traceback.print_exc(file=sv.cerr)
382 374 raise
383 375
384 376 class chgunixservice(commandserver.unixservice):
385 377 def init(self):
386 378 # drop options set for "hg serve --cmdserver" command
387 379 self.ui.setconfig('progress', 'assume-tty', None)
388 380 signal.signal(signal.SIGHUP, self._reloadconfig)
389 381 class cls(SocketServer.ForkingMixIn, SocketServer.UnixStreamServer):
390 382 ui = self.ui
391 383 repo = self.repo
392 384 self.server = cls(self.address, _requesthandler)
393 385 # avoid writing "listening at" message to stdout before attachio
394 386 # request, which calls setvbuf()
395 387
396 388 def _reloadconfig(self, signum, frame):
397 389 self.ui = self.server.ui = _renewui(self.ui)
398 390
399 391 def uisetup(ui):
400 392 commandserver._servicemap['chgunix'] = chgunixservice
General Comments 0
You need to be logged in to leave comments. Login now