Mercurial > hg
comparison tests/sshprotoext.py @ 35938:80a2b8ae42a1
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) |