Fix boundary check in Bitmap::BitOp (#1727)

src/commands/cmd_bit.cc

@@ -203,6 +203,7 @@ class CommandBitOp : public Commander {
 
   Status Execute(Server *svr, Connection *conn, std::string *output) override {
     std::vector<Slice> op_keys;
+    op_keys.reserve(args_.size() - 2);
     for (uint64_t i = 3; i < args_.size(); i++) {
       op_keys.emplace_back(args_[i]);
     }

src/types/redis_bitmap.cc

@@ -359,7 +359,7 @@ rocksdb::Status Bitmap::BitOp(BitOpFlags op_flag, const std::string &op_name, co
       return rocksdb::Status::InvalidArgument(kErrMsgWrongType);
     }
     if (metadata.size > max_size) max_size = metadata.size;
-    meta_pairs.emplace_back(ns_op_key, metadata);
+    meta_pairs.emplace_back(std::move(ns_op_key), metadata);
   }
 
   auto batch = storage_->GetWriteBatchBase();
@@ -376,7 +376,8 @@ rocksdb::Status Bitmap::BitOp(BitOpFlags op_flag, const std::string &op_name, co
 
   BitmapMetadata res_metadata;
   if (num_keys == op_keys.size() || op_flag != kBitOpAnd) {
-    uint64_t frag_numkeys = num_keys, stop_index = (max_size - 1) / kBitmapSegmentBytes;
+    uint64_t frag_numkeys = num_keys;
+    uint64_t stop_index = (max_size - 1) / kBitmapSegmentBytes;
     std::unique_ptr<unsigned char[]> frag_res(new unsigned char[kBitmapSegmentBytes]);
     uint16_t frag_maxlen = 0, frag_minlen = 0;
     std::string fragment;
@@ -404,7 +405,7 @@ rocksdb::Status Bitmap::BitOp(BitOpFlags op_flag, const std::string &op_name, co
         } else {
           if (frag_maxlen < fragment.size()) frag_maxlen = fragment.size();
           if (fragment.size() < frag_minlen || frag_minlen == 0) frag_minlen = fragment.size();
-          fragments.emplace_back(fragment);
+          fragments.emplace_back(std::move(fragment));
         }
       }
 
@@ -502,7 +503,14 @@ rocksdb::Status Bitmap::BitOp(BitOpFlags op_flag, const std::string &op_name, co
 
         if (op_flag == kBitOpNot) {
           if (frag_index == stop_index) {
-            frag_maxlen = max_size % kBitmapSegmentBytes;
+            if (max_size == (frag_index + 1) * kBitmapSegmentBytes) {
+              // If the last fragment is full, `max_size % kBitmapSegmentBytes`
+              // would be 0. In this case, we should set `frag_maxlen` to
+              // `kBitmapSegmentBytes` to avoid writing an empty fragment.
+              frag_maxlen = kBitmapSegmentBytes;
+            } else {
+              frag_maxlen = max_size % kBitmapSegmentBytes;
+            }
           } else {
             frag_maxlen = kBitmapSegmentBytes;
           }

tests/cppunit/disk_test.cc

@@ -168,6 +168,46 @@ TEST_F(RedisDiskTest, BitmapDisk) {
   bitmap->Del(key_);
 }
 
+TEST_F(RedisDiskTest, BitmapDisk2) {
+  const int64_t kGroupSize = 8192;
+  for (size_t num_bits : {8192, 16384}) {
+    for (bool set_op : {false, true}) {
+      std::unique_ptr<redis::Bitmap> bitmap = std::make_unique<redis::Bitmap>(storage_, "disk_ns_bitmap2");
+      std::unique_ptr<redis::Disk> disk = std::make_unique<redis::Disk>(storage_, "disk_ns_bitmap2");
+      key_ = "bitmapdisk_key2";
+      bitmap->Del(key_);
+      bool bit = false;
+
+      for (size_t i = 0; i < num_bits; i += kGroupSize) {
+        // Set all first bit of group to `!set_op`
+        EXPECT_TRUE(bitmap->SetBit(key_, i, !set_op, &bit).ok());
+        // Set all last bit of group to `set_op`.
+        EXPECT_TRUE(bitmap->SetBit(key_, i + kGroupSize - 1, set_op, &bit).ok());
+      }
+
+      auto bit_not_dest_key = "bit_op_not_dest_key";
+
+      int64_t len = 0;
+      EXPECT_TRUE(bitmap->BitOp(BitOpFlags::kBitOpNot, "NOT", bit_not_dest_key, {key_}, &len).ok());
+
+      for (size_t i = 0; i < num_bits; i += kGroupSize) {
+        bool result = false;
+        // Check all first bit of group is `set_op`
+        EXPECT_TRUE(bitmap->GetBit(bit_not_dest_key, i, &result).ok());
+        EXPECT_EQ(set_op, result);
+        // Check all last bit of group is `!set_op`
+        EXPECT_TRUE(bitmap->GetBit(bit_not_dest_key, i + kGroupSize - 1, &result).ok());
+        EXPECT_EQ(!set_op, result);
+        // Check bit in group between (first, last) is "1".
+        for (size_t j = i + 1; j < i + kGroupSize - 1; ++j) {
+          EXPECT_TRUE(bitmap->GetBit(bit_not_dest_key, j, &result).ok());
+          EXPECT_TRUE(result) << j << " is not true";
+        }
+      }
+    }
+  }
+}
+
 TEST_F(RedisDiskTest, SortedintDisk) {
   std::unique_ptr<redis::Sortedint> sortedint = std::make_unique<redis::Sortedint>(storage_, "disk_ns_sortedint");
   std::unique_ptr<redis::Disk> disk = std::make_unique<redis::Disk>(storage_, "disk_ns_sortedint");

tests/gocase/unit/type/bitmap/bitmap_test.go

@@ -234,6 +234,14 @@ func TestBitmap(t *testing.T) {
 		}
 	})
 
+	t.Run("BITOP Boundary Check", func(t *testing.T) {
+		require.NoError(t, rdb.Del(ctx, "str").Err())
+		str := util.RandStringWithSeed(0, 1000, util.Binary, 2701)
+		Set2SetBit(t, rdb, ctx, "str", []byte(str))
+		require.NoError(t, rdb.BitOpNot(ctx, "target", "str").Err())
+		require.EqualValues(t, SimulateBitOp(NOT, []byte(str)), rdb.Get(ctx, "target").Val())
+	})
+
 	t.Run("BITOP with non string source key", func(t *testing.T) {
 		require.NoError(t, rdb.Del(ctx, "c").Err())
 		Set2SetBit(t, rdb, ctx, "a", []byte("\xaa\x00\xff\x55"))

tests/gocase/util/random.go

@@ -58,7 +58,11 @@ const (
 )
 
 func RandString(min, max int, typ RandStringType) string {
-	r := rand.New(rand.NewSource(time.Now().UnixNano()))
+	return RandStringWithSeed(min, max, typ, time.Now().UnixNano())
+}
+
+func RandStringWithSeed(min, max int, typ RandStringType, seed int64) string {
+	r := rand.New(rand.NewSource(seed))
 	length := min + r.Intn(max-min+1)
 
 	var minVal, maxVal int
@@ -71,7 +75,7 @@ func RandString(min, max int, typ RandStringType) string {
 
 	var sb strings.Builder
 	for ; length > 0; length-- {
-		s := fmt.Sprintf("%c", minVal+rand.Intn(maxVal-minVal+1))
+		s := fmt.Sprintf("%c", minVal+int(r.Int31n(int32(maxVal-minVal+1))))
 		sb.WriteString(s)
 	}
 	return sb.String()