Update redis_sets.cc

OpenAtomFoundation · Nov 7, 2023 · c55d644 · c55d644
1 parent dccf192
commit c55d644
Showing 1 changed file with 27 additions and 0 deletions.
diff --git a/src/storage/src/redis_sets.cc b/src/storage/src/redis_sets.cc
@@ -217,6 +217,9 @@ rocksdb::Status RedisSets::SAdd(const Slice& key, const std::vector<std::string>
     ParsedSetsMetaValue parsed_sets_meta_value(&meta_value);
     if (parsed_sets_meta_value.IsStale() || parsed_sets_meta_value.count() == 0) {
       version = parsed_sets_meta_value.InitialMetaValue();
+      if (!parsed_sets_meta_value.check_set_count(static_cast<int32_t>(filtered_members.size()))) {
+        return Status::InvalidArgument("set size overflow");
+      }
       parsed_sets_meta_value.set_count(static_cast<int32_t>(filtered_members.size()));
       batch.Put(handles_[0], key, meta_value);
       for (const auto& member : filtered_members) {
@@ -243,6 +246,9 @@ rocksdb::Status RedisSets::SAdd(const Slice& key, const std::vector<std::string>
       if (cnt == 0) {
         return rocksdb::Status::OK();
       } else {
+        if (!parsed_sets_meta_value.CheckModifyCount(cnt)){
+          return Status::InvalidArgument("set size overflow");
+        }
         parsed_sets_meta_value.ModifyCount(cnt);
         batch.Put(handles_[0], key, meta_value);
       }
@@ -420,6 +426,9 @@ rocksdb::Status RedisSets::SDiffstore(const Slice& destination, const std::vecto
     ParsedSetsMetaValue parsed_sets_meta_value(&meta_value);
     statistic = parsed_sets_meta_value.count();
     version = parsed_sets_meta_value.InitialMetaValue();
+      if (!parsed_sets_meta_value.check_set_count(static_cast<int32_t>(members.size()))) {
+        return Status::InvalidArgument("set size overflow");
+      }
     parsed_sets_meta_value.set_count(static_cast<int32_t>(members.size()));
     batch.Put(handles_[0], destination, meta_value);
   } else if (s.IsNotFound()) {
@@ -603,6 +612,9 @@ rocksdb::Status RedisSets::SInterstore(const Slice& destination, const std::vect
     ParsedSetsMetaValue parsed_sets_meta_value(&meta_value);
     statistic = parsed_sets_meta_value.count();
     version = parsed_sets_meta_value.InitialMetaValue();
+    if (!parsed_sets_meta_value.check_set_count(static_cast<int32_t>(members.size()))) {
+      return Status::InvalidArgument("set size overflow");
+    }
     parsed_sets_meta_value.set_count(static_cast<int32_t>(members.size()));
     batch.Put(handles_[0], destination, meta_value);
   } else if (s.IsNotFound()) {
@@ -714,6 +726,9 @@ rocksdb::Status RedisSets::SMove(const Slice& source, const Slice& destination,
       s = db_->Get(default_read_options_, handles_[1], sets_member_key.Encode(), &member_value);
       if (s.ok()) {
         *ret = 1;
+        if (!parsed_sets_meta_value.CheckModifyCount(-1)){
+          return Status::InvalidArgument("set size overflow");
+        }
         parsed_sets_meta_value.ModifyCount(-1);
         batch.Put(handles_[0], source, meta_value);
         batch.Delete(handles_[1], sets_member_key.Encode());
@@ -747,6 +762,9 @@ rocksdb::Status RedisSets::SMove(const Slice& source, const Slice& destination,
       SetsMemberKey sets_member_key(destination, version, member);
       s = db_->Get(default_read_options_, handles_[1], sets_member_key.Encode(), &member_value);
       if (s.IsNotFound()) {
+        if (!parsed_sets_meta_value.CheckModifyCount(1)){
+          return Status::InvalidArgument("set size overflow");
+        }
         parsed_sets_meta_value.ModifyCount(1);
         batch.Put(handles_[0], destination, meta_value);
         batch.Put(handles_[1], sets_member_key.Encode(), Slice());
@@ -843,6 +861,9 @@ rocksdb::Status RedisSets::SPop(const Slice& key, std::vector<std::string>* memb
           }
         }
 
+        if (!parsed_sets_meta_value.CheckModifyCount(static_cast<int32_t>(-cnt))){
+          return Status::InvalidArgument("set size overflow");
+        }
         parsed_sets_meta_value.ModifyCount(static_cast<int32_t>(-cnt));
         batch.Put(handles_[0], key, meta_value);
         delete iter;
@@ -974,6 +995,9 @@ rocksdb::Status RedisSets::SRem(const Slice& key, const std::vector<std::string>
         }
       }
       *ret = cnt;
+      if (!parsed_sets_meta_value.CheckModifyCount(-cnt)){
+        return Status::InvalidArgument("set size overflow");
+      }
       parsed_sets_meta_value.ModifyCount(-cnt);
       batch.Put(handles_[0], key, meta_value);
     }
@@ -1086,6 +1110,9 @@ rocksdb::Status RedisSets::SUnionstore(const Slice& destination, const std::vect
     ParsedSetsMetaValue parsed_sets_meta_value(&meta_value);
     statistic = parsed_sets_meta_value.count();
     version = parsed_sets_meta_value.InitialMetaValue();
+    if (!parsed_sets_meta_value.check_set_count(static_cast<int32_t>(members.size()))) {
+      return Status::InvalidArgument("set size overflow");
+    }
     parsed_sets_meta_value.set_count(static_cast<int32_t>(members.size()));
     batch.Put(handles_[0], destination, meta_value);
   } else if (s.IsNotFound()) {