From 5989d31ebf1a84bc6c6bdb7b32b4a9ea5437f36d Mon Sep 17 00:00:00 2001 From: Insaf Nureev Date: Thu, 8 Aug 2024 13:05:45 +0300 Subject: [PATCH] Add namespace value to TokenRedis source --- tests/test_token_plugins.py | 91 +++++++++++++++++++++++++++++++++---- websockify/token_plugins.py | 31 ++++++++++--- 2 files changed, 108 insertions(+), 14 deletions(-) diff --git a/tests/test_token_plugins.py b/tests/test_token_plugins.py index f09ae8a..a63e11b 100644 --- a/tests/test_token_plugins.py +++ b/tests/test_token_plugins.py @@ -267,6 +267,42 @@ class TokenRedisTestCase(unittest.TestCase): instance.get.assert_called_once_with('testhost') self.assertIsNone(result) + @patch('redis.Redis') + def test_token_without_namespace(self, mock_redis): + plugin = TokenRedis('127.0.0.1:1234') + token = 'testhost' + + def mock_redis_get(key): + self.assertEqual(key, token) + return b'remote_host:remote_port' + + instance = mock_redis.return_value + instance.get = mock_redis_get + + result = plugin.lookup(token) + + self.assertIsNotNone(result) + self.assertEqual(result[0], 'remote_host') + self.assertEqual(result[1], 'remote_port') + + @patch('redis.Redis') + def test_token_with_namespace(self, mock_redis): + plugin = TokenRedis('127.0.0.1:1234:::namespace') + token = 'testhost' + + def mock_redis_get(key): + self.assertEqual(key, "namespace:" + token) + return b'remote_host:remote_port' + + instance = mock_redis.return_value + instance.get = mock_redis_get + + result = plugin.lookup(token) + + self.assertIsNotNone(result) + self.assertEqual(result[0], 'remote_host') + self.assertEqual(result[1], 'remote_port') + def test_src_only_host(self): plugin = TokenRedis('127.0.0.1') @@ -274,6 +310,7 @@ class TokenRedisTestCase(unittest.TestCase): self.assertEqual(plugin._port, 6379) self.assertEqual(plugin._db, 0) self.assertEqual(plugin._password, None) + self.assertEqual(plugin._namespace, "") def test_src_with_host_port(self): plugin = TokenRedis('127.0.0.1:1234') @@ -282,6 +319,7 @@ class TokenRedisTestCase(unittest.TestCase): self.assertEqual(plugin._port, 1234) self.assertEqual(plugin._db, 0) self.assertEqual(plugin._password, None) + self.assertEqual(plugin._namespace, "") def test_src_with_host_port_db(self): plugin = TokenRedis('127.0.0.1:1234:2') @@ -290,6 +328,7 @@ class TokenRedisTestCase(unittest.TestCase): self.assertEqual(plugin._port, 1234) self.assertEqual(plugin._db, 2) self.assertEqual(plugin._password, None) + self.assertEqual(plugin._namespace, "") def test_src_with_host_port_db_pass(self): plugin = TokenRedis('127.0.0.1:1234:2:verysecret') @@ -298,67 +337,103 @@ class TokenRedisTestCase(unittest.TestCase): self.assertEqual(plugin._port, 1234) self.assertEqual(plugin._db, 2) self.assertEqual(plugin._password, 'verysecret') + self.assertEqual(plugin._namespace, "") - def test_src_with_host_empty_port_empty_db_pass(self): + def test_src_with_host_port_db_pass_namespace(self): + plugin = TokenRedis('127.0.0.1:1234:2:verysecret:namespace') + + self.assertEqual(plugin._server, '127.0.0.1') + self.assertEqual(plugin._port, 1234) + self.assertEqual(plugin._db, 2) + self.assertEqual(plugin._password, 'verysecret') + self.assertEqual(plugin._namespace, "namespace:") + + def test_src_with_host_empty_port_empty_db_pass_no_namespace(self): plugin = TokenRedis('127.0.0.1:::verysecret') self.assertEqual(plugin._server, '127.0.0.1') self.assertEqual(plugin._port, 6379) self.assertEqual(plugin._db, 0) self.assertEqual(plugin._password, 'verysecret') + self.assertEqual(plugin._namespace, "") - def test_src_with_host_empty_port_empty_db_empty_pass(self): + def test_src_with_host_empty_port_empty_db_empty_pass_empty_namespace(self): + plugin = TokenRedis('127.0.0.1::::') + + self.assertEqual(plugin._server, '127.0.0.1') + self.assertEqual(plugin._port, 6379) + self.assertEqual(plugin._db, 0) + self.assertEqual(plugin._password, None) + self.assertEqual(plugin._namespace, "") + + def test_src_with_host_empty_port_empty_db_empty_pass_no_namespace(self): plugin = TokenRedis('127.0.0.1:::') self.assertEqual(plugin._server, '127.0.0.1') self.assertEqual(plugin._port, 6379) self.assertEqual(plugin._db, 0) self.assertEqual(plugin._password, None) + self.assertEqual(plugin._namespace, "") - def test_src_with_host_empty_port_empty_db_no_pass(self): + def test_src_with_host_empty_port_empty_db_no_pass_no_namespace(self): plugin = TokenRedis('127.0.0.1::') self.assertEqual(plugin._server, '127.0.0.1') self.assertEqual(plugin._port, 6379) self.assertEqual(plugin._db, 0) self.assertEqual(plugin._password, None) + self.assertEqual(plugin._namespace, "") - def test_src_with_host_empty_port_no_db_no_pass(self): + def test_src_with_host_empty_port_no_db_no_pass_no_namespace(self): plugin = TokenRedis('127.0.0.1:') self.assertEqual(plugin._server, '127.0.0.1') self.assertEqual(plugin._port, 6379) self.assertEqual(plugin._db, 0) self.assertEqual(plugin._password, None) + self.assertEqual(plugin._namespace, "") - def test_src_with_host_empty_port_db_no_pass(self): + def test_src_with_host_empty_port_empty_db_empty_pass_namespace(self): + plugin = TokenRedis('127.0.0.1::::namespace') + + self.assertEqual(plugin._server, '127.0.0.1') + self.assertEqual(plugin._port, 6379) + self.assertEqual(plugin._db, 0) + self.assertEqual(plugin._password, None) + self.assertEqual(plugin._namespace, "namespace:") + + def test_src_with_host_empty_port_db_no_pass_no_namespace(self): plugin = TokenRedis('127.0.0.1::2') self.assertEqual(plugin._server, '127.0.0.1') self.assertEqual(plugin._port, 6379) self.assertEqual(plugin._db, 2) self.assertEqual(plugin._password, None) + self.assertEqual(plugin._namespace, "") - def test_src_with_host_port_empty_db_pass(self): + def test_src_with_host_port_empty_db_pass_no_namespace(self): plugin = TokenRedis('127.0.0.1:1234::verysecret') self.assertEqual(plugin._server, '127.0.0.1') self.assertEqual(plugin._port, 1234) self.assertEqual(plugin._db, 0) self.assertEqual(plugin._password, 'verysecret') + self.assertEqual(plugin._namespace, "") - def test_src_with_host_empty_port_db_pass(self): + def test_src_with_host_empty_port_db_pass_no_namespace(self): plugin = TokenRedis('127.0.0.1::2:verysecret') self.assertEqual(plugin._server, '127.0.0.1') self.assertEqual(plugin._port, 6379) self.assertEqual(plugin._db, 2) self.assertEqual(plugin._password, 'verysecret') + self.assertEqual(plugin._namespace, "") - def test_src_with_host_empty_port_db_empty_pass(self): + def test_src_with_host_empty_port_db_empty_pass_no_namespace(self): plugin = TokenRedis('127.0.0.1::2:') self.assertEqual(plugin._server, '127.0.0.1') self.assertEqual(plugin._port, 6379) self.assertEqual(plugin._db, 2) self.assertEqual(plugin._password, None) + self.assertEqual(plugin._namespace, "") diff --git a/websockify/token_plugins.py b/websockify/token_plugins.py index 36a1dbc..bb9abe5 100644 --- a/websockify/token_plugins.py +++ b/websockify/token_plugins.py @@ -178,9 +178,9 @@ class TokenRedis(BasePlugin): The token source is in the format: - host[:port[:db[:password]]] + host[:port[:db[:password[:namespace]]]] - where port, db and password are optional. If port or db are left empty + where port, db, password and namespace are optional. If port or db are left empty they will take its default value, ie. 6379 and 0 respectively. If your redis server is using the default port (6379) then you can use: @@ -192,9 +192,14 @@ class TokenRedis(BasePlugin): my-redis-host:::verysecretpass + You can also specify a namespace. In this case, the tokens + will be stored in the format '{namespace}:{token}' + + my-redis-host::::my-app-namespace + In the more general case you will use: - my-redis-host:6380:1:verysecretpass + my-redis-host:6380:1:verysecretpass:my-app-namespace The TokenRedis plugin expects the format of the target in one of these two formats: @@ -234,6 +239,7 @@ class TokenRedis(BasePlugin): self._port = 6379 self._db = 0 self._password = None + self._namespace = "" try: fields = src.split(":") if len(fields) == 1: @@ -256,15 +262,28 @@ class TokenRedis(BasePlugin): self._db = 0 if not self._password: self._password = None + elif len(fields) == 5: + self._server, self._port, self._db, self._password, self._namespace = fields + if not self._port: + self._port = 6379 + if not self._db: + self._db = 0 + if not self._password: + self._password = None + if not self._namespace: + self._namespace = "" else: raise ValueError self._port = int(self._port) self._db = int(self._db) - logger.info("TokenRedis backend initilized (%s:%s)" % + if self._namespace: + self._namespace += ":" + + logger.info("TokenRedis backend initialized (%s:%s)" % (self._server, self._port)) except ValueError: logger.error("The provided --token-source='%s' is not in the " - "expected format [:[:[:]]]" % + "expected format [:[:[:[:]]]]" % src) sys.exit() @@ -278,7 +297,7 @@ class TokenRedis(BasePlugin): logger.info("resolving token '%s'" % token) client = redis.Redis(host=self._server, port=self._port, db=self._db, password=self._password) - stuff = client.get(token) + stuff = client.get(self._namespace + token) if stuff is None: return None else: