Merge pull request #154 from alisaifee/fix-cluster-auth-from-uri

Support authentication through storage string for redis cluster
alisaifee · Jan 24, 2023 · a92ea21 · a92ea21
2 parents 68ee7f6 + c06039b
commit a92ea21
Show file tree

Hide file tree

Showing 7 changed files with 119 additions and 5 deletions.
diff --git a/docker-compose.yml b/docker-compose.yml
@@ -165,6 +165,28 @@ services:
     depends_on: [redis-ssl-cluster-1, redis-ssl-cluster-2, redis-ssl-cluster-3, redis-ssl-cluster-4, redis-ssl-cluster-5, redis-ssl-cluster-6]
     volumes:
       - ./tests/tls:/tls
+  redis-cluster-auth-1:
+    image: "redis:${LIMITS_REDIS_SERVER_VERSION:-latest}"
+    command: redis-server --port 8400 ${DEFAULT_ARGS---enable-debug-command yes} --protected-mode no --cluster-enabled yes --loglevel verbose --cluster-announce-ip ${HOST_IP} --requirepass sekret
+    ports:
+      - '8400:8400'
+      - '18400:18400'
+  redis-cluster-auth-2:
+    image: "redis:${LIMITS_REDIS_SERVER_VERSION:-latest}"
+    command: redis-server --port 8401 ${DEFAULT_ARGS---enable-debug-command yes} --protected-mode no --cluster-enabled yes --loglevel verbose --cluster-announce-ip ${HOST_IP} --requirepass sekret
+    ports:
+      - '8401:8401'
+      - '18401:18401'
+  redis-cluster-auth-3:
+    image: "redis:${LIMITS_REDIS_SERVER_VERSION:-latest}"
+    command: redis-server --port 8402 ${DEFAULT_ARGS---enable-debug-command yes} --protected-mode no --cluster-enabled yes --loglevel verbose --cluster-announce-ip ${HOST_IP} --requirepass sekret
+    ports:
+      - '8402:8402'
+      - '18402:18402'
+  redis-cluster-auth-init:
+    image: "redis:${LIMITS_REDIS_SERVER_VERSION:-latest}"
+    command: bash -c "echo yes | redis-cli --cluster create --cluster-replicas 0 ${HOST_IP}:8400 ${HOST_IP}:8401 ${HOST_IP}:8402 -a sekret"
+    depends_on: [redis-cluster-auth-1, redis-cluster-auth-2, redis-cluster-auth-3]
   redis-basic:
     image: "redis:${LIMITS_REDIS_SERVER_VERSION:-latest}"
     command: redis-server --port 7379

diff --git a/limits/aio/storage/redis.py b/limits/aio/storage/redis.py
@@ -295,9 +295,17 @@ def __init__(self, uri: str, **options: Union[float, str, bool]) -> None:
          available or if the redis host cannot be pinged.
         """
         parsed = urllib.parse.urlparse(uri)
+        parsed_auth: Dict[str, Union[float, str, bool]] = {}
+
+        if parsed.username:
+            parsed_auth["username"] = parsed.username
+        if parsed.password:
+            parsed_auth["password"] = parsed.password
+
+        sep = parsed.netloc.find("@") + 1
         cluster_hosts = []
 
-        for loc in parsed.netloc.split(","):
+        for loc in parsed.netloc[sep:].split(","):
             host, port = loc.split(":")
             cluster_hosts.append({"host": host, "port": int(port)})
 
@@ -306,7 +314,8 @@ def __init__(self, uri: str, **options: Union[float, str, bool]) -> None:
         self.dependency = self.dependencies["coredis"].module
 
         self.storage: "coredis.RedisCluster[str]" = self.dependency.RedisCluster(
-            startup_nodes=cluster_hosts, **{**self.DEFAULT_OPTIONS, **options}
+            startup_nodes=cluster_hosts,
+            **{**self.DEFAULT_OPTIONS, **parsed_auth, **options},
         )
         self.initialize_storage(uri)
 

diff --git a/limits/storage/redis_cluster.py b/limits/storage/redis_cluster.py
@@ -50,14 +50,24 @@ def __init__(self, uri: str, **options: Union[float, str, bool]) -> None:
          available or if the redis cluster cannot be reached.
         """
         parsed = urllib.parse.urlparse(uri)
+        parsed_auth: Dict[str, Union[float, str, bool]] = {}
+
+        if parsed.username:
+            parsed_auth["username"] = parsed.username
+        if parsed.password:
+            parsed_auth["password"] = parsed.password
+
+        sep = parsed.netloc.find("@") + 1
         cluster_hosts = []
-        for loc in parsed.netloc.split(","):
+        for loc in parsed.netloc[sep:].split(","):
             host, port = loc.split(":")
             cluster_hosts.append((host, int(port)))
 
         self.storage = None
         self.using_redis_py = False
-        self.__pick_storage(cluster_hosts, **{**self.DEFAULT_OPTIONS, **options})
+        self.__pick_storage(
+            cluster_hosts, **{**self.DEFAULT_OPTIONS, **parsed_auth, **options}
+        )
         assert self.storage
         self.initialize_storage(uri)
         super(RedisStorage, self).__init__(uri, **options)
