From f44be01a0412986108cb928f33be0bf50ff2bf27 Mon Sep 17 00:00:00 2001 From: caiyu <15260903118@163.com> Date: Thu, 12 Jun 2025 15:23:59 +0800 Subject: [PATCH] Fix the error bugs in getrange and setrange in pika 3.5 version --- conf/pika.conf | 5 + include/pika_cache.h | 2 +- include/pika_conf.h | 10 ++ src/pika_cache.cc | 265 +++++++++++++++++++------------ src/pika_conf.cc | 4 + src/pika_kv.cc | 11 ++ tests/integration/string_test.go | 18 +++ 7 files changed, 211 insertions(+), 104 deletions(-) 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_cache.h b/include/pika_cache.h index a6960cef02..11d5754a65 100644 --- a/include/pika_cache.h +++ b/include/pika_cache.h @@ -133,7 +133,7 @@ class PikaCache : public pstd::noncopyable, public std::enable_shared_from_this< rocksdb::Status RPushnx(std::string& key, std::vector &values, int64_t ttl); rocksdb::Status RPushnxWithoutTTL(std::string& key, std::vector &values); rocksdb::Status LPushIfKeyExist(std::string& key, std::vector &values); - + // Set Commands rocksdb::Status SAdd(std::string& key, std::vector& members); rocksdb::Status SAddIfKeyExist(std::string& key, std::vector& members); 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 7bb9e89916..db0193657f 100644 --- a/src/pika_cache.cc +++ b/src/pika_cache.cc @@ -141,54 +141,69 @@ std::map PikaCache::TTL(std::string &key, std::map ret; int64_t timestamp = 0; - std::string CacheKeyPrefixK = PCacheKeyPrefixK + key; - int cache_indexk = CacheIndex(CacheKeyPrefixK); - s = caches_[cache_indexk]->TTL(CacheKeyPrefixK, ×tamp); - if (s.ok() || s.IsNotFound()) { - ret[storage::DataType::kStrings] = timestamp; - } else if (!s.IsNotFound()) { - ret[storage::DataType::kStrings] = -3; - (*type_status)[storage::DataType::kStrings] = s; - } - - std::string CacheKeyPrefixH = PCacheKeyPrefixH + key; - int cache_indexh = CacheIndex(CacheKeyPrefixH); - s = caches_[cache_indexh]->TTL(CacheKeyPrefixH, ×tamp); - if (s.ok() || s.IsNotFound()) { - ret[storage::DataType::kHashes] = timestamp; - } else if (!s.IsNotFound()) { - ret[storage::DataType::kHashes] = -3; - (*type_status)[storage::DataType::kHashes] = s; - } - - std::string CacheKeyPrefixL = PCacheKeyPrefixL + key; - int cache_indexl = CacheIndex(CacheKeyPrefixL); - s = caches_[cache_indexl]->TTL(CacheKeyPrefixL, ×tamp); - if (s.ok() || s.IsNotFound()) { - ret[storage::DataType::kLists] = timestamp; - } else if (!s.IsNotFound()) { - ret[storage::DataType::kLists] = -3; - (*type_status)[storage::DataType::kLists] = s; - } - - std::string CacheKeyPrefixS = PCacheKeyPrefixS + key; - int cache_indexs = CacheIndex(CacheKeyPrefixS); - s = caches_[cache_indexs]->TTL(CacheKeyPrefixS, ×tamp); - if (s.ok() || s.IsNotFound()) { - ret[storage::DataType::kSets] = timestamp; - } else if (!s.IsNotFound()) { - ret[storage::DataType::kSets] = -3; - (*type_status)[storage::DataType::kSets] = s; - } - - std::string CacheKeyPrefixZ = PCacheKeyPrefixZ + key; - int cache_indexz = CacheIndex(CacheKeyPrefixZ); - s = caches_[cache_indexz]->TTL(CacheKeyPrefixZ, ×tamp); - if (s.ok() || s.IsNotFound()) { - ret[storage::DataType::kZSets] = timestamp; - } else if (!s.IsNotFound()) { - ret[storage::DataType::kZSets] = -3; - (*type_status)[storage::DataType::kZSets] = s; + { + std::string CacheKeyPrefixK = PCacheKeyPrefixK + key; + int cache_indexk = CacheIndex(CacheKeyPrefixK); + std::lock_guard lm(*cache_mutexs_[cache_indexk]); + s = caches_[cache_indexk]->TTL(CacheKeyPrefixK, ×tamp); + if (s.ok() || s.IsNotFound()) { + ret[storage::DataType::kStrings] = timestamp; + } else if (!s.IsNotFound()) { + ret[storage::DataType::kStrings] = -3; + (*type_status)[storage::DataType::kStrings] = s; + } + } + + { + std::string CacheKeyPrefixH = PCacheKeyPrefixH + key; + int cache_indexh = CacheIndex(CacheKeyPrefixH); + std::lock_guard lm(*cache_mutexs_[cache_indexh]); + s = caches_[cache_indexh]->TTL(CacheKeyPrefixH, ×tamp); + if (s.ok() || s.IsNotFound()) { + ret[storage::DataType::kHashes] = timestamp; + } else if (!s.IsNotFound()) { + ret[storage::DataType::kHashes] = -3; + (*type_status)[storage::DataType::kHashes] = s; + } + } + + { + std::string CacheKeyPrefixL = PCacheKeyPrefixL + key; + int cache_indexl = CacheIndex(CacheKeyPrefixL); + std::lock_guard lm(*cache_mutexs_[cache_indexl]); + s = caches_[cache_indexl]->TTL(CacheKeyPrefixL, ×tamp); + if (s.ok() || s.IsNotFound()) { + ret[storage::DataType::kLists] = timestamp; + } else if (!s.IsNotFound()) { + ret[storage::DataType::kLists] = -3; + (*type_status)[storage::DataType::kLists] = s; + } + } + + { + std::string CacheKeyPrefixS = PCacheKeyPrefixS + key; + int cache_indexs = CacheIndex(CacheKeyPrefixS); + std::lock_guard lm(*cache_mutexs_[cache_indexs]); + s = caches_[cache_indexs]->TTL(CacheKeyPrefixS, ×tamp); + if (s.ok() || s.IsNotFound()) { + ret[storage::DataType::kSets] = timestamp; + } else if (!s.IsNotFound()) { + ret[storage::DataType::kSets] = -3; + (*type_status)[storage::DataType::kSets] = s; + } + } + + { + std::string CacheKeyPrefixZ = PCacheKeyPrefixZ + key; + int cache_indexz = CacheIndex(CacheKeyPrefixZ); + std::lock_guard lm(*cache_mutexs_[cache_indexz]); + s = caches_[cache_indexz]->TTL(CacheKeyPrefixZ, ×tamp); + if (s.ok() || s.IsNotFound()) { + ret[storage::DataType::kZSets] = timestamp; + } else if (!s.IsNotFound()) { + ret[storage::DataType::kZSets] = -3; + (*type_status)[storage::DataType::kZSets] = s; + } } return ret; } @@ -226,68 +241,83 @@ Status PikaCache::GetType(const std::string& key, bool single, std::vectorGet(CacheKeyPrefixK, &value); - if (s.ok()) { - types.emplace_back("string"); - } else if (!s.IsNotFound()) { - return s; - } - if (single && !types.empty()) { - return s; + { + std::string CacheKeyPrefixK = PCacheKeyPrefixK + key; + int cache_indexk = CacheIndex(CacheKeyPrefixK); + std::lock_guard lm(*cache_mutexs_[cache_indexk]); + s = caches_[cache_indexk]->Get(CacheKeyPrefixK, &value); + if (s.ok()) { + types.emplace_back("string"); + } else if (!s.IsNotFound()) { + return s; + } + if (single && !types.empty()) { + return s; + } } - uint64_t hashes_len = 0; - std::string CacheKeyPrefixH = PCacheKeyPrefixH + key; - int cache_indexh = CacheIndex(CacheKeyPrefixH); - s = caches_[cache_indexh]->HLen(CacheKeyPrefixH, &hashes_len); - if (s.ok() && hashes_len != 0) { - types.emplace_back("hash"); - } else if (!s.IsNotFound()) { - return s; - } - if (single && !types.empty()) { - return s; + { + uint64_t hashes_len = 0; + std::string CacheKeyPrefixH = PCacheKeyPrefixH + key; + int cache_indexh = CacheIndex(CacheKeyPrefixH); + std::lock_guard lm(*cache_mutexs_[cache_indexh]); + s = caches_[cache_indexh]->HLen(CacheKeyPrefixH, &hashes_len); + if (s.ok() && hashes_len != 0) { + types.emplace_back("hash"); + } else if (!s.IsNotFound()) { + return s; + } + if (single && !types.empty()) { + return s; + } } - uint64_t lists_len = 0; - std::string CacheKeyPrefixL = PCacheKeyPrefixL + key; - int cache_indexl = CacheIndex(CacheKeyPrefixL); - s = caches_[cache_indexl]->LLen(CacheKeyPrefixL, &lists_len); - if (s.ok() && lists_len != 0) { - types.emplace_back("list"); - } else if (!s.IsNotFound()) { - return s; - } - if (single && !types.empty()) { - return s; + { + uint64_t lists_len = 0; + std::string CacheKeyPrefixL = PCacheKeyPrefixL + key; + int cache_indexl = CacheIndex(CacheKeyPrefixL); + std::lock_guard lm(*cache_mutexs_[cache_indexl]); + s = caches_[cache_indexl]->LLen(CacheKeyPrefixL, &lists_len); + if (s.ok() && lists_len != 0) { + types.emplace_back("list"); + } else if (!s.IsNotFound()) { + return s; + } + if (single && !types.empty()) { + return s; + } } - uint64_t zsets_size = 0; - std::string CacheKeyPrefixZ = PCacheKeyPrefixZ + key; - int cache_indexz = CacheIndex(CacheKeyPrefixZ); - s = caches_[cache_indexz]->ZCard(CacheKeyPrefixZ, &zsets_size); - if (s.ok() && zsets_size != 0) { - types.emplace_back("zset"); - } else if (!s.IsNotFound()) { - return s; - } - if (single && !types.empty()) { - return s; + { + uint64_t zsets_size = 0; + std::string CacheKeyPrefixZ = PCacheKeyPrefixZ + key; + int cache_indexz = CacheIndex(CacheKeyPrefixZ); + std::lock_guard lm(*cache_mutexs_[cache_indexz]); + s = caches_[cache_indexz]->ZCard(CacheKeyPrefixZ, &zsets_size); + if (s.ok() && zsets_size != 0) { + types.emplace_back("zset"); + } else if (!s.IsNotFound()) { + return s; + } + if (single && !types.empty()) { + return s; + } } - uint64_t sets_size = 0; - std::string CacheKeyPrefixS = PCacheKeyPrefixS + key; - int cache_indexs = CacheIndex(CacheKeyPrefixS); - s = caches_[cache_indexs]->SCard(CacheKeyPrefixS, &sets_size); - if (s.ok() && sets_size != 0) { - types.emplace_back("set"); - } else if (!s.IsNotFound()) { - return s; - } - if (single && types.empty()) { - types.emplace_back("none"); + { + uint64_t sets_size = 0; + std::string CacheKeyPrefixS = PCacheKeyPrefixS + key; + int cache_indexs = CacheIndex(CacheKeyPrefixS); + std::lock_guard lm(*cache_mutexs_[cache_indexs]); + s = caches_[cache_indexs]->SCard(CacheKeyPrefixS, &sets_size); + if (s.ok() && sets_size != 0) { + types.emplace_back("set"); + } else if (!s.IsNotFound()) { + return s; + } + if (single && types.empty()) { + types.emplace_back("none"); + } } return Status::OK(); } @@ -410,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())