Fix the error bugs in getrange and setrange in pika 3.5 version (#3106)

Co-authored-by: caiyu <[email protected]>

Author: YuCai18

conf/pika.conf

@@ -86,6 +86,11 @@ db-path : ./db/
 # Supported Units [K|M|G], write-buffer-size default unit is in [bytes].
 write-buffer-size : 256M
 
+# The maximum size of a single bulk string in Pika protocol.
+# This value is used to limit the size of a single bulk string in Pika protocol.
+# The default value is 512M.
+proto-max-bulk-len : 512M
+
 # The size of one block in arena memory allocation.
 # If <= 0, a proper value is automatically calculated.
 # (usually 1/8 of writer-buffer-size, rounded up to a multiple of 4KB)

include/pika_conf.h

@@ -124,6 +124,10 @@ class PikaConf : public pstd::BaseConf {
     std::shared_lock l(rwlock_);
     return write_buffer_size_;
   }
+  int64_t proto_max_bulk_len() {
+    std::shared_lock l(rwlock_);
+    return proto_max_bulk_len_;
+  }
   int64_t arena_block_size() {
     std::shared_lock l(rwlock_);
     return arena_block_size_;
@@ -745,6 +749,11 @@ class PikaConf : public pstd::BaseConf {
     rsync_timeout_ms_.store(value);
   }
 
+  void SetProtoMaxBulkLen(const int64_t value) {
+    std::lock_guard l(rwlock_);
+    TryPushDiffCommands("proto-max-bulk-len", std::to_string(value));
+    proto_max_bulk_len_ = value;
+  }
   int RocksDBPerfLevel() const {
     return rocksdb_perf_level_.load();
   }
@@ -907,6 +916,7 @@ class PikaConf : public pstd::BaseConf {
   int64_t least_free_disk_to_resume_ = 268435456; // 256 MB
   double min_check_resume_ratio_ = 0.7;
   int64_t write_buffer_size_ = 0;
+  int64_t proto_max_bulk_len_ = 0;
   int64_t arena_block_size_ = 0;
   int64_t slotmigrate_thread_num_ = 0;
   int64_t thread_migrate_keys_num_ = 0;

src/pika_cache.cc

@@ -440,10 +440,39 @@ Status PikaCache::Appendxx(std::string& key, std::string &value) {
   return Status::NotFound("key not exist");
 }
 
+/*
+  Added boundary checks for start and end parameters to the PikaCache::GetRange function,
+  and used the full_value variable to store the actual length of string type, 
+  avoiding excessive memory allocation by sdsnewlen.
+*/
 Status PikaCache::GetRange(std::string& key, int64_t start, int64_t end, std::string *value) {
   int cache_index = CacheIndex(key);
   std::lock_guard lm(*cache_mutexs_[cache_index]);
-  return caches_[cache_index]->GetRange(key, start, end, value);
+  std::string full_value;
+  auto s = caches_[cache_index]->Get(key, &full_value);
+  if (!s.ok()) {
+    return s;
+  }
+  int64_t strlen = full_value.size();
+
+  if (start < 0) {
+    start = strlen + start;
+  }
+  if (end < 0) {
+    end = strlen + end;
+  }
+
+  if (start < 0) start = 0;
+  if (end < 0) end = 0;
+  if (end >= strlen) end = strlen - 1;
+
+  if (start > end || strlen == 0) {
+    value->clear();
+    return Status::OK();
+  }
+
+  *value = full_value.substr(start, end - start + 1);
+  return Status::OK();
 }
 
 Status PikaCache::SetRangeIfKeyExist(std::string& key, int64_t start, std::string &value) {

src/pika_conf.cc

@@ -323,6 +323,10 @@ int PikaConf::Load() {
     write_buffer_size_ = 268435456;  // 256Mb
   }
 
+  GetConfInt64Human("proto-max-bulk-len", &proto_max_bulk_len_);
+  if (proto_max_bulk_len_ <= 0) {
+    proto_max_bulk_len_ = 512 * 1024 * 1024;  // 512MB
+  }
   // arena_block_size
   GetConfInt64Human("arena-block-size", &arena_block_size_);
   if (arena_block_size_ <= 0) {

src/pika_kv.cc

@@ -1215,6 +1215,17 @@ void SetrangeCmd::DoInitial() {
     return;
   }
   value_ = argv_[3];
+  // Read the proto-max-bulk-len parameter settings in the pika configuration file pika_conf
+  const int64_t PROTO_MAX_BULK_LEN = g_pika_conf->proto_max_bulk_len();
+  //Handle the overflow issue of offset_
+  if (offset_ < 0) {
+    res_.SetRes(CmdRes::kInvalidInt, "offset is out of range");
+    return;
+  }
+  if (offset_ > PROTO_MAX_BULK_LEN - static_cast<int64_t>(value_.size())) {
+    res_.SetRes(CmdRes::kErrOther, "string exceeds maximum allowed size (proto-max-bulk-len)");
+    return;
+  }
 }
 
 void SetrangeCmd::Do() {

tests/integration/string_test.go

@@ -300,6 +300,24 @@ var _ = Describe("String Commands", func() {
 			Expect(getRange.Val()).To(Equal("string"))
 		})
 
+		//Caiyu's test cases for GETRANGE and SETRANGE to fix bug #3092.
+		It("should not crash on huge GETRANGE", func() {
+			set := client.Set(ctx, "key1", "abc", 0)
+			Expect(set.Err()).NotTo(HaveOccurred())
+			Expect(set.Val()).To(Equal("OK"))
+			getRange1 := client.GetRange(ctx, "key1", 1, 4294967296)
+    		Expect(getRange1.Val()).To(Equal("bc"))
+    		getRange2 := client.GetRange(ctx, "key1", 1, 4294967296)
+    		Expect(getRange2.Val()).To(Equal("bc"))
+		})
+		It("should not crash on huge SETRANGE", func() {
+			set := client.Set(ctx, "key1", "abc", 0)
+			Expect(set.Err()).NotTo(HaveOccurred())
+			Expect(set.Val()).To(Equal("OK"))
+			setRange := client.SetRange(ctx, "key1", 9223372036854775757, "value2")
+			Expect(setRange.Err()).To(HaveOccurred())
+		})
+
 		It("should GetSet", func() {
 			incr := client.Incr(ctx, "key")
 			Expect(incr.Err()).NotTo(HaveOccurred())