fix: keyspace causes heap-buffer-overflow (#2749)

* fix keyspace error about heap-buffer-overflow * fix by ai review comments --------- Co-authored-by: wangshaoyi <[email protected]>
OpenAtomFoundation · Jun 20, 2024 · 6f85ab5 · 6f85ab5
1 parent 55de8b3
commit 6f85ab5
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 9 deletions.
diff --git a/src/pika_admin.cc b/src/pika_admin.cc
@@ -1190,9 +1190,9 @@ void InfoCmd::InfoKeyspace(std::string& info) {
   if (argv_.size() > 1 && strcasecmp(argv_[1].data(), kAllSection.data()) == 0) {
     tmp_stream << "# Start async statistics\r\n";
   } else if (argv_.size() == 3 && strcasecmp(argv_[1].data(), kKeyspaceSection.data()) == 0) {
-    tmp_stream << "# Start async statistics\r\n";    
+    tmp_stream << "# Start async statistics\r\n";
   } else {
-    tmp_stream << "# Use \"info keyspace 1\" to do async statistics\r\n";  
+    tmp_stream << "# Use \"info keyspace 1\" to do async statistics\r\n";
   }
   std::shared_lock rwl(g_pika_server->dbs_rw_);
   for (const auto& db_item : g_pika_server->dbs_) {
@@ -1201,7 +1201,8 @@ void InfoCmd::InfoKeyspace(std::string& info) {
       key_scan_info = db_item.second->GetKeyScanInfo();
       key_infos = key_scan_info.key_infos;
       duration = key_scan_info.duration;
-      if (key_infos.size() != (size_t)(storage::DataType::kNones)) {
+      if (key_infos.size() != (size_t)(storage::DataTypeNum)) {
+        LOG(ERROR) << "key_infos size is not equal with expected, potential data inconsistency";
         info.append("info keyspace error\r\n");
         return;
       }
@@ -1216,7 +1217,7 @@ void InfoCmd::InfoKeyspace(std::string& info) {
         tmp_stream << "# Duration: " << std::to_string(duration) + "s"
                    << "\r\n";
       }
-      
+
       tmp_stream << db_name << " Strings_keys=" << key_infos[0].keys << ", expires=" << key_infos[0].expires
                  << ", invalid_keys=" << key_infos[0].invaild_keys << "\r\n";
       tmp_stream << db_name << " Hashes_keys=" << key_infos[1].keys << ", expires=" << key_infos[1].expires
@@ -2911,8 +2912,8 @@ void DbsizeCmd::Do() {
     }
     KeyScanInfo key_scan_info = dbs->GetKeyScanInfo();
     std::vector<storage::KeyInfo> key_infos = key_scan_info.key_infos;
-    if (key_infos.size() != (size_t)(storage::DataType::kNones)) {
-      res_.SetRes(CmdRes::kErrOther, "keyspace error");
+    if (key_infos.size() != (size_t)(storage::DataTypeNum)) {
+      res_.SetRes(CmdRes::kErrOther, "Mismatch in expected data types and actual key info count");
       return;
     }
     uint64_t dbsize = 0;

diff --git a/src/storage/src/base_value_format.h b/src/storage/src/base_value_format.h
@@ -19,6 +19,7 @@
 namespace storage {
 
 enum class DataType : uint8_t { kStrings = 0, kHashes = 1, kSets = 2, kLists = 3, kZSets = 4, kStreams = 5, kNones = 6, kAll = 7 };
+constexpr int DataTypeNum = int(DataType::kNones);
 
 constexpr char DataTypeTag[] = { 'k', 'h', 's', 'l', 'z', 'x', 'n', 'a'};
 constexpr char* DataTypeStrings[] = { "string", "hash", "set", "list", "zset", "streams", "none", "all"};

diff --git a/src/storage/src/redis.cc b/src/storage/src/redis.cc
@@ -352,7 +352,7 @@ void Redis::SetCompactRangeOptions(const bool is_canceled) {
     default_compact_range_options_.canceled = new std::atomic<bool>(is_canceled);
   } else {
     default_compact_range_options_.canceled->store(is_canceled);
-  } 
+  }
 }
 
 Status Redis::GetProperty(const std::string& property, uint64_t* out) {
@@ -365,7 +365,7 @@ Status Redis::GetProperty(const std::string& property, uint64_t* out) {
 }
 
 Status Redis::ScanKeyNum(std::vector<KeyInfo>* key_infos) {
-  key_infos->resize(5);
+  key_infos->resize(DataTypeNum);
   rocksdb::Status s;
   s = ScanStringsKeyNum(&((*key_infos)[0]));
   if (!s.ok()) {

diff --git a/src/storage/src/storage.cc b/src/storage/src/storage.cc
@@ -1828,7 +1828,7 @@ uint64_t Storage::GetProperty(const std::string& property) {
 
 Status Storage::GetKeyNum(std::vector<KeyInfo>* key_infos) {
   KeyInfo key_info;
-  key_infos->resize(size_t(DataType::kNones));
+  key_infos->resize(DataTypeNum);
   for (const auto& db : insts_) {
     std::vector<KeyInfo> db_key_infos;
     // check the scanner was stopped or not, before scanning the next db