Mercurial > public > mercurial-scm > hg

comparison tests/sshprotoext.py @ 35938:80a2b8ae42a1

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
sshpeer: move handshake outside of sshpeer With the handshake now performed before a peer class is instantiated, we can now instantiate a different peer class depending on the results of the handshake. Our test extension had to change to cope with the new API. Because we now issue the command via raw I/O calls and don't call _callstream(), we no longer have to register the fake command. (_callstream() uses the command registration to see what args to send). Differential Revision: https://phab.mercurial-scm.org/D2034
author Gregory Szorc <gregory.szorc@gmail.com>
date Mon, 05 Feb 2018 09:14:32 -0800
parents a9cffd14aa04
children 48a3a9283f09
comparison
equal deleted inserted replaced
35937:a9cffd14aa04 35938:80a2b8ae42a1
10 10
11 from __future__ import absolute_import 11 from __future__ import absolute_import
12 12
13 from mercurial import ( 13 from mercurial import (
14 error, 14 error,
15 extensions,
15 registrar, 16 registrar,
16 sshpeer, 17 sshpeer,
17 wireproto, 18 wireproto,
18 wireprotoserver, 19 wireprotoserver,
19 ) 20 )
50 rsp = wireproto.dispatch(self._repo, self, b'between') 51 rsp = wireproto.dispatch(self._repo, self, b'between')
51 self._handlers[rsp.__class__](self, rsp) 52 self._handlers[rsp.__class__](self, rsp)
52 53
53 super(prehelloserver, self).serve_forever() 54 super(prehelloserver, self).serve_forever()
54 55
55 class extrahandshakecommandspeer(sshpeer.sshpeer): 56 def performhandshake(orig, ui, stdin, stdout, stderr):
56 """An ssh peer that sends extra commands as part of initial handshake.""" 57 """Wrapped version of sshpeer._performhandshake to send extra commands."""
57 def _validaterepo(self): 58 mode = ui.config(b'sshpeer', b'handshake-mode')
58 mode = self._ui.config(b'sshpeer', b'handshake-mode') 59 if mode == b'pre-no-args':
59 if mode == b'pre-no-args': 60 ui.debug(b'sending no-args command\n')
60 self._callstream(b'no-args') 61 stdin.write(b'no-args\n')
61 return super(extrahandshakecommandspeer, self)._validaterepo() 62 stdin.flush()
62 elif mode == b'pre-multiple-no-args': 63 return orig(ui, stdin, stdout, stderr)
63 self._callstream(b'unknown1') 64 elif mode == b'pre-multiple-no-args':
64 self._callstream(b'unknown2') 65 ui.debug(b'sending unknown1 command\n')
65 self._callstream(b'unknown3') 66 stdin.write(b'unknown1\n')
66 return super(extrahandshakecommandspeer, self)._validaterepo() 67 ui.debug(b'sending unknown2 command\n')
67 else: 68 stdin.write(b'unknown2\n')
68 raise error.ProgrammingError(b'unknown HANDSHAKECOMMANDMODE: %s' % 69 ui.debug(b'sending unknown3 command\n')
69 mode) 70 stdin.write(b'unknown3\n')
70 71 stdin.flush()
71 def registercommands(): 72 return orig(ui, stdin, stdout, stderr)
72 def dummycommand(repo, proto): 73 else:
73 raise error.ProgrammingError('this should never be called') 74 raise error.ProgrammingError(b'unknown HANDSHAKECOMMANDMODE: %s' %
74 75 mode)
75 wireproto.wireprotocommand(b'no-args', b'')(dummycommand)
76 wireproto.wireprotocommand(b'unknown1', b'')(dummycommand)
77 wireproto.wireprotocommand(b'unknown2', b'')(dummycommand)
78 wireproto.wireprotocommand(b'unknown3', b'')(dummycommand)
79 76
80 def extsetup(ui): 77 def extsetup(ui):
81 # It's easier for tests to define the server behavior via environment 78 # It's easier for tests to define the server behavior via environment
82 # variables than config options. This is because `hg serve --stdio` 79 # variables than config options. This is because `hg serve --stdio`
83 # has to be invoked with a certain form for security reasons and 80 # has to be invoked with a certain form for security reasons and
92 raise error.ProgrammingError(b'unknown server mode: %s' % servermode) 89 raise error.ProgrammingError(b'unknown server mode: %s' % servermode)
93 90
94 peermode = ui.config(b'sshpeer', b'mode') 91 peermode = ui.config(b'sshpeer', b'mode')
95 92
96 if peermode == b'extra-handshake-commands': 93 if peermode == b'extra-handshake-commands':
97 sshpeer.sshpeer = extrahandshakecommandspeer 94 extensions.wrapfunction(sshpeer, '_performhandshake', performhandshake)
98 registercommands()
99 elif peermode: 95 elif peermode:
100 raise error.ProgrammingError(b'unknown peer mode: %s' % peermode) 96 raise error.ProgrammingError(b'unknown peer mode: %s' % peermode)