Add namespace value to TokenRedis source

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

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 <host>[:<port>[:<db>[:<password>]]]" %
+                         "expected format <host>[:<port>[:<db>[:<password>[:<namespace>]]]]" %
                          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: