diff --git a/conf/pika.conf b/conf/pika.conf index 446d91cebe..97d171d419 100644 --- a/conf/pika.conf +++ b/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) diff --git a/include/pika_conf.h b/include/pika_conf.h index 51ab3a1a27..80d5abe8f0 100644 --- a/include/pika_conf.h +++ b/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; diff --git a/src/pika_cache.cc b/src/pika_cache.cc index 76edf80155..db0193657f 100644 --- a/src/pika_cache.cc +++ b/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) { diff --git a/src/pika_conf.cc b/src/pika_conf.cc index f0435b4c58..acf90e3a6f 100644 --- a/src/pika_conf.cc +++ b/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) { diff --git a/src/pika_kv.cc b/src/pika_kv.cc index 321cd9f546..1c1abdd4cf 100644 --- a/src/pika_kv.cc +++ b/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(value_.size())) { + res_.SetRes(CmdRes::kErrOther, "string exceeds maximum allowed size (proto-max-bulk-len)"); + return; + } } void SetrangeCmd::Do() { diff --git a/tests/integration/string_test.go b/tests/integration/string_test.go index 9a8e4c4a69..5f699549bc 100644 --- a/tests/integration/string_test.go +++ b/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())