feat: introduce configurable flags for I/O limits

Introduced `max_multi_bulk_len` as a max limit when parsing RESP arrays
as well as `max_client_iobuf_len` as a max limit on the iobuf used to
read from a socket.

Signed-off-by: Roman Gershman <[email protected]>

Author: romange

src/facade/dragonfly_connection.cc

@@ -21,30 +21,38 @@
 #include "util/tls/tls_socket.h"
 #endif
 
+using namespace std;
+
 ABSL_FLAG(bool, tcp_nodelay, false,
           "Configures dragonfly connections with socket option TCP_NODELAY");
 ABSL_FLAG(bool, primary_port_http_enabled, true,
           "If true allows accessing http console on main TCP port");
 
-ABSL_FLAG(
-    std::uint16_t, admin_port, 0,
-    "If set, would enable admin access to console on the assigned port. This supports both HTTP "
-    "and RESP protocols");
-ABSL_FLAG(std::string, admin_bind, "",
+ABSL_FLAG(uint16_t, admin_port, 0,
+          "If set, would enable admin access to console on the assigned port. "
+          "This supports both HTTP and RESP protocols");
+
+ABSL_FLAG(string, admin_bind, "",
           "If set, the admin consol TCP connection would be bind the given address. "
-          "This supports both HTTP and RESP "
-          "protocols");
+          "This supports both HTTP and RESP protocols");
 
