Remove recursion of redis command parser

sunhao.inf · sunhao.inf · commit 571ead76b373 · 2025-11-03T17:41:25.000+08:00
diff --git a/src/brpc/redis_command.cpp b/src/brpc/redis_command.cpp
@@ -377,136 +377,138 @@ size_t RedisCommandParser::ParsedArgsSize() {
 ParseError RedisCommandParser::Consume(butil::IOBuf& buf,
                                        std::vector<butil::StringPiece>* args,
                                        butil::Arena* arena) {
-    const auto pfc = static_cast<const char *>(buf.fetch1());
-    if (pfc == NULL) {
-        return PARSE_ERROR_NOT_ENOUGH_DATA;
-    }
-    // '*' stands for array "*<size>\r\n<sub-reply1><sub-reply2>..."
-    if (!_parsing_array && *pfc != '*') {
-        if (!std::isalpha(static_cast<unsigned char>(*pfc))) {
-            return PARSE_ERROR_TRY_OTHERS;
+    do {
+        const auto pfc = static_cast<const char *>(buf.fetch1());
+        if (pfc == NULL) {
+            return PARSE_ERROR_NOT_ENOUGH_DATA;
         }
-        const size_t buf_size = buf.size();
-        const auto copy_str = static_cast<char *>(arena->allocate(buf_size + 1));
-        buf.copy_to(copy_str, buf_size);
-        if (*copy_str == ' ') {
+        // '*' stands for array "*<size>\r\n<sub-reply1><sub-reply2>..."
+        if (!_parsing_array && *pfc != '*') {
+            if (!std::isalpha(static_cast<unsigned char>(*pfc))) {
+                return PARSE_ERROR_TRY_OTHERS;
+            }
+            const size_t buf_size = buf.size();
+            const auto copy_str = static_cast<char *>(arena->allocate(buf_size + 1));
+            buf.copy_to(copy_str, buf_size);
+            if (*copy_str == ' ') {
+                return PARSE_ERROR_ABSOLUTELY_WRONG;
+            }
+            copy_str[buf_size] = '\0';
+            const size_t crlf_pos = butil::StringPiece(copy_str, buf_size).find("\r\n");
+            if (crlf_pos == butil::StringPiece::npos) {  // not enough data
+                return PARSE_ERROR_NOT_ENOUGH_DATA;
+            }
+            args->clear();
+            size_t offset = 0;
+            while (offset < crlf_pos && copy_str[offset] != ' ') {
+                ++offset;
+            }
+            const auto first_arg = static_cast<char*>(arena->allocate(offset));
+            memcpy(first_arg, copy_str, offset);
+            for (size_t i = 0; i < offset; ++i) {
+                first_arg[i] = tolower(first_arg[i]);
+            }
+            args->push_back(butil::StringPiece(first_arg, offset));
+            if (offset == crlf_pos) {
+                // only one argument, directly return
+                buf.pop_front(crlf_pos + 2);
+                return PARSE_OK;
+            }
+            size_t arg_start_pos = ++offset;
+
+            for (; offset < crlf_pos; ++offset) {
+                if (copy_str[offset] != ' ') {
+                    continue;
+                }
+                const auto arg_length = offset - arg_start_pos;
+                const auto arg = static_cast<char *>(arena->allocate(arg_length));
+                memcpy(arg, copy_str + arg_start_pos, arg_length);
+                args->push_back(butil::StringPiece(arg, arg_length));
+                arg_start_pos = ++offset;
+            }
+
+            if (arg_start_pos < crlf_pos) {
+                // process the last argument
+                const auto arg_length = crlf_pos - arg_start_pos;
+                const auto arg = static_cast<char *>(arena->allocate(arg_length));
+                memcpy(arg, copy_str + arg_start_pos, arg_length);
+                args->push_back(butil::StringPiece(arg, arg_length));
+            }
+
+            buf.pop_front(crlf_pos + 2);
+            return PARSE_OK;
+        }
+        // '$' stands for bulk string "$<length>\r\n<string>\r\n"
+        if (_parsing_array && *pfc != '$') {
             return PARSE_ERROR_ABSOLUTELY_WRONG;
         }
-        copy_str[buf_size] = '\0';
-        const size_t crlf_pos = butil::StringPiece(copy_str, buf_size).find("\r\n");
+        char intbuf[32];  // enough for fc + 64-bit decimal + \r\n
+        const size_t ncopied = buf.copy_to(intbuf, sizeof(intbuf) - 1);
+        intbuf[ncopied] = '\0';
+        const size_t crlf_pos = butil::StringPiece(intbuf, ncopied).find("\r\n");
         if (crlf_pos == butil::StringPiece::npos) {  // not enough data
             return PARSE_ERROR_NOT_ENOUGH_DATA;
         }
-        args->clear();
-        size_t offset = 0;
-        while (offset < crlf_pos && copy_str[offset] != ' ') {
-            ++offset;
-        }
-        const auto first_arg = static_cast<char*>(arena->allocate(offset));
-        memcpy(first_arg, copy_str, offset);
-        for (size_t i = 0; i < offset; ++i) {
-            first_arg[i] = tolower(first_arg[i]);
+        char* endptr = NULL;
+        int64_t value = strtoll(intbuf + 1/*skip fc*/, &endptr, 10);
+        if (endptr != intbuf + crlf_pos) {
+            LOG(ERROR) << '`' << intbuf + 1 << "' is not a valid 64-bit decimal";
+            return PARSE_ERROR_ABSOLUTELY_WRONG;
         }
-        args->push_back(butil::StringPiece(first_arg, offset));
-        if (offset == crlf_pos) {
-            // only one argument, directly return
-            buf.pop_front(crlf_pos + 2);
-            return PARSE_OK;
+        if (value < 0) {
+            LOG(ERROR) << "Invalid len=" << value << " in redis command";
+            return PARSE_ERROR_ABSOLUTELY_WRONG;
         }
-        size_t arg_start_pos = ++offset;
-
-        for (; offset < crlf_pos; ++offset) {
-            if (copy_str[offset] != ' ') {
-                continue;
+        if (!_parsing_array) {
+            if (value > (int64_t)(FLAGS_redis_max_allocation_size / sizeof(butil::StringPiece))) {
+                LOG(ERROR) << "command array size exceeds limit! max="
+                        << (FLAGS_redis_max_allocation_size / sizeof(butil::StringPiece))
+                        << ", actually=" << value;
+                return PARSE_ERROR_ABSOLUTELY_WRONG;
             }
-            const auto arg_length = offset - arg_start_pos;
-            const auto arg = static_cast<char *>(arena->allocate(arg_length));
-            memcpy(arg, copy_str + arg_start_pos, arg_length);
-            args->push_back(butil::StringPiece(arg, arg_length));
-            arg_start_pos = ++offset;
+            buf.pop_front(crlf_pos + 2/*CRLF*/);
+            _parsing_array = true;
+            _length = value;
+            _index = 0;
+            _args.resize(value);
+            continue;
         }
-
-        if (arg_start_pos < crlf_pos) {
-            // process the last argument
-            const auto arg_length = crlf_pos - arg_start_pos;
-            const auto arg = static_cast<char *>(arena->allocate(arg_length));
-            memcpy(arg, copy_str + arg_start_pos, arg_length);
-            args->push_back(butil::StringPiece(arg, arg_length));
+        CHECK(_index < _length) << "a complete command has been parsed. "
+                "impl of RedisCommandParser::Parse is buggy";
+        const int64_t len = value;  // `value' is length of the string
+        if (len < 0) {
+            LOG(ERROR) << "string in command is nil!";
+            return PARSE_ERROR_ABSOLUTELY_WRONG;
         }
-
-        buf.pop_front(crlf_pos + 2);
-        return PARSE_OK;
-    }
-    // '$' stands for bulk string "$<length>\r\n<string>\r\n"
-    if (_parsing_array && *pfc != '$') {
-        return PARSE_ERROR_ABSOLUTELY_WRONG;
-    }
-    char intbuf[32];  // enough for fc + 64-bit decimal + \r\n
-    const size_t ncopied = buf.copy_to(intbuf, sizeof(intbuf) - 1);
-    intbuf[ncopied] = '\0';
-    const size_t crlf_pos = butil::StringPiece(intbuf, ncopied).find("\r\n");
-    if (crlf_pos == butil::StringPiece::npos) {  // not enough data
-        return PARSE_ERROR_NOT_ENOUGH_DATA;
-    }
-    char* endptr = NULL;
-    int64_t value = strtoll(intbuf + 1/*skip fc*/, &endptr, 10);
-    if (endptr != intbuf + crlf_pos) {
-        LOG(ERROR) << '`' << intbuf + 1 << "' is not a valid 64-bit decimal";
-        return PARSE_ERROR_ABSOLUTELY_WRONG;
-    }
-    if (value < 0) {
-        LOG(ERROR) << "Invalid len=" << value << " in redis command";
-        return PARSE_ERROR_ABSOLUTELY_WRONG;
-    }
-    if (!_parsing_array) {
-        if (value > (int64_t)(FLAGS_redis_max_allocation_size / sizeof(butil::StringPiece))) {
-            LOG(ERROR) << "command array size exceeds limit! max=" 
-                       << (FLAGS_redis_max_allocation_size / sizeof(butil::StringPiece)) 
-                       << ", actually=" << value;
+        if (len > FLAGS_redis_max_allocation_size) {
+            LOG(ERROR) << "command string exceeds max allocation size! max="
+                    << FLAGS_redis_max_allocation_size << ", actually=" << len;
             return PARSE_ERROR_ABSOLUTELY_WRONG;
         }
+        if (buf.size() < crlf_pos + 2 + (size_t)len + 2/*CRLF*/) {
+            return PARSE_ERROR_NOT_ENOUGH_DATA;
+        }
         buf.pop_front(crlf_pos + 2/*CRLF*/);
-        _parsing_array = true;
-        _length = value;
-        _index = 0;
-        _args.resize(value);
-        return Consume(buf, args, arena);
-    }
-    CHECK(_index < _length) << "a complete command has been parsed. "
-            "impl of RedisCommandParser::Parse is buggy";
-    const int64_t len = value;  // `value' is length of the string
-    if (len < 0) {
-        LOG(ERROR) << "string in command is nil!";
-        return PARSE_ERROR_ABSOLUTELY_WRONG;
-    }
-    if (len > FLAGS_redis_max_allocation_size) {
-        LOG(ERROR) << "command string exceeds max allocation size! max=" 
-                   << FLAGS_redis_max_allocation_size << ", actually=" << len;
-        return PARSE_ERROR_ABSOLUTELY_WRONG;
-    }
-    if (buf.size() < crlf_pos + 2 + (size_t)len + 2/*CRLF*/) {
-        return PARSE_ERROR_NOT_ENOUGH_DATA;
-    }
-    buf.pop_front(crlf_pos + 2/*CRLF*/);
-    char* d = (char*)arena->allocate((len/8 + 1) * 8);
-    buf.cutn(d, len);
-    d[len] = '\0';
-    _args[_index].set(d, len);
-    if (_index == 0) {
-        // convert it to lowercase when it is command name
-        for (int i = 0; i < len; ++i) {
-            d[i] = ::tolower(d[i]);
+        char* d = (char*)arena->allocate((len/8 + 1) * 8);
+        buf.cutn(d, len);
+        d[len] = '\0';
+        _args[_index].set(d, len);
+        if (_index == 0) {
+            // convert it to lowercase when it is command name
+            for (int i = 0; i < len; ++i) {
+                d[i] = ::tolower(d[i]);
+            }
         }
-    }
-    char crlf[2];
-    buf.cutn(crlf, sizeof(crlf));
-    if (crlf[0] != '\r' || crlf[1] != '\n') {
-        LOG(ERROR) << "string in command is not ended with CRLF";
-        return PARSE_ERROR_ABSOLUTELY_WRONG;
-    }
-    if (++_index < _length) {
-        return Consume(buf, args, arena);
-    }
+        char crlf[2];
+        buf.cutn(crlf, sizeof(crlf));
+        if (crlf[0] != '\r' || crlf[1] != '\n') {
+            LOG(ERROR) << "string in command is not ended with CRLF";
+            return PARSE_ERROR_ABSOLUTELY_WRONG;
+        }
+        if (++_index == _length) {
+            break;
+        }
+    } while (true);
     args->swap(_args);
     Reset();
     return PARSE_OK;
diff --git a/test/brpc_redis_unittest.cpp b/test/brpc_redis_unittest.cpp
@@ -1423,6 +1423,25 @@ TEST_F(RedisTest, memory_allocation_limits) {
         brpc::ParseError err = parser.Consume(buf, &args, &arena);
         ASSERT_EQ(brpc::PARSE_ERROR_ABSOLUTELY_WRONG, err);
     }
+
+    {
+        // Test large command array work
+        int32_t original_limit_tmp = brpc::FLAGS_redis_max_allocation_size;
+        brpc::FLAGS_redis_max_allocation_size = 1024 * 1024;
+        brpc::RedisCommandParser parser;
+        butil::IOBuf buf;
+        int32_t large_array_size = brpc::FLAGS_redis_max_allocation_size / sizeof(butil::StringPiece);
+        std::string large_array_cmd = "*" + std::to_string(large_array_size) + "\r\n";
+        for(int i = 0; i < large_array_size; i++){
+            large_array_cmd.append("$1\r\n1\r\n");
+        }
+        buf.append(large_array_cmd);
+
+        std::vector<butil::StringPiece> args;
+        brpc::ParseError err = parser.Consume(buf, &args, &arena);
+        ASSERT_EQ(brpc::PARSE_OK, err);
+        brpc::FLAGS_redis_max_allocation_size = original_limit_tmp;
+    }
     
     // Test valid cases within limits
     {
