Fix module call CLUSTER SLOTS / SHARDS fake client check crash

The reason is VM_Call will use a fake client without connection, so we also need to check if c->conn is NULL. Fixes valkey-io#1054. Signed-off-by: Binbin <[email protected]>
enjoy-binbin · Sep 23, 2024 · be7a566 · be7a566
1 parent d07c297
commit be7a566
Show file tree

Hide file tree

Showing 4 changed files with 76 additions and 4 deletions.
diff --git a/src/networking.c b/src/networking.c
@@ -3246,7 +3246,7 @@ char *getClientSockname(client *c) {
 int isClientConnIpV6(client *c) {
     /* The cached client peer id is on the form "[IPv6]:port" for IPv6
      * addresses, so we just check for '[' here. */
-    if (c->conn->type == NULL && server.current_client) {
+    if ((c->conn == NULL || c->conn->type == NULL) && server.current_client) {
         /* Fake client? Use current client instead. */
         c = server.current_client;
     }

diff --git a/tests/modules/Makefile b/tests/modules/Makefile
@@ -63,7 +63,8 @@ TEST_MODULES = \
     postnotifications.so \
     moduleauthtwo.so \
     rdbloadsave.so \
-    crash.so
+    crash.so \
+    cluster.so
 
 .PHONY: all
 

diff --git a/tests/modules/cluster.c b/tests/modules/cluster.c
@@ -0,0 +1,51 @@
+#include "valkeymodule.h"
+
+#define UNUSED(x) (void)(x)
+
+int test_cluster_slots(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int argc) {
+    UNUSED(argv);
+
+    if (argc != 1) return ValkeyModule_WrongArity(ctx);
+
+    ValkeyModuleCallReply *rep = ValkeyModule_Call(ctx, "CLUSTER", "c", "SLOTS");
+    if (!rep) {
+        ValkeyModule_ReplyWithError(ctx, "ERR NULL reply returned");
+    } else {
+        ValkeyModule_ReplyWithCallReply(ctx, rep);
+        ValkeyModule_FreeCallReply(rep);
+    }
+
+    return VALKEYMODULE_OK;
+}
+
+int test_cluster_shards(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int argc) {
+    UNUSED(argv);
+q
+    if (argc != 1) return ValkeyModule_WrongArity(ctx);
+
+    ValkeyModuleCallReply *rep = ValkeyModule_Call(ctx, "CLUSTER", "c", "SHARDS");
+    if (!rep) {
+        ValkeyModule_ReplyWithError(ctx, "ERR NULL reply returned");
+    } else {
+        ValkeyModule_ReplyWithCallReply(ctx, rep);
+        ValkeyModule_FreeCallReply(rep);
+    }
+
+    return VALKEYMODULE_OK;
+}
+
+int ValkeyModule_OnLoad(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int argc) {
+    VALKEYMODULE_NOT_USED(argv);
+    VALKEYMODULE_NOT_USED(argc);
+
+    if (ValkeyModule_Init(ctx, "cluster", 1, VALKEYMODULE_APIVER_1)== VALKEYMODULE_ERR)
+        return VALKEYMODULE_ERR;
+
+    if (ValkeyModule_CreateCommand(ctx, "test.cluster_slots", test_cluster_slots, "",0,0,0) == VALKEYMODULE_ERR)
+        return VALKEYMODULE_ERR;
+
+    if (ValkeyModule_CreateCommand(ctx, "test.cluster_shards", test_cluster_shards, "", 0, 0, 0) == VALKEYMODULE_ERR)
+        return VALKEYMODULE_ERR;
+
+    return VALKEYMODULE_OK;
+}
diff --git a/tests/unit/moduleapi/cluster.tcl b/tests/unit/moduleapi/cluster.tcl
@@ -219,8 +219,6 @@ start_cluster 2 2 [list config_lines $modules] {
     }
 }
 
-}
-
 set testmodule [file normalize tests/modules/basics.so]
 set modules [list loadmodule $testmodule]
 start_cluster 3 0 [list config_lines $modules] {
@@ -234,3 +232,25 @@ start_cluster 3 0 [list config_lines $modules] {
         assert_equal {PONG} [$node3 PING]
     }
 }
+
+set testmodule [file normalize tests/modules/cluster.so]
+set modules [list loadmodule $testmodule]
+start_cluster 3 0 [list config_lines $modules] {
+    set node1 [srv 0 client]
+    set node2 [srv -1 client]
+    set node3 [srv -2 client]
+
+    test "RM_Call with cluster slots" {
+        assert_equal [lsort [$node1 cluster slots]] [lsort [$node1 test.cluster_slots]]
+        assert_equal [lsort [$node2 cluster slots]] [lsort [$node2 test.cluster_slots]]
+        assert_equal [lsort [$node3 cluster slots]] [lsort [$node3 test.cluster_slots]]
+    }
+
+    test "RM_Call with cluster shards" {
+        assert_equal [lsort [$node1 cluster shards]] [lsort [$node1 test.cluster_shards]]
+        assert_equal [lsort [$node2 cluster shards]] [lsort [$node2 test.cluster_shards]]
+        assert_equal [lsort [$node3 cluster shards]] [lsort [$node3 test.cluster_shards]]
+    }
+}
+
+} ;# end tag