From f68e1ef7e3139340e5eef13941a1c6080ed30527 Mon Sep 17 00:00:00 2001 From: Kostas Kyrimis Date: Mon, 6 Nov 2023 11:27:46 +0200 Subject: [PATCH] fix(memcached): parsing multi key get command (#2122) * remove limit of 8 keys per command * refactor (small) of the parsing logic * add test --- src/facade/memcache_parser.cc | 63 ++++++++++++------------------ src/facade/memcache_parser.h | 2 - src/facade/memcache_parser_test.cc | 18 +++++++++ src/server/main_service.cc | 3 +- 4 files changed, 45 insertions(+), 41 deletions(-) diff --git a/src/facade/memcache_parser.cc b/src/facade/memcache_parser.cc index e17bb0139..2e1852146 100644 --- a/src/facade/memcache_parser.cc +++ b/src/facade/memcache_parser.cc @@ -4,9 +4,13 @@ #include "facade/memcache_parser.h" #include +#include #include #include +#include +#include +#include "base/logging.h" #include "base/stl_util.h" namespace facade { @@ -32,7 +36,10 @@ MP::CmdType From(string_view token) { return it->second; } -MP::Result ParseStore(const std::string_view* tokens, unsigned num_tokens, MP::Command* res) { +using TokensView = absl::Span; + +MP::Result ParseStore(TokensView tokens, MP::Command* res) { + const size_t num_tokens = tokens.size(); unsigned opt_pos = 3; if (res->type == MP::CAS) { if (num_tokens <= opt_pos) @@ -63,8 +70,9 @@ MP::Result ParseStore(const std::string_view* tokens, unsigned num_tokens, MP::C return MP::OK; } -MP::Result ParseValueless(const std::string_view* tokens, unsigned num_tokens, MP::Command* res) { - unsigned key_pos = 0; +MP::Result ParseValueless(TokensView tokens, MP::Command* res) { + const size_t num_tokens = tokens.size(); + size_t key_pos = 0; if (res->type == MP::GAT || res->type == MP::GATS) { if (!absl::SimpleAtoi(tokens[0], &res->expire_ts)) { return MP::BAD_INT; @@ -103,52 +111,29 @@ MP::Result ParseValueless(const std::string_view* tokens, unsigned num_tokens, M auto MP::Parse(string_view str, uint32_t* consumed, Command* cmd) -> Result { cmd->no_reply = false; // re-initialize - auto pos = str.find('\n'); + auto pos = str.find("\r\n"); *consumed = 0; if (pos == string_view::npos) { - // TODO: it's over simplified since we may process GET/GAT command that is not limited to - // 300 characters. - return str.size() > 300 ? PARSE_ERROR : INPUT_PENDING; + return INPUT_PENDING; } if (pos == 0) { return PARSE_ERROR; } - *consumed = pos + 1; + *consumed = pos + 2; + + std::string_view tokens_expression = str.substr(0, pos); // cas [noreply]\r\n // get *\r\n - string_view tokens[8]; - unsigned num_tokens = 0; - uint32_t cur = 0; + absl::InlinedVector tokens = + absl::StrSplit(tokens_expression, ' ', absl::SkipWhitespace()); - while (cur < pos && str[cur] == ' ') - ++cur; - - uint32_t s = cur; - for (; cur <= pos; ++cur) { - if (absl::ascii_isspace(str[cur])) { - if (cur != s) { - tokens[num_tokens++] = str.substr(s, cur - s); - if (num_tokens == ABSL_ARRAYSIZE(tokens)) { - ++cur; - s = cur; - break; - } - } - s = cur + 1; - } - } + const size_t num_tokens = tokens.size(); if (num_tokens == 0) return PARSE_ERROR; - while (cur < pos - 1) { - if (str[cur] != ' ') - return PARSE_ERROR; - ++cur; - } - cmd->type = From(tokens[0]); if (cmd->type == INVALID) { return UNKNOWN_CMD; @@ -159,19 +144,21 @@ auto MP::Parse(string_view str, uint32_t* consumed, Command* cmd) -> Result { return MP::PARSE_ERROR; } - // memcpy(single_key_, tokens[0].data(), tokens[0].size()); // we copy the key cmd->key = string_view{tokens[1].data(), tokens[1].size()}; - return ParseStore(tokens + 2, num_tokens - 2, cmd); + TokensView tokens_view{tokens.begin() + 2, num_tokens - 2}; + return ParseStore(tokens_view, cmd); } if (num_tokens == 1) { - if (base::_in(cmd->type, {MP::STATS, MP::FLUSHALL, MP::QUIT, MP::VERSION})) + if (base::_in(cmd->type, {MP::STATS, MP::FLUSHALL, MP::QUIT, MP::VERSION})) { return MP::OK; + } return MP::PARSE_ERROR; } - return ParseValueless(tokens + 1, num_tokens - 1, cmd); + TokensView tokens_view{tokens.begin() + 1, num_tokens - 1}; + return ParseValueless(tokens_view, cmd); }; } // namespace facade diff --git a/src/facade/memcache_parser.h b/src/facade/memcache_parser.h index 1b433cf8a..4f403ca44 100644 --- a/src/facade/memcache_parser.h +++ b/src/facade/memcache_parser.h @@ -72,8 +72,6 @@ class MemcacheParser { } Result Parse(std::string_view str, uint32_t* consumed, Command* res); - - private: }; } // namespace facade diff --git a/src/facade/memcache_parser_test.cc b/src/facade/memcache_parser_test.cc index 1d482ef23..9e20d4049 100644 --- a/src/facade/memcache_parser_test.cc +++ b/src/facade/memcache_parser_test.cc @@ -128,4 +128,22 @@ TEST_F(MCParserNoreplyTest, Other) { RunTest("flush_all\r\n", false); } +TEST_F(MCParserNoreplyTest, LargeGetRequest) { + std::string large_request = "get"; + for (size_t i = 0; i < 100; ++i) { + absl::StrAppend(&large_request, " mykey", i, " "); + } + absl::StrAppend(&large_request, "\r\n"); + + RunTest(large_request, false); + + EXPECT_EQ(cmd_.type, MemcacheParser::CmdType::GET); + EXPECT_EQ(cmd_.key, "mykey0"); + const auto& keys = cmd_.keys_ext; + EXPECT_TRUE(std::all_of(keys.begin(), keys.end(), [](const auto& elem) { + static size_t i = 1; + return elem == absl::StrCat("mykey", i++); + })); +} + } // namespace facade diff --git a/src/server/main_service.cc b/src/server/main_service.cc index 6b9d0ea80..393598425 100644 --- a/src/server/main_service.cc +++ b/src/server/main_service.cc @@ -1145,7 +1145,8 @@ class ReplyGuard { const bool is_one_of = absl::flat_hash_set({"REPLCONF", "DFLY"}).contains(cid_name); auto* maybe_mcache = dynamic_cast(cntx->reply_builder()); - const bool is_no_reply_memcache = maybe_mcache && maybe_mcache->NoReply(); + const bool is_no_reply_memcache = + maybe_mcache && (maybe_mcache->NoReply() || cid_name == "QUIT"); const bool should_dcheck = !is_one_of && !is_script && !is_no_reply_memcache; if (should_dcheck) { builder_ = cntx->reply_builder();