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) |