@@ -99,7 +109,7 @@ def __use_legacy_cluster_implementation(
         )
         self.storage = redis_cluster.RedisCluster(
             startup_nodes=[{"host": c[0], "port": c[1]} for c in cluster_hosts],
-            **options
+            **options,
         )
 
     def reset(self) -> Optional[int]:

diff --git a/tests/aio/test_storage.py b/tests/aio/test_storage.py
@@ -158,6 +158,14 @@ async def get_moving_window(self, *a, **k):
             marks=pytest.mark.redis_cluster,
             id="redis-cluster",
         ),
+        pytest.param(
+            "async+redis+cluster://:sekret@localhost:8400/",
+            {},
+            RedisClusterStorage,
+            pytest.lazy_fixture("redis_auth_cluster"),
+            marks=pytest.mark.redis_cluster,
+            id="redis-cluster-auth",
+        ),
         pytest.param(
             "async+mongodb://localhost:37017/",
             {},

diff --git a/tests/conftest.py b/tests/conftest.py
@@ -19,6 +19,16 @@ def check_redis_cluster_ready(host, port):
         return False
 
 
+def check_redis_auth_cluster_ready(host, port):
+    try:
+        return (
+            redis.Redis(host, port, password="sekret").cluster("info")["cluster_state"]
+            == "ok"
+        )
+    except Exception:
+        return False
+
+
 def check_redis_ssl_cluster_ready(host, port):
 
     storage_url = (
@@ -148,6 +158,18 @@ def redis_cluster_client(docker_services):
     return redis.cluster.RedisCluster("localhost", 7001)
 
 
+@pytest.fixture(scope="session")
+def redis_auth_cluster_client(docker_services):
+    docker_services.start("redis-cluster-auth-init")
+    docker_services.wait_for_service(
+        "redis-cluster-auth-3", 8402, check_redis_auth_cluster_ready
+    )
+    if os.environ.get("CI") == "True":
+        time.sleep(10)
+
+    return redis.cluster.RedisCluster("localhost", 8400, password="sekret")
+
+
 @pytest.fixture(scope="session")
 def redis_ssl_cluster_client(docker_services):
     docker_services.start("redis-ssl-cluster-init")
@@ -276,6 +298,13 @@ def redis_cluster(redis_cluster_client):
     return redis_cluster_client
 
 
+@pytest.fixture
+def redis_auth_cluster(redis_auth_cluster_client):
+    redis_auth_cluster_client.flushall()
+
+    return redis_auth_cluster_client
+
+
 @pytest.fixture
 def redis_ssl_cluster(redis_ssl_cluster_client):
     redis_ssl_cluster_client.flushall()

diff --git a/tests/test_storage.py b/tests/test_storage.py
@@ -172,6 +172,14 @@ def get_moving_window(self, *a, **k):
             marks=pytest.mark.redis_cluster,
             id="redis-cluster",
         ),
+        pytest.param(
+            "redis+cluster://:sekret@localhost:8400/",
+            {},
+            RedisClusterStorage,
+            pytest.lazy_fixture("redis_auth_cluster"),
+            marks=pytest.mark.redis_cluster,
+            id="redis-cluster-auth",
+        ),
         pytest.param(
             "mongodb://localhost:37017/",
             {},

diff --git a/tests/utils.py b/tests/utils.py
@@ -80,6 +80,13 @@ async def async_window(delay_end: float, delay: Optional[float] = None):
             marks=pytest.mark.redis_cluster,
             id="redis-cluster",
         ),
+        pytest.param(
+            "redis+cluster://:sekret@localhost:8400/",
+            {},
+            pytest.lazy_fixture("redis_auth_cluster"),
+            marks=pytest.mark.redis_cluster,
+            id="redis-cluster-auth",
+        ),
         pytest.param(
             "redis+cluster://localhost:8301",
             {
@@ -135,6 +142,13 @@ async def async_window(delay_end: float, delay: Optional[float] = None):
             marks=pytest.mark.redis_cluster,
             id="redis-cluster",
         ),
+        pytest.param(
+            "redis+cluster://:sekret@localhost:8400/",
+            {},
+            pytest.lazy_fixture("redis_auth_cluster"),
+            marks=pytest.mark.redis_cluster,
+            id="redis-cluster-auth",
+        ),
         pytest.param(
             "redis+cluster://localhost:8301",
             {
@@ -197,6 +211,13 @@ async def async_window(delay_end: float, delay: Optional[float] = None):
             marks=pytest.mark.redis_cluster,
             id="redis-cluster",
         ),
+        pytest.param(
+            "async+redis+cluster://:sekret@localhost:8400/",
+            {},
+            pytest.lazy_fixture("redis_auth_cluster"),
+            marks=pytest.mark.redis_cluster,
+            id="redis-cluster-auth",
+        ),
         pytest.param(
             "async+redis+cluster://localhost:8301",
             {
@@ -252,6 +273,13 @@ async def async_window(delay_end: float, delay: Optional[float] = None):
             marks=pytest.mark.redis_cluster,
             id="redis-cluster",
         ),
+        pytest.param(
+            "async+redis+cluster://:sekret@localhost:8400/",
+            {},
+            pytest.lazy_fixture("redis_auth_cluster"),
+            marks=pytest.mark.redis_cluster,
+            id="redis-cluster-auth",
+        ),
         pytest.param(
             "async+redis+cluster://localhost:8301",
             {