diff --git a/src/wormhole/server/cli.py b/src/wormhole/server/cli.py index 59dbd9a..9596dbd 100644 --- a/src/wormhole/server/cli.py +++ b/src/wormhole/server/cli.py @@ -1,4 +1,5 @@ from __future__ import print_function +import json import click from ..cli.cli import Config, _compose @@ -19,6 +20,21 @@ def server(ctx): # this is the setuptools entrypoint for bin/wormhole-server # server commands don't use ctx.obj = Config() +def _validate_websocket_protocol_options(ctx, param, value): + return list(_validate_websocket_protocol_option(option) for option in value) + +def _validate_websocket_protocol_option(option): + try: + key, value = option.split("=", 1) + except ValueError: + raise click.BadParameter("format options as OPTION=VALUE") + + try: + value = json.loads(value) + except: + raise click.BadParameter("could not parse JSON value for {}".format(key)) + + return (key, value) LaunchArgs = _compose( click.option( @@ -58,6 +74,11 @@ LaunchArgs = _compose( "--stats-json-path", default="stats.json", metavar="PATH", help="location to write the relay stats file", ), + click.option( + "--websocket-protocol-option", multiple=True, metavar="OPTION=VALUE", + callback=_validate_websocket_protocol_options, + help="a websocket server protocol option to configure", + ), ) diff --git a/src/wormhole/test/test_cli.py b/src/wormhole/test/test_cli.py index ba532bb..8e9d5bb 100644 --- a/src/wormhole/test/test_cli.py +++ b/src/wormhole/test/test_cli.py @@ -1238,3 +1238,45 @@ class Server(unittest.TestCase): relay = plugin.makeService(None) self.assertEqual('relay.sqlite', relay._db_url) self.assertEqual('stats.json', relay._stats_file) + + @mock.patch("wormhole.server.cmd_server.start_server") + def test_websocket_protocol_options(self, fake_start_server): + result = self.runner.invoke( + server, [ + 'start', + '--websocket-protocol-option=a=3', + '--websocket-protocol-option=b=true', + '--websocket-protocol-option=c=3.5', + '--websocket-protocol-option=d=["foo","bar"]', + ]) + self.assertEqual(0, result.exit_code) + cfg = fake_start_server.mock_calls[0][1][0] + self.assertEqual( + cfg.websocket_protocol_option, + [("a", 3), ("b", True), ("c", 3.5), ("d", ['foo', 'bar'])], + ) + + def test_broken_websocket_protocol_options(self): + result = self.runner.invoke( + server, [ + 'start', + '--websocket-protocol-option=a', + ]) + self.assertNotEqual(0, result.exit_code) + self.assertIn( + 'Error: Invalid value for "--websocket-protocol-option": ' + 'format options as OPTION=VALUE', + result.output, + ) + + result = self.runner.invoke( + server, [ + 'start', + '--websocket-protocol-option=a=foo', + ]) + self.assertNotEqual(0, result.exit_code) + self.assertIn( + 'Error: Invalid value for "--websocket-protocol-option": ' + 'could not parse JSON value for a', + result.output, + )