1
0
Fork 0
mirror of https://github.com/dragonflydb/dragonfly.git synced 2024-12-14 11:58:02 +00:00

fix(memcached): parsing multi key get command (#2122)

* remove limit of 8 keys per command
* refactor (small) of the parsing logic
* add test
This commit is contained in:
Kostas Kyrimis 2023-11-06 11:27:46 +02:00 committed by GitHub
parent efeae543a1
commit f68e1ef7e3
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 45 additions and 41 deletions

View file

@ -4,9 +4,13 @@
#include "facade/memcache_parser.h" #include "facade/memcache_parser.h"
#include <absl/container/flat_hash_map.h> #include <absl/container/flat_hash_map.h>
#include <absl/container/inlined_vector.h>
#include <absl/strings/ascii.h> #include <absl/strings/ascii.h>
#include <absl/strings/numbers.h> #include <absl/strings/numbers.h>
#include <absl/strings/str_split.h>
#include <absl/types/span.h>
#include "base/logging.h"
#include "base/stl_util.h" #include "base/stl_util.h"
namespace facade { namespace facade {
@ -32,7 +36,10 @@ MP::CmdType From(string_view token) {
return it->second; return it->second;
} }
MP::Result ParseStore(const std::string_view* tokens, unsigned num_tokens, MP::Command* res) { using TokensView = absl::Span<std::string_view>;
MP::Result ParseStore(TokensView tokens, MP::Command* res) {
const size_t num_tokens = tokens.size();
unsigned opt_pos = 3; unsigned opt_pos = 3;
if (res->type == MP::CAS) { if (res->type == MP::CAS) {
if (num_tokens <= opt_pos) 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; return MP::OK;
} }
MP::Result ParseValueless(const std::string_view* tokens, unsigned num_tokens, MP::Command* res) { MP::Result ParseValueless(TokensView tokens, MP::Command* res) {
unsigned key_pos = 0; const size_t num_tokens = tokens.size();
size_t key_pos = 0;
if (res->type == MP::GAT || res->type == MP::GATS) { if (res->type == MP::GAT || res->type == MP::GATS) {
if (!absl::SimpleAtoi(tokens[0], &res->expire_ts)) { if (!absl::SimpleAtoi(tokens[0], &res->expire_ts)) {
return MP::BAD_INT; 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 { auto MP::Parse(string_view str, uint32_t* consumed, Command* cmd) -> Result {
cmd->no_reply = false; // re-initialize cmd->no_reply = false; // re-initialize
auto pos = str.find('\n'); auto pos = str.find("\r\n");
*consumed = 0; *consumed = 0;
if (pos == string_view::npos) { if (pos == string_view::npos) {
// TODO: it's over simplified since we may process GET/GAT command that is not limited to return INPUT_PENDING;
// 300 characters.
return str.size() > 300 ? PARSE_ERROR : INPUT_PENDING;
} }
if (pos == 0) { if (pos == 0) {
return PARSE_ERROR; return PARSE_ERROR;
} }
*consumed = pos + 1; *consumed = pos + 2;
std::string_view tokens_expression = str.substr(0, pos);
// cas <key> <flags> <exptime> <bytes> <cas unique> [noreply]\r\n // cas <key> <flags> <exptime> <bytes> <cas unique> [noreply]\r\n
// get <key>*\r\n // get <key>*\r\n
string_view tokens[8]; absl::InlinedVector<std::string_view, 32> tokens =
unsigned num_tokens = 0; absl::StrSplit(tokens_expression, ' ', absl::SkipWhitespace());
uint32_t cur = 0;
while (cur < pos && str[cur] == ' ') const size_t num_tokens = tokens.size();
++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;
}
}
if (num_tokens == 0) if (num_tokens == 0)
return PARSE_ERROR; return PARSE_ERROR;
while (cur < pos - 1) {
if (str[cur] != ' ')
return PARSE_ERROR;
++cur;
}
cmd->type = From(tokens[0]); cmd->type = From(tokens[0]);
if (cmd->type == INVALID) { if (cmd->type == INVALID) {
return UNKNOWN_CMD; return UNKNOWN_CMD;
@ -159,19 +144,21 @@ auto MP::Parse(string_view str, uint32_t* consumed, Command* cmd) -> Result {
return MP::PARSE_ERROR; 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()}; 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 (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::OK;
}
return MP::PARSE_ERROR; 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 } // namespace facade

View file

@ -72,8 +72,6 @@ class MemcacheParser {
} }
Result Parse(std::string_view str, uint32_t* consumed, Command* res); Result Parse(std::string_view str, uint32_t* consumed, Command* res);
private:
}; };
} // namespace facade } // namespace facade

View file

@ -128,4 +128,22 @@ TEST_F(MCParserNoreplyTest, Other) {
RunTest("flush_all\r\n", false); 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 } // namespace facade

View file

@ -1145,7 +1145,8 @@ class ReplyGuard {
const bool is_one_of = const bool is_one_of =
absl::flat_hash_set<std::string_view>({"REPLCONF", "DFLY"}).contains(cid_name); absl::flat_hash_set<std::string_view>({"REPLCONF", "DFLY"}).contains(cid_name);
auto* maybe_mcache = dynamic_cast<MCReplyBuilder*>(cntx->reply_builder()); auto* maybe_mcache = dynamic_cast<MCReplyBuilder*>(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; const bool should_dcheck = !is_one_of && !is_script && !is_no_reply_memcache;
if (should_dcheck) { if (should_dcheck) {
builder_ = cntx->reply_builder(); builder_ = cntx->reply_builder();