fix: Specia Redis directives can cause pika service to crash (#3100)

* Specia Redis directives can cause pika service to crash

* Special Redis directives can cause pika service to crash

* Special Redis directives can cause pika service to crash

* Add support for reading proto-max-bulk-len from configuration file

* optimize boundary checks and empty string handling in GetRange function of pika_cache

* Complete the required items for pika_conf.cc and pika_conf.h files

---------

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

Author: YuCai18

conf/pika.conf

@@ -90,6 +90,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

@@ -167,6 +167,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_;
+  }
   int min_write_buffer_number_to_merge() {
     std::shared_lock l(rwlock_);
     return min_write_buffer_number_to_merge_;
@@ -858,6 +862,11 @@ class PikaConf : public pstd::BaseConf {
     TryPushDiffCommands("rsync-timeout-ms", std::to_string(value));
     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();
@@ -1035,6 +1044,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

@@ -364,12 +364,41 @@ 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) {
   int cache_index = CacheIndex(key);
   std::lock_guard lm(*cache_mutexs_[cache_index]);

src/pika_conf.cc

@@ -366,7 +366,10 @@ int PikaConf::Load() {
   if (write_buffer_size_ <= 0) {
     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
+  }
   GetConfInt("level0-stop-writes-trigger", &level0_stop_writes_trigger_);
   if (level0_stop_writes_trigger_ < 36) {
     level0_stop_writes_trigger_ = 36;

src/pika_kv.cc

@@ -1146,6 +1146,7 @@ void GetrangeCmd::Do() {
   std::string substr;
   STAGE_TIMER_GUARD(storage_duration_ms, true);
   s_= db_->storage()->Getrange(key_, start_, end_, &substr);
+  
   if (s_.ok() || s_.IsNotFound()) {
     res_.AppendStringLenUint64(substr.size());
     res_.AppendContent(substr);
@@ -1194,12 +1195,25 @@ void SetrangeCmd::DoInitial() {
     res_.SetRes(CmdRes::kWrongNum, kCmdNameSetrange);
     return;
   }
-  key_ = argv_[1];
+  key_ = argv_[1];  
   if (pstd::string2int(argv_[2].data(), argv_[2].size(), &offset_) == 0) {
     res_.SetRes(CmdRes::kInvalidInt);
     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

@@ -277,7 +277,7 @@ var _ = Describe("String Commands", func() {
 			Expect(getBit.Err()).NotTo(HaveOccurred())
 			Expect(getBit.Val()).To(Equal(int64(0)))
 		})
-				
+
 		It("should GetRange", func() {
 			set := client.Set(ctx, "key", "This is a string", 0)
 			Expect(set.Err()).NotTo(HaveOccurred())
@@ -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())
@@ -798,24 +816,24 @@ var _ = Describe("String Commands", func() {
 		})
 
 		It("should SetEX ten seconds", func() {
-			err := client.SetEx(ctx, "x", "y", 10 * time.Second).Err()
+			err := client.SetEx(ctx, "x", "y", 10*time.Second).Err()
 			Expect(err).NotTo(HaveOccurred())
-		
+
 			val, err := client.Get(ctx, "x").Result()
 			Expect(err).NotTo(HaveOccurred())
 			Expect(val).To(Equal("y"))
-		
+
 			time.Sleep(11 * time.Second)
 			//sleep 10 second x still exists
-		
+
 			err = client.Do(ctx, "compact").Err()
 			Expect(err).NotTo(HaveOccurred())
-		
+
 			keys, err := client.Keys(ctx, "x").Result()
 			Expect(err).NotTo(HaveOccurred())
 			Expect(keys).To(BeEmpty())
 		})
-		
+
 		It("should SetNX", func() {
 			_, err := client.Del(ctx, "key").Result()
 			Expect(err).NotTo(HaveOccurred())