Fix cluster Pub/Sub length truncation (#5030)

Validate the full publish message length before converting it to the cluster bus 32-bit length field. Reject oversized messages before allocation and add integration coverage for PUBLISH and SPUBLISH.

(cherry picked from commit ff765915e1)

Signed-off-by: Roshan Khatri <rvkhatri@amazon.com>
Signed-off-by: Sarthak Aggarwal <sarthagg@amazon.com>

src/cluster.h

@@ -68,6 +68,8 @@ void clusterUpdateMyselfHostname(void);
 void clusterUpdateMyselfAnnouncedPorts(void);
 void clusterUpdateMyselfHumanNodename(void);
 
+int clusterPublishMessageLengthValid(robj *channel, robj *message);
+void clusterSetPublishMessageLengthLimitForTesting(size_t limit);
 void clusterPropagatePublish(robj *channel, robj *message, int sharded);
 void clusterBroadcastPong(int target);
 

src/cluster_legacy.c

@@ -4832,7 +4832,26 @@ void clusterBroadcastPong(int target) {
     dictReleaseIterator(di);
 }
 
-/* Create a PUBLISH message block.
+static size_t cluster_publish_message_length_limit = 0;
+
+void clusterSetPublishMessageLengthLimitForTesting(size_t limit) {
+    cluster_publish_message_length_limit = limit;
+}
+
+int clusterPublishMessageLengthValid(robj *channel, robj *message) {
+    size_t channel_len = stringObjectLen(channel);
+    size_t message_len = stringObjectLen(message);
+    size_t max_payload_len = (size_t)UINT32_MAX - offsetof(clusterMsgSendBlock, data) - CLUSTERMSG_MIN_LEN -
+                             offsetof(clusterMsgDataPublish, bulk_data);
+
+    if (cluster_publish_message_length_limit && cluster_publish_message_length_limit < max_payload_len)
+        max_payload_len = cluster_publish_message_length_limit;
+
+    return channel_len <= max_payload_len && message_len <= max_payload_len - channel_len;
+}
+
+/* Create a PUBLISH message block, or return NULL if it exceeds the supported
+ * cluster bus publish size.
  *
  * Sanitizer suppression: In clusterMsgDataPublish, sizeof(bulk_data) is 8.
  * As all the struct is used as a buffer, when more than 8 bytes are copied into
@@ -4840,13 +4859,14 @@ void clusterBroadcastPong(int target) {
  * positive in this context. */
 VALKEY_NO_SANITIZE("bounds")
 clusterMsgSendBlock *clusterCreatePublishMsgBlock(robj *channel, robj *message, int is_light, int is_sharded) {
-    uint32_t channel_len, message_len;
     uint16_t type = is_sharded ? CLUSTERMSG_TYPE_PUBLISHSHARD : CLUSTERMSG_TYPE_PUBLISH;
 
+    if (!clusterPublishMessageLengthValid(channel, message)) return NULL;
+
     channel = getDecodedObject(channel);
     message = getDecodedObject(message);
-    channel_len = sdslen(channel->ptr);
-    message_len = sdslen(message->ptr);
+    size_t channel_len = sdslen(channel->ptr);
+    size_t message_len = sdslen(message->ptr);
     size_t msglen;
 
     if (is_light) {
@@ -4867,10 +4887,10 @@ clusterMsgSendBlock *clusterCreatePublishMsgBlock(robj *channel, robj *message,
         clusterMsg *hdr = getMessageFromSendBlock(msgblock);
         hdr_data_msg = &hdr->data.publish.msg;
     }
-    hdr_data_msg->channel_len = htonl(channel_len);
-    hdr_data_msg->message_len = htonl(message_len);
-    memcpy(hdr_data_msg->bulk_data, channel->ptr, sdslen(channel->ptr));
-    memcpy(hdr_data_msg->bulk_data + sdslen(channel->ptr), message->ptr, sdslen(message->ptr));
+    hdr_data_msg->channel_len = htonl((uint32_t)channel_len);
+    hdr_data_msg->message_len = htonl((uint32_t)message_len);
+    memcpy(hdr_data_msg->bulk_data, channel->ptr, channel_len);
+    memcpy(hdr_data_msg->bulk_data + channel_len, message->ptr, message_len);
 
     decrRefCount(channel);
     decrRefCount(message);
@@ -5029,6 +5049,7 @@ int clusterSendModuleMessageToTarget(const char *target,
 void clusterPropagatePublish(robj *channel, robj *message, int sharded) {
     clusterMsgSendBlock *msgblock, *msgblock_light;
     msgblock_light = clusterCreatePublishMsgBlock(channel, message, 1, sharded);
+    if (msgblock_light == NULL) return;
     /* We will only create msgblock with normal hdr if there are any nodes that do not support light hdr */
     msgblock = NULL;
     ClusterNodeIterator iter;

src/debug.c

@@ -481,6 +481,9 @@ void debugCommand(client *c) {
             "QUICKLIST-PACKED-THRESHOLD <size>",
             "    Sets the threshold for elements to be inserted as plain vs packed nodes",
             "    Default value is 1GB, allows values up to 4GB. Setting to 0 restores to default.",
+            "CLUSTER-PUBLISH-MESSAGE-LIMIT <size>",
+            "    Sets a lower cluster publish payload limit for testing.",
+            "    Setting to 0 restores the default.",
             "SET-SKIP-CHECKSUM-VALIDATION <0|1>",
             "    Enables or disables checksum checks for RDB files and RESTORE's payload.",
             "SLEEP <seconds>",
@@ -894,6 +897,15 @@ void debugCommand(client *c) {
         } else {
             addReply(c, shared.ok);
         }
+    } else if (!strcasecmp(c->argv[1]->ptr, "cluster-publish-message-limit") && c->argc == 3) {
+        int memerr;
+        unsigned long long sz = memtoull((const char *)c->argv[2]->ptr, &memerr);
+        if (memerr || sz > SIZE_MAX) {
+            addReplyError(c, "argument must be a valid memory value");
+        } else {
+            clusterSetPublishMessageLengthLimitForTesting((size_t)sz);
+            addReply(c, shared.ok);
+        }
     } else if (!strcasecmp(c->argv[1]->ptr, "set-skip-checksum-validation") && c->argc == 3) {
         server.skip_checksum_validation = atoi(c->argv[2]->ptr);
         addReply(c, shared.ok);

src/pubsub.c

@@ -639,8 +639,12 @@ void punsubscribeCommand(client *c) {
 }
 
 /* This function wraps pubsubPublishMessage and also propagates the message to cluster.
- * Used by the commands PUBLISH/SPUBLISH and their respective module APIs.*/
+ * Used by the commands PUBLISH/SPUBLISH and their respective module APIs.
+ * Returns zero without publishing if the message exceeds the supported cluster
+ * bus publish size. */
 int pubsubPublishMessageAndPropagateToCluster(robj *channel, robj *message, int sharded) {
+    if (server.cluster_enabled && !clusterPublishMessageLengthValid(channel, message)) return 0;
+
     int receivers = pubsubPublishMessage(channel, message, sharded);
     if (server.cluster_enabled) clusterPropagatePublish(channel, message, sharded);
     return receivers;
@@ -653,6 +657,10 @@ void publishCommand(client *c) {
         return;
     }
 
+    if (server.cluster_enabled && !clusterPublishMessageLengthValid(c->argv[1], c->argv[2])) {
+        addReplyError(c, "Publish message is too large for the cluster bus");
+        return;
+    }
     int receivers = pubsubPublishMessageAndPropagateToCluster(c->argv[1], c->argv[2], 0);
     if (!server.cluster_enabled) forceCommandPropagation(c, PROPAGATE_REPL);
     addReplyLongLong(c, receivers);
@@ -743,6 +751,10 @@ void channelList(client *c, sds pat, kvstore *pubsub_channels) {
 
 /* SPUBLISH <shardchannel> <message> */
 void spublishCommand(client *c) {
+    if (server.cluster_enabled && !clusterPublishMessageLengthValid(c->argv[1], c->argv[2])) {
+        addReplyError(c, "Publish message is too large for the cluster bus");
+        return;
+    }
     int receivers = pubsubPublishMessageAndPropagateToCluster(c->argv[1], c->argv[2], 1);
     if (!server.cluster_enabled) forceCommandPropagation(c, PROPAGATE_REPL);
     addReplyLongLong(c, receivers);

tests/unit/cluster/pubsub.tcl

@@ -34,6 +34,30 @@ test "Test publishing to master" {
 test "Test publishing to slave" {
     test_cluster_publish 5 10
 }
+
+test "Oversized PUBLISH is rejected before local or cluster delivery" {
+    set channel publish-limit
+    set local_subscriber [valkey_deferring_client_by_addr [srv 0 host] [srv 0 port]]
+    set remote_subscriber [valkey_deferring_client_by_addr [srv -1 host] [srv -1 port]]
+
+    foreach subscriber [list $local_subscriber $remote_subscriber] {
+        $subscriber subscribe $channel
+        $subscriber read
+    }
+
+    R 0 DEBUG CLUSTER-PUBLISH-MESSAGE-LIMIT 32
+    catch {R 0 PUBLISH $channel [string repeat x 32]} err
+    R 0 DEBUG CLUSTER-PUBLISH-MESSAGE-LIMIT 0
+    assert_match {*Publish message is too large for the cluster bus*} $err
+
+    set data [randomValue]
+    R 0 PUBLISH $channel $data
+    foreach subscriber [list $local_subscriber $remote_subscriber] {
+        set msg [$subscriber read]
+        assert_equal $data [lindex $msg 2]
+        $subscriber close
+    }
+}
 } ;# start_cluster
 
 start_cluster 3 0 {tags {external:skip cluster}} {

tests/unit/cluster/pubsubshard.tcl

@@ -42,6 +42,29 @@ test "Pub/Sub shard basics" {
     $anotherclient close
 }
 
+test "Oversized SPUBLISH is rejected before local delivery" {
+    set slot [$cluster cluster keyslot "channel.0"]
+    array set publishnode [$cluster masternode_for_slot $slot]
+    set publishclient [valkey_client_by_addr $publishnode(host) $publishnode(port)]
+    set subscribeclient [valkey_deferring_client_by_addr $publishnode(host) $publishnode(port)]
+
+    $subscribeclient ssubscribe channel.0
+    $subscribeclient read
+
+    $publishclient debug cluster-publish-message-limit 32
+    catch {$publishclient spublish channel.0 [string repeat x 32]} err
+    $publishclient debug cluster-publish-message-limit 0
+    assert_match {*Publish message is too large for the cluster bus*} $err
+
+    set data [randomValue]
+    $publishclient spublish channel.0 $data
+    set msg [$subscribeclient read]
+    assert_equal $data [lindex $msg 2]
+
+    $publishclient close
+    $subscribeclient close
+}
+
 test "client can't subscribe to multiple shard channels across different slots in same call" {
     catch {$cluster ssubscribe channel.0 channel.1} err
     assert_match {CROSSSLOT Keys*} $err