Fix the bitmap range is smaller than that of Redis (apache#465)

ShooterIT · Jan 27, 2022 · 52eb756 · 52eb756
1 parent 9747c3a
commit 52eb756
Show file tree

Hide file tree

Showing 8 changed files with 79 additions and 46 deletions.
diff --git a/src/redis_bitmap.cc b/src/redis_bitmap.cc
@@ -195,7 +195,7 @@ rocksdb::Status Bitmap::SetBit(const Slice &user_key, uint32_t offset, bool new_
   return storage_->Write(rocksdb::WriteOptions(), &batch);
 }
 
-rocksdb::Status Bitmap::BitCount(const Slice &user_key, int start, int stop, uint32_t *cnt) {
+rocksdb::Status Bitmap::BitCount(const Slice &user_key, int64_t start, int64_t stop, uint32_t *cnt) {
   *cnt = 0;
   std::string ns_key, raw_value;
   AppendNamespacePrefix(user_key, &ns_key);
@@ -211,33 +211,37 @@ rocksdb::Status Bitmap::BitCount(const Slice &user_key, int start, int stop, uin
 
   if (start < 0) start += metadata.size + 1;
   if (stop < 0) stop += metadata.size + 1;
-  if (stop > static_cast<int>(metadata.size)) stop = metadata.size;
+  if (stop > static_cast<int64_t>(metadata.size)) stop = metadata.size;
   if (start < 0 || stop <= 0 || start >= stop) return rocksdb::Status::OK();
 
+  uint32_t u_start = static_cast<uint32_t>(start);
+  uint32_t u_stop = static_cast<uint32_t>(stop);
+
   LatestSnapShot ss(db_);
   rocksdb::ReadOptions read_options;
   read_options.snapshot = ss.GetSnapShot();
-  int start_index = start / kBitmapSegmentBytes;
-  int stop_index = stop / kBitmapSegmentBytes;
+  uint32_t start_index = u_start / kBitmapSegmentBytes;
+  uint32_t stop_index = u_stop / kBitmapSegmentBytes;
   // Don't use multi get to prevent large range query, and take too much memory
   std::string sub_key, value;
-  for (int i = start_index; i <= stop_index; i++) {
+  for (uint32_t i = start_index; i <= stop_index; i++) {
     InternalKey(ns_key, std::to_string(i * kBitmapSegmentBytes), metadata.version,
                   storage_->IsSlotIdEncoded()).Encode(&sub_key);
     s = db_->Get(read_options, sub_key, &value);
     if (!s.ok() && !s.IsNotFound()) return s;
     if (s.IsNotFound()) continue;
     size_t j = 0;
-    if (i == start_index) j = start % kBitmapSegmentBytes;
+    if (i == start_index) j = u_start % kBitmapSegmentBytes;
     for (; j < value.size(); j++) {
-      if (i == stop_index && j > (stop % kBitmapSegmentBytes)) break;
+      if (i == stop_index && j > (u_stop % kBitmapSegmentBytes)) break;
       *cnt += kNum2Bits[static_cast<uint8_t>(value[j])];
     }
   }
   return rocksdb::Status::OK();
 }
 
-rocksdb::Status Bitmap::BitPos(const Slice &user_key, bool bit, int start, int stop, bool stop_given, int *pos) {
+rocksdb::Status Bitmap::BitPos(const Slice &user_key, bool bit, int64_t start,
+                                  int64_t stop, bool stop_given, int64_t *pos) {
   std::string ns_key, raw_value;
   AppendNamespacePrefix(user_key, &ns_key);
 
@@ -260,6 +264,8 @@ rocksdb::Status Bitmap::BitPos(const Slice &user_key, bool bit, int start, int s
     *pos = -1;
     return rocksdb::Status::OK();
   }
+  uint32_t u_start = static_cast<uint32_t>(start);
+  uint32_t u_stop = static_cast<uint32_t>(stop);
 
   auto bitPosInByte = [](char byte, bool bit) -> int {
     for (int i = 0; i < 8; i++) {
@@ -272,11 +278,11 @@ rocksdb::Status Bitmap::BitPos(const Slice &user_key, bool bit, int start, int s
   LatestSnapShot ss(db_);
   rocksdb::ReadOptions read_options;
   read_options.snapshot = ss.GetSnapShot();
-  int start_index = start / kBitmapSegmentBytes;
-  int stop_index = stop / kBitmapSegmentBytes;
+  uint32_t start_index = u_start / kBitmapSegmentBytes;
+  uint32_t stop_index = u_stop / kBitmapSegmentBytes;
   // Don't use multi get to prevent large range query, and take too much memory
   std::string sub_key, value;
-  for (int i = start_index; i <= stop_index; i++) {
+  for (uint32_t i = start_index; i <= stop_index; i++) {
     InternalKey(ns_key, std::to_string(i * kBitmapSegmentBytes), metadata.version,
                   storage_->IsSlotIdEncoded()).Encode(&sub_key);
     s = db_->Get(read_options, sub_key, &value);
@@ -289,21 +295,21 @@ rocksdb::Status Bitmap::BitPos(const Slice &user_key, bool bit, int start, int s
       continue;
     }
     size_t j = 0;
-    if (i == start_index) j = start % kBitmapSegmentBytes;
+    if (i == start_index) j = u_start % kBitmapSegmentBytes;
     for (; j < value.size(); j++) {
-      if (i == stop_index && j > (stop % kBitmapSegmentBytes)) break;
+      if (i == stop_index && j > (u_stop % kBitmapSegmentBytes)) break;
       if (bitPosInByte(value[j], bit) != -1) {
-        *pos = static_cast<int>(i * kBitmapSegmentBits + j * 8 + bitPosInByte(value[j], bit));
+        *pos = static_cast<int64_t>(i * kBitmapSegmentBits + j * 8 + bitPosInByte(value[j], bit));
         return rocksdb::Status::OK();
       }
     }
     if (!bit && value.size() < kBitmapSegmentBytes) {
-      *pos = static_cast<int>(i * kBitmapSegmentBits + value.size() * 8);
+      *pos = static_cast<int64_t>(i * kBitmapSegmentBits + value.size() * 8);
       return rocksdb::Status::OK();
     }
   }
   // bit was not found
-  *pos = bit ? -1 : static_cast<int>(metadata.size * 8);
+  *pos = bit ? -1 : static_cast<int64_t>(metadata.size * 8);
   return rocksdb::Status::OK();
 }
 

diff --git a/src/redis_bitmap.h b/src/redis_bitmap.h
@@ -14,8 +14,8 @@ class Bitmap : public Database {
   rocksdb::Status GetBit(const Slice &user_key, uint32_t offset, bool *bit);
   rocksdb::Status GetString(const Slice &user_key, const uint32_t max_btos_size, std::string *value);
   rocksdb::Status SetBit(const Slice &user_key, uint32_t offset, bool new_bit, bool *old_bit);
-  rocksdb::Status BitCount(const Slice &user_key, int start, int stop, uint32_t *cnt);
-  rocksdb::Status BitPos(const Slice &user_key, bool bit, int start, int stop, bool stop_given, int *pos);
+  rocksdb::Status BitCount(const Slice &user_key, int64_t start, int64_t stop, uint32_t *cnt);
+  rocksdb::Status BitPos(const Slice &user_key, bool bit, int64_t start, int64_t stop, bool stop_given, int64_t *pos);
   static bool GetBitFromValueAndOffset(const std::string &value, const uint32_t offset);
   static bool IsEmptySegment(const Slice &segment);
  private:

diff --git a/src/redis_bitmap_string.cc b/src/redis_bitmap_string.cc
@@ -45,7 +45,7 @@ rocksdb::Status BitmapString::SetBit(const Slice &ns_key,
   return storage_->Write(rocksdb::WriteOptions(), &batch);
 }
 
-rocksdb::Status BitmapString::BitCount(const std::string &raw_value, int start, int stop, uint32_t *cnt) {
+rocksdb::Status BitmapString::BitCount(const std::string &raw_value, int64_t start, int64_t stop, uint32_t *cnt) {
   *cnt = 0;
   auto string_value = raw_value.substr(STRING_HDR_SIZE, raw_value.size() - STRING_HDR_SIZE);
   /* Convert negative indexes */
@@ -57,36 +57,36 @@ rocksdb::Status BitmapString::BitCount(const std::string &raw_value, int start,
   if (stop < 0) stop = strlen + stop;
   if (start < 0) start = 0;
   if (stop < 0) stop = 0;
-  if (stop >= static_cast<int>(strlen)) stop = strlen - 1;
+  if (stop >= static_cast<int64_t>(strlen)) stop = strlen - 1;
 
   /* Precondition: end >= 0 && end < strlen, so the only condition where
      * zero can be returned is: start > stop. */
   if (start <= stop) {
-    int bytes = stop - start + 1;
+    int64_t bytes = stop - start + 1;
     *cnt = redisPopcount((unsigned char *) (&string_value[0] + start), bytes);
   }
   return rocksdb::Status::OK();
 }
 
 rocksdb::Status BitmapString::BitPos(const std::string &raw_value,
                                      bool bit,
-                                     int start,
-                                     int stop,
+                                     int64_t start,
+                                     int64_t stop,
                                      bool stop_given,
-                                     int *pos) {
+                                     int64_t *pos) {
   auto string_value = raw_value.substr(STRING_HDR_SIZE, raw_value.size() - STRING_HDR_SIZE);
   auto strlen = string_value.size();
   /* Convert negative indexes */
   if (start < 0) start = strlen + start;
   if (stop < 0) stop = strlen + stop;
   if (start < 0) start = 0;
   if (stop < 0) stop = 0;
-  if (stop >= static_cast<int>(strlen)) stop = strlen - 1;
+  if (stop >= static_cast<int64_t>(strlen)) stop = strlen - 1;
 
   if (start > stop) {
     *pos = -1;
   } else {
-    int bytes = stop - start + 1;
+    int64_t bytes = stop - start + 1;
     *pos = redisBitpos((unsigned char *) (&string_value[0] + start), bytes, bit);
 
     /* If we are looking for clear bits, and the user specified an exact
@@ -113,7 +113,7 @@ rocksdb::Status BitmapString::BitPos(const std::string &raw_value,
  * This function started out as:
  * https://github.com/antirez/redis/blob/94f2e7f/src/bitops.c#L40
  * */
-size_t BitmapString::redisPopcount(unsigned char *p, int count) {
+size_t BitmapString::redisPopcount(unsigned char *p, int64_t count) {
   size_t bits = 0;
   uint32_t *p4;
   static const unsigned char bitsinbyte[256] =
@@ -185,10 +185,10 @@ size_t BitmapString::redisPopcount(unsigned char *p, int count) {
  * This function started out as:
  * https://github.com/antirez/redis/blob/94f2e7f/src/bitops.c#L101
  * */
-int BitmapString::redisBitpos(unsigned char *c, int count, int bit) {
+int64_t BitmapString::redisBitpos(unsigned char *c, int64_t count, int bit) {
   uint64_t *l;
   uint64_t skipval, word = 0, one;
-  int pos = 0; /* Position of bit, to return to the caller. */
+  int64_t pos = 0; /* Position of bit, to return to the caller. */
   uint64_t j;
   int found;
 

diff --git a/src/redis_bitmap_string.h b/src/redis_bitmap_string.h
@@ -13,11 +13,12 @@ class BitmapString : public Database {
   BitmapString(Engine::Storage *storage, const std::string &ns) : Database(storage, ns) {}
   rocksdb::Status GetBit(const std::string &raw_value, uint32_t offset, bool *bit);
   rocksdb::Status SetBit(const Slice &ns_key, std::string *raw_value, uint32_t offset, bool new_bit, bool *old_bit);
-  rocksdb::Status BitCount(const std::string &raw_value, int start, int stop, uint32_t *cnt);
-  rocksdb::Status BitPos(const std::string &raw_value, bool bit, int start, int stop, bool stop_given, int *pos);
+  rocksdb::Status BitCount(const std::string &raw_value, int64_t start, int64_t stop, uint32_t *cnt);
+  rocksdb::Status BitPos(const std::string &raw_value, bool bit,
+                          int64_t start, int64_t stop, bool stop_given, int64_t *pos);
  private:
-  size_t redisPopcount(unsigned char *p, int count);
-  int redisBitpos(unsigned char *c, int count, int bit);
+  size_t redisPopcount(unsigned char *p, int64_t count);
+  int64_t redisBitpos(unsigned char *c, int64_t count, int bit);
 };
 
 }  // namespace Redis
diff --git a/src/redis_cmd.cc b/src/redis_cmd.cc
@@ -729,7 +729,7 @@ Status getBitOffsetFromArgument(std::string arg, uint32_t *offset) {
   } catch (std::exception &e) {
     return Status(Status::RedisParseErr, errValueNotInterger);
   }
-  if (offset_arg < 0 || offset_arg > INT_MAX) {
+  if (offset_arg < 0 || offset_arg > UINT_MAX) {
     return Status(Status::RedisParseErr, "bit offset is out of range");
   }
   *offset = static_cast<uint32_t>(offset_arg);
@@ -788,11 +788,16 @@ class CommandSetBit : public Commander {
 class CommandBitCount : public Commander {
  public:
   Status Parse(const std::vector<std::string> &args) override {
-    try {
-      if (args.size() >= 3) start_ = std::stoi(args[2]);
-      if (args.size() >= 4) stop_ = std::stoi(args[3]);
-    } catch (std::exception &e) {
-      return Status(Status::RedisParseErr, errValueNotInterger);
+    if (args.size() == 3) {
+      return Status(Status::RedisParseErr, errInvalidSyntax);
+    }
+    if (args.size() == 4) {
+      try {
+        start_ = std::stol(args[2]);
+        stop_ = std::stol(args[3]);
+      } catch (std::exception &e) {
+        return Status(Status::RedisParseErr, errValueNotInterger);
+      }
     }
     return Commander::Parse(args);
   }
@@ -805,18 +810,19 @@ class CommandBitCount : public Commander {
     *output = Redis::Integer(cnt);
     return Status::OK();
   }
+
  private:
-  int start_ = 0, stop_ = -1;
+  int64_t start_ = 0, stop_ = -1;
 };
 
 class CommandBitPos: public Commander {
  public:
   Status Parse(const std::vector<std::string> &args) override {
     try {
-      if (args.size() >= 4) start_ = std::stoi(args[3]);
+      if (args.size() >= 4) start_ = std::stol(args[3]);
       if (args.size() >= 5) {
         stop_given_ = true;
-        stop_ = std::stoi(args[4]);
+        stop_ = std::stol(args[4]);
       }
     } catch (std::exception &e) {
       return Status(Status::RedisParseErr, errValueNotInterger);
@@ -832,7 +838,7 @@ class CommandBitPos: public Commander {
   }
 
   Status Execute(Server *svr, Connection *conn, std::string *output) override {
-    int pos;
+    int64_t pos;
     Redis::Bitmap bitmap_db(svr->storage_, conn->GetNamespace());
     rocksdb::Status s = bitmap_db.BitPos(args_[1], bit_, start_, stop_, stop_given_, &pos);
     if (!s.ok()) return Status(Status::RedisExecErr, s.ToString());
@@ -841,7 +847,7 @@ class CommandBitPos: public Commander {
   }
 
  private:
-  int start_ = 0, stop_ = -1;
+  int64_t start_ = 0, stop_ = -1;
   bool bit_ = false, stop_given_ = false;
 };
 

diff --git a/tests/t_bitmap_test.cc b/tests/t_bitmap_test.cc
@@ -48,7 +48,7 @@ TEST_F(RedisBitmapTest, BitCount) {
 }
 
 TEST_F(RedisBitmapTest, BitPosClearBit) {
-  int pos;
+  int64_t pos;
   bool old_bit;
   for (int i = 0; i < 1024+16;i ++) {
     bitmap->BitPos(key_, false, 0, -1, true, &pos);
@@ -65,7 +65,7 @@ TEST_F(RedisBitmapTest, BitPosSetBit) {
     bool bit = false;
     bitmap->SetBit(key_, offset, true, &bit);
   }
-  int pos;
+  int64_t pos;
   int start_indexes[] = {0, 1, 124, 1025, 1027, 3*1024+1};
   for (size_t i = 0; i < sizeof(start_indexes)/ sizeof(start_indexes[0]); i++) {
     bitmap->BitPos(key_, true, start_indexes[i], -1, true, &pos);

diff --git a/tests/tcl/tests/unit/type/bitmap.tcl b/tests/tcl/tests/unit/type/bitmap.tcl
@@ -21,4 +21,14 @@ start_server {tags {"bitmap"}} {
       catch {r get b0} e
       set e
   } {ERR Operation aborted: The size of the bitmap *}
+
+  test "SETBIT/GETBIT/BITCOUNT/BITPOS boundary check (type bitmap)" {
+      r del b0
+      set max_offset [expr 4*1024*1024*1024-1]
+      assert_error "*out of range*" {r setbit b0 [expr $max_offset+1] 1}
+      r setbit b0 $max_offset 1
+      assert_equal 1 [r getbit b0 $max_offset ]
+      assert_equal 1 [r bitcount b0 0 [expr $max_offset / 8] ]
+      assert_equal $max_offset  [r bitpos b0 1 ]
+  }
 }
diff --git a/tests/tcl/tests/unit/type/string.tcl b/tests/tcl/tests/unit/type/string.tcl
@@ -293,6 +293,16 @@ start_server {tags {"string"}} {
         assert_error "*out of range*" {r setbit mykey 0 20}
     }
 
+    test "SETBIT/GETBIT/BITCOUNT/BITPOS boundary check (type string)" {
+        r del mykey
+        r set mykey ""
+        set max_offset [expr 4*1024*1024*1024-1]
+        r setbit mykey $max_offset 1
+        assert_equal 1 [r getbit mykey $max_offset ]
+        assert_equal 1 [r bitcount mykey 0 [expr $max_offset / 8] ]
+        assert_equal $max_offset  [r bitpos mykey 1 ]
+    }
+
     # test "SETBIT fuzzing" {
     #     set str ""
     #     set len [expr 256*8]