Avoid heap buffer overflow in valkeyAsyncFormattedCommand

src/async.c

@@ -61,6 +61,7 @@
 
 typedef struct {
     sds command;
+    size_t len;
     valkeyCallbackFn *user_callback;
     void *user_priv_data;
 } ssubscribeCallbackData;
@@ -808,21 +809,51 @@ void valkeyAsyncHandleTimeout(valkeyAsyncContext *ac) {
     valkeyAsyncDisconnectInternal(ac);
 }
 
-/* Sets a pointer to the first argument and its length starting at p. Returns
- * the number of bytes to skip to get to the following argument. */
-static const char *nextArgument(const char *start, const char **str, size_t *len) {
-    const char *p = start;
+/* Find the next argument in a command buffer, i.e. find the next bulkstring
+ * in an array of bulkstrings.
+ * Returns a pointer to the end of a found argument, which can be used when
+ * finding following arguments, or NULL when an argument is not found.
+ * The found string is returned by pointer via `str` and length in `strlen`. */
+static const char *nextArgument(const char *buf, size_t buflen, const char **str, size_t *strlen) {
+    const char *p = buf;
+    const size_t minlen = 7; /* '$1\r\n \r\n'*/
+    if (buf == NULL || buflen < minlen)
+        goto error;
+
+    /* Find a bulkstring identifier. */
     if (p[0] != '$') {
         p = strchr(p, '$');
-        if (p == NULL)
-            return NULL;
+        /* Fail if not found, or if buffer can't contain a complete bulkstring. */
+        if (p == NULL || (buflen - (p - buf)) < minlen)
+            goto error;
     }
+    p++; /* Skip found '$' */
+
+    /* Get length and next separator, fail if negative value or empty string. */
+    long int len;
+    char *sep;
+    if ((len = strtol(p, &sep, 10)) < 1)
+        goto error;
+    /* Calculate end ptr, includes length and two separators (`\r\n`). */
+    const char *end = sep + 2 + len + 2;
+
+    /* Validate the parsed length and field separators. */
+    if ((size_t)(end - buf) > buflen ||
+        sep[0] != '\r' || sep[len + 2] != '\r')
+        goto error;
+
+    /* Return pointer to the string, length, and pointer to next element. */
+    *str = sep + 2;
+    *strlen = len;
+
+    if ((size_t)(end - buf) == buflen) /* No more data in buffer? */
+        return NULL;
+    return end;
 
-    *len = (int)strtol(p + 1, NULL, 10);
-    p = strchr(p, '\r');
-    assert(p);
-    *str = p + 2;
-    return p + 2 + (*len) + 2;
+error:
+    *str = NULL;
+    *strlen = 0;
+    return NULL;
 }
 
 void valkeySsubscribeCallback(struct valkeyAsyncContext *ac, void *reply, void *privdata) {
@@ -846,8 +877,8 @@ void valkeySsubscribeCallback(struct valkeyAsyncContext *ac, void *reply, void *
     assert(r != NULL);
     if (r->type == VALKEY_REPLY_ERROR) {
         /*/ On CROSSSLOT, MOVED and other errors */
-        p = nextArgument(data->command, &cstr, &clen);
-        while ((p = nextArgument(p, &astr, &alen)) != NULL) {
+        p = nextArgument(data->command, data->len, &cstr, &clen);
+        while ((p = nextArgument(p, data->len - (p - data->command), &astr, &alen)) != NULL || astr != NULL) {
             sname = sdsnewlen(astr, alen);
             if (sname == NULL)
                 goto oom;
@@ -865,8 +896,8 @@ void valkeySsubscribeCallback(struct valkeyAsyncContext *ac, void *reply, void *
         }
     } else {
         if ((r->type == VALKEY_REPLY_ARRAY || r->type == VALKEY_REPLY_PUSH) && strncasecmp(r->element[0]->str, "ssubscribe", 10) == 0) {
-            p = nextArgument(data->command, &cstr, &clen);
-            while ((p = nextArgument(p, &astr, &alen)) != NULL) {
+            p = nextArgument(data->command, data->len, &cstr, &clen);
+            while ((p = nextArgument(p, data->len - (p - data->command), &astr, &alen)) != NULL || astr != NULL) {
                 sname = sdsnewlen(astr, alen);
                 if (sname == NULL)
                     goto oom;
@@ -921,29 +952,31 @@ static int valkeyAsyncAppendCmdLen(valkeyAsyncContext *ac, valkeyCallbackFn *fn,
     if (c->flags & (VALKEY_DISCONNECTING | VALKEY_FREEING))
         return VALKEY_ERR;
 
-    /* Setup callback */
-    cb.fn = fn;
-    cb.privdata = privdata;
-    cb.pending_subs = 1;
-    cb.unsubscribe_sent = 0;
-    cb.subscribed = 0;
+    /* Get the first string in the command, and don't accept empty commands. */
+    p = nextArgument(cmd, len, &cstr, &clen);
+    if (cstr == NULL)
+        return VALKEY_ERR;
 
-    /* Find out which command will be appended. */
-    p = nextArgument(cmd, &cstr, &clen);
-    assert(p != NULL);
-    hasnext = (p[0] == '$');
+    hasnext = (p && (p[0] == '$'));
     pvariant = (tolower(cstr[0]) == 'p') ? 1 : 0;
     svariant = valkeyIsShardedVariant(cstr);
     hasprefix = svariant || pvariant;
     cstr += hasprefix;
     clen -= hasprefix;
 
+    /* Setup callback */
+    cb.fn = fn;
+    cb.privdata = privdata;
+    cb.pending_subs = 1;
+    cb.unsubscribe_sent = 0;
+    cb.subscribed = 0;
+
     if (hasnext && strncasecmp(cstr, "subscribe\r\n", 11) == 0) {
         int was_subscribed = c->flags & VALKEY_SUBSCRIBED;
         c->flags |= VALKEY_SUBSCRIBED;
 
         /* Add every channel/pattern to the list of subscription callbacks. */
-        while ((p = nextArgument(p, &astr, &alen)) != NULL) {
+        while ((p = nextArgument(p, len - (p - cmd), &astr, &alen)) != NULL || astr != NULL) {
             sname = sdsnewlen(astr, alen);
             if (sname == NULL)
                 goto oom;
@@ -977,15 +1010,12 @@ static int valkeyAsyncAppendCmdLen(valkeyAsyncContext *ac, valkeyCallbackFn *fn,
             if (ssubscribe_data == NULL)
                 goto oom;
 
-            /* copy command to iterate over all channels.
-             * actual length of cmd is actually len + 1 (see valkeyvFormatCommand).
-             * last byte important in nextArgument function.
-             */
-            ssubscribe_data->command = vk_malloc(len + 1);
+            /* Copy command to iterate over all channels. */
+            ssubscribe_data->command = vk_malloc(len);
             if (ssubscribe_data->command == NULL)
                 goto oom;
-            memcpy(ssubscribe_data->command, cmd, len + 1);
-
+            memcpy(ssubscribe_data->command, cmd, len);
+            ssubscribe_data->len = len;
             ssubscribe_data->user_callback = fn;
             ssubscribe_data->user_priv_data = privdata;
 
@@ -1015,7 +1045,7 @@ static int valkeyAsyncAppendCmdLen(valkeyAsyncContext *ac, valkeyCallbackFn *fn,
         if (hasnext) {
             /* Send an unsubscribe with specific channels/patterns.
              * Bookkeeping the number of expected replies */
-            while ((p = nextArgument(p, &astr, &alen)) != NULL) {
+            while ((p = nextArgument(p, len - (p - cmd), &astr, &alen)) != NULL || astr != NULL) {
                 sname = sdsnewlen(astr, alen);
                 if (sname == NULL)
                     goto oom;

tests/client_test.c

@@ -1884,6 +1884,54 @@ void null_cb(valkeyAsyncContext *ac, void *r, void *privdata) {
     state->checkpoint++;
 }
 
+/* Test the command parsing, required for pub/sub in the async API. */
+void test_async_command_parsing(struct config config) {
+    test("Async command parsing: ");
+    valkeyOptions options = get_server_tcp_options(config);
+    valkeyAsyncContext *ac = valkeyAsyncConnectWithOptions(&options);
+    assert(ac);
+
+    /* Null ptr. */
+    assert(VALKEY_ERR == valkeyAsyncFormattedCommand(ac, NULL, NULL, NULL, 45));
+    /* Empty command. */
+    assert(VALKEY_ERR == valkeyAsyncFormattedCommand(ac, NULL, NULL, "", 0));
+    assert(VALKEY_ERR == valkeyAsyncFormattedCommand(ac, NULL, NULL, "  $", 2));
+    assert(VALKEY_ERR == valkeyAsyncFormattedCommand(ac, NULL, NULL, "*0\r\n", 4));
+    assert(VALKEY_ERR == valkeyAsyncFormattedCommand(ac, NULL, NULL, "*1\r\n$0\r\n\r\n", 10));
+    assert(VALKEY_ERR == valkeyAsyncFormattedCommand(ac, NULL, NULL, "*1\r\n$-1\r\n", 9));
+    assert(VALKEY_ERR == valkeyAsyncFormattedCommand(ac, NULL, NULL, "*1\r\n$-1\r\nUNSUBSCRIBE\r\n", 22));
+    /* Protocol error: erroneous bulkstring length and data. */
+    assert(VALKEY_ERR == valkeyAsyncFormattedCommand(ac, NULL, NULL, "*1\r\n$100000", 11));
+    assert(VALKEY_ERR == valkeyAsyncFormattedCommand(ac, NULL, NULL, "*1\r\n$100000\r", 12));
+    assert(VALKEY_ERR == valkeyAsyncFormattedCommand(ac, NULL, NULL, "*1\r\n$100000\r\n", 13));
+    assert(VALKEY_ERR == valkeyAsyncFormattedCommand(ac, NULL, NULL, "*1\r\n$10HELP\r\n\r\n", 15));
+    assert(VALKEY_ERR == valkeyAsyncFormattedCommand(ac, NULL, NULL, "*1\r\n$100000\r\nTO-SHORT\r\n", 23));
+    assert(VALKEY_ERR == valkeyAsyncFormattedCommand(ac, NULL, NULL, "*1\r\n$1\r\nTO-LONG\r\n", 17));
+    assert(VALKEY_ERR == valkeyAsyncFormattedCommand(ac, NULL, NULL, "*1\r\n$123456789\r\n", 11));
+
+    /* Faulty length given to function. */
+    for (int i = 0; i < 19; i++) {
+        assert(VALKEY_ERR == valkeyAsyncFormattedCommand(ac, NULL, NULL, "*2\r\n$9\r\nSUBSCRIBE\r\n$7\r\nCHANNEL\r\n", i));
+    }
+    for (int i = 0; i < 21; i++) {
+        assert(VALKEY_ERR == valkeyAsyncFormattedCommand(ac, NULL, NULL, "*1\r\n$11\r\nUNSUBSCRIBE\r\n", i));
+    }
+
+    /* Complete command. */
+    assert(VALKEY_OK == valkeyAsyncFormattedCommand(ac, NULL, NULL, "*2\r\n$9\r\nSUBSCRIBE\r\n$7\r\nCHANNEL\r\n", 32));
+    assert(VALKEY_OK == valkeyAsyncFormattedCommand(ac, NULL, NULL, "*1\r\n$11\r\nUNSUBSCRIBE\r\n", 22));
+
+    // Heap allocate command without NULL terminator.
+    const char ping[] = "*1\r\n$4\r\nPING\r\n";
+    size_t len = sizeof(ping) - 1;
+    char *buf = vk_malloc_safe(len);
+    memcpy(buf, ping, len);
+    assert(VALKEY_OK == valkeyAsyncFormattedCommand(ac, NULL, NULL, buf, len));
+    free(buf);
+
+    valkeyAsyncFree(ac);
+}
+
 static void test_pubsub_handling(struct config config) {
     test("Subscribe, handle published message and unsubscribe: ");
     /* Setup event dispatcher with a testcase timeout */
@@ -2817,6 +2865,7 @@ int main(int argc, char **argv) {
     get_server_version(c, &major, NULL);
     disconnect(c, 0);
 
+    test_async_command_parsing(cfg);
     test_pubsub_handling(cfg);
     test_pubsub_multiple_channels(cfg);
     test_monitor(cfg);