-ABSL_FLAG(std::uint64_t, request_cache_limit, 1ULL << 26,  // 64MB
+ABSL_FLAG(uint64_t, request_cache_limit, 1ULL << 26,  // 64MB
           "Amount of memory to use for request cache in bytes - per IO thread.");
 
 ABSL_FLAG(bool, no_tls_on_admin_port, false, "Allow non-tls connections on admin port");
 
-ABSL_FLAG(std::uint64_t, pipeline_squash, 0,
+ABSL_FLAG(uint64_t, pipeline_squash, 0,
           "Number of queued pipelined commands above which squashing is enabled, 0 means disabled");
 
+// When changing this constant, also update `test_large_cmd` test in connection_test.py.
+ABSL_FLAG(uint32_t, max_multi_bulk_len, 1u << 16,
+          "Maximum multi-bulk (array) length that is "
+          "allowed to be accepted when parsing RESP protocol");
+
+ABSL_FLAG(size_t, max_client_iobuf_len, 1u << 16,
+          "Maximum io buffer length that is used to read client requests.");
+
 using namespace util;
-using namespace std;
 using nonstd::make_unexpected;
 
 namespace facade {
@@ -79,8 +87,6 @@ bool MatchHttp11Line(string_view line) {
 }
 
 constexpr size_t kMinReadSize = 256;
-constexpr size_t kMaxReadSize = 64_KB;
-
 constexpr size_t kMaxDispatchQMemory = 5_MB;
 
 thread_local uint32_t free_req_release_weight = 0;
@@ -235,7 +241,7 @@ Connection::Connection(Protocol protocol, util::HttpListenerBase* http_listener,
 
   switch (protocol) {
     case Protocol::REDIS:
-      redis_parser_.reset(new RedisParser);
+      redis_parser_.reset(new RedisParser(absl::GetFlag(FLAGS_max_multi_bulk_len)));
       break;
     case Protocol::MEMCACHE:
       memcache_parser_.reset(new MemcacheParser);
@@ -676,6 +682,8 @@ auto Connection::IoLoop(util::FiberSocketBase* peer, SinkReplyBuilder* orig_buil
   error_code ec;
   ParserStatus parse_status = OK;
 
+  size_t max_iobfuf_len = absl::GetFlag(FLAGS_max_client_iobuf_len);
+
   do {
     FetchBuilderStats(stats_, orig_builder);
 
@@ -709,13 +717,13 @@ auto Connection::IoLoop(util::FiberSocketBase* peer, SinkReplyBuilder* orig_buil
       parse_status = OK;
 
       size_t capacity = io_buf_.Capacity();
-      if (capacity < kMaxReadSize) {
+      if (capacity < max_iobfuf_len) {
         size_t parser_hint = 0;
         if (redis_parser_)
           parser_hint = redis_parser_->parselen_hint();  // Could be done for MC as well.
 
         if (parser_hint > capacity) {
-          io_buf_.Reserve(std::min(kMaxReadSize, parser_hint));
+          io_buf_.Reserve(std::min(max_iobfuf_len, parser_hint));
         }
 
         if (io_buf_.AppendLen() < 64u) {

src/facade/memcache_parser.cc

@@ -154,8 +154,8 @@ auto MP::Parse(string_view str, uint32_t* consumed, Command* cmd) -> Result {
     return UNKNOWN_CMD;
   }
 
-  if (cmd->type <= CAS) {  // Store command
-    if (num_tokens < 5 || tokens[1].size() > 250) {
+  if (cmd->type <= CAS) {                            // Store command
+    if (num_tokens < 5 || tokens[1].size() > 250) {  // key length limit
       return MP::PARSE_ERROR;
     }
 

src/facade/redis_parser.cc

@@ -11,14 +11,6 @@ namespace facade {
 
 using namespace std;
 
-namespace {
-
-// When changing this constant, also update `test_large_cmd` test in connection_test.py.
-constexpr int kMaxArrayLen = 65536;
-constexpr int64_t kMaxBulkLen = 256 * (1ul << 20);  // 256MB.
-
-}  // namespace
-
 auto RedisParser::Parse(Buffer str, uint32_t* consumed, RespExpr::Vec* res) -> Result {
   *consumed = 0;
   res->clear();
@@ -251,8 +243,8 @@ auto RedisParser::ConsumeArrayLen(Buffer str) -> Result {
     case BAD_INT:
       return BAD_ARRAYLEN;
     case OK:
-      if (len < -1 || len > kMaxArrayLen) {
-        LOG_IF(WARNING, len > kMaxArrayLen) << "Multi bulk len is too big " << len;
+      if (len < -1 || len > max_arr_len_) {
+        LOG_IF(WARNING, len > max_arr_len_) << "Multibulk len is too large " << len;
 
         return BAD_ARRAYLEN;
       }

src/facade/redis_parser.h

@@ -22,10 +22,13 @@ namespace facade {
  */
 class RedisParser {
  public:
+  constexpr static long kMaxBulkLen = 256 * (1ul << 20);  // 256MB.
+
   enum Result { OK, INPUT_PENDING, BAD_ARRAYLEN, BAD_BULKLEN, BAD_STRING, BAD_INT, BAD_DOUBLE };
   using Buffer = RespExpr::Buffer;
 
-  explicit RedisParser(bool server_mode = true) : server_mode_(server_mode) {
+  explicit RedisParser(uint32_t max_arr_len = UINT32_MAX, bool server_mode = true)
+      : max_arr_len_(max_arr_len), server_mode_(server_mode) {
   }
 
   /**
@@ -98,6 +101,8 @@ class RedisParser {
   using BlobPtr = std::unique_ptr<uint8_t[]>;
   std::vector<BlobPtr> buf_stash_;
   RespVec* cached_expr_ = nullptr;
+  uint32_t max_arr_len_;
+
   bool is_broken_token_ = false;
   bool server_mode_ = true;
 };

src/server/protocol_client.cc

@@ -425,7 +425,8 @@ error_code ProtocolClient::SendCommandAndReadResponse(string_view command) {
 }
 
 void ProtocolClient::ResetParser(bool server_mode) {
-  parser_.reset(new RedisParser(server_mode));
+  // We accept any length for the parser because it has been approved by the master.
+  parser_.reset(new RedisParser(UINT32_MAX, server_mode));
 }
 
 uint64_t ProtocolClient::LastIoTime() const {

src/server/test_utils.cc

@@ -423,7 +423,7 @@ RespVec BaseFamilyTest::TestConnWrapper::ParseResponse(bool fully_consumed) {
   auto s_copy = s;
 
   uint32_t consumed = 0;
-  parser_.reset(new RedisParser{false});  // Client mode.
+  parser_.reset(new RedisParser{UINT32_MAX, false});  // Client mode.
   RespVec res;
   RedisParser::Result st = parser_->Parse(buf, &consumed, &res);