Limit Redis inline command allocation

src/brpc/redis_command.cpp

@@ -26,6 +26,23 @@
 namespace {
 
 const size_t CTX_WIDTH = 5;
+const size_t CRLF_NOT_FOUND = (size_t)-1;
+
+size_t FindCRLF(const butil::IOBuf& buf, size_t max_scan_size) {
+    butil::IOBufBytesIterator it(buf);
+    bool prev_cr = false;
+    size_t pos = 0;
+    while (it && pos < max_scan_size) {
+        const char c = static_cast<char>(*it);
+        if (prev_cr && c == '\n') {
+            return pos - 1;
+        }
+        prev_cr = (c == '\r');
+        ++it;
+        ++pos;
+    }
+    return CRLF_NOT_FOUND;
+}
 
 } // namespace
 
@@ -425,7 +442,35 @@ RedisCommandConsumeState RedisCommandParser::ConsumeImpl(butil::IOBuf& buf,
                 return CONSUME_STATE_ERROR;
             }
         }
-        const size_t buf_size = buf.size();
+        const size_t max_inline_size =
+            FLAGS_redis_max_allocation_size > 0 ?
+            static_cast<size_t>(FLAGS_redis_max_allocation_size) : 0;
+        const size_t crlf_pos = FindCRLF(buf, max_inline_size + 2);
+        if (crlf_pos == CRLF_NOT_FOUND) {
+            if (buf.size() <= max_inline_size) {
+                *err = PARSE_ERROR_NOT_ENOUGH_DATA;
+                return CONSUME_STATE_ERROR;
+            }
+            if (buf.size() == max_inline_size + 1) {
+                char last_char = '\0';
+                buf.copy_to(&last_char, 1, max_inline_size);
+                if (last_char == '\r') {
+                    *err = PARSE_ERROR_NOT_ENOUGH_DATA;
+                    return CONSUME_STATE_ERROR;
+                }
+            }
+            LOG(ERROR) << "inline command exceeds max allocation size! max="
+                       << FLAGS_redis_max_allocation_size << ", actually=" << buf.size();
+            *err = PARSE_ERROR_ABSOLUTELY_WRONG;
+            return CONSUME_STATE_ERROR;
+        }
+        if (crlf_pos > max_inline_size) {
+            LOG(ERROR) << "inline command exceeds max allocation size! max="
+                       << FLAGS_redis_max_allocation_size << ", actually=" << crlf_pos;
+            *err = PARSE_ERROR_ABSOLUTELY_WRONG;
+            return CONSUME_STATE_ERROR;
+        }
+        const size_t buf_size = crlf_pos;
         const auto copy_str = static_cast<char *>(arena->allocate(buf_size + 1));
         // arena->allocate() may return NULL on allocation failure
         if (copy_str == NULL) {
@@ -439,11 +484,6 @@ RedisCommandConsumeState RedisCommandParser::ConsumeImpl(butil::IOBuf& buf,
             return CONSUME_STATE_ERROR;
         }
         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
-            *err = PARSE_ERROR_NOT_ENOUGH_DATA;
-            return CONSUME_STATE_ERROR;
-        }
         size_t offset = 0;
         while (offset < crlf_pos && copy_str[offset] != ' ') {
             ++offset;

test/brpc_redis_unittest.cpp

@@ -1512,6 +1512,57 @@ TEST_F(RedisTest, memory_allocation_limits) {
         ASSERT_EQ(brpc::PARSE_ERROR_ABSOLUTELY_WRONG, err);
     }
 
+    {
+        // Test inline command exceeding limit before CRLF arrives.
+        brpc::RedisCommandParser parser;
+        butil::IOBuf buf;
+        std::string large_inline_cmd = "get ";
+        large_inline_cmd.append(brpc::FLAGS_redis_max_allocation_size + 1, 'k');
+        buf.append(large_inline_cmd);
+
+        std::vector<butil::StringPiece> args;
+        brpc::ParseError err = parser.Consume(buf, &args, &arena);
+        ASSERT_EQ(brpc::PARSE_ERROR_ABSOLUTELY_WRONG, err);
+    }
+
+    {
+        // A command line exactly at the limit may have CRLF split across reads.
+        brpc::RedisCommandParser parser;
+        butil::IOBuf buf;
+        std::string boundary_inline_cmd = "get ";
+        boundary_inline_cmd.append(brpc::FLAGS_redis_max_allocation_size -
+                                   boundary_inline_cmd.size(), 'k');
+        boundary_inline_cmd.push_back('\r');
+        buf.append(boundary_inline_cmd);
+
+        std::vector<butil::StringPiece> args;
+        brpc::ParseError err = parser.Consume(buf, &args, &arena);
+        ASSERT_EQ(brpc::PARSE_ERROR_NOT_ENOUGH_DATA, err);
+
+        buf.push_back('\n');
+        err = parser.Consume(buf, &args, &arena);
+        ASSERT_EQ(brpc::PARSE_OK, err);
+        ASSERT_EQ(2, (int)args.size());
+        ASSERT_EQ("get", args[0].as_string());
+        ASSERT_EQ(brpc::FLAGS_redis_max_allocation_size - 4, (int)args[1].size());
+    }
+
+    {
+        // Test that only the current inline command line is limited and copied.
+        brpc::RedisCommandParser parser;
+        butil::IOBuf buf;
+        std::string pipelined_inline_cmd = "PING\r\nget ";
+        pipelined_inline_cmd.append(brpc::FLAGS_redis_max_allocation_size + 1, 'k');
+        buf.append(pipelined_inline_cmd);
+
+        std::vector<butil::StringPiece> args;
+        brpc::ParseError err = parser.Consume(buf, &args, &arena);
+        ASSERT_EQ(brpc::PARSE_OK, err);
+        ASSERT_EQ(1, (int)args.size());
+        ASSERT_EQ("ping", args[0].as_string());
+        ASSERT_EQ(pipelined_inline_cmd.size() - 6, buf.size());
+    }
+
     {
         // Test large command array work
         int32_t original_limit_tmp = brpc::FLAGS_redis_max_allocation_size;