Fix: Random partitioner hash values should be in the range [1, 2^127-1]

Author: Michael Penick

src/md5.cpp

@@ -102,38 +102,6 @@ void Md5::update(const uint8_t* data, size_t size) {
 }
 
 void Md5::final(uint8_t* result) {
-  final();
-
-  result[0]  = a_;
-  result[1]  = a_ >> 8;
-  result[2]  = a_ >> 16;
-  result[3]  = a_ >> 24;
-  result[4]  = b_;
-  result[5]  = b_ >> 8;
-  result[6]  = b_ >> 16;
-  result[7]  = b_ >> 24;
-  result[8]  = c_;
-  result[9]  = c_ >> 8;
-  result[10] = c_ >> 16;
-  result[11] = c_ >> 24;
-  result[12] = d_;
-  result[13] = d_ >> 8;
-  result[14] = d_ >> 16;
-  result[15] = d_ >> 24;
-
-  memset(this, 0, sizeof(Md5));
-}
-
-void Md5::final(uint64_t* hi, uint64_t* lo) {
-  final();
-
-  *hi = static_cast<uint64_t>(a_) << 32 | (static_cast<uint64_t>(b_) & 0xFFFFFFFF);
-  *lo  = static_cast<uint64_t>(c_) << 32 | (static_cast<uint64_t>(d_) & 0xFFFFFFFF);
-
-  memset(this, 0, sizeof(Md5));
-}
-
-void Md5::final() {
   unsigned long used, free;
 
   used = lo_ & 0x3f;
@@ -162,6 +130,25 @@ void Md5::final() {
   buffer_[63] = hi_ >> 24;
 
   body(buffer_, 64);
+
+  result[0]  = a_;
+  result[1]  = a_ >> 8;
+  result[2]  = a_ >> 16;
+  result[3]  = a_ >> 24;
+  result[4]  = b_;
+  result[5]  = b_ >> 8;
+  result[6]  = b_ >> 16;
+  result[7]  = b_ >> 24;
+  result[8]  = c_;
+  result[9]  = c_ >> 8;
+  result[10] = c_ >> 16;
+  result[11] = c_ >> 24;
+  result[12] = d_;
+  result[13] = d_ >> 8;
+  result[14] = d_ >> 16;
+  result[15] = d_ >> 24;
+
+  memset(this, 0, sizeof(Md5));
 }
 
 // This processes one or more 64-byte data blocks, but does NOT update

src/md5.hpp

@@ -33,10 +33,8 @@ class Md5 {
 
   void update(const uint8_t* data, size_t size);
   void final(uint8_t* result);
-  void final(uint64_t* hi, uint64_t* lo);
 
 private:
-  void final();
   const uint8_t* body(const uint8_t* data, size_t size);
 
 private:

src/token_map_impl.cpp

@@ -100,11 +100,45 @@ RandomPartitioner::Token RandomPartitioner::from_string(const StringRef& str) {
   return token;
 }
 
+uint64_t RandomPartitioner::encode(uint8_t* bytes) {
+  uint64_t result = 0;
+  const size_t num_bytes = sizeof(uint64_t);
+  for (size_t i = 0; i < num_bytes; ++i) {
+    result |= (static_cast<uint64_t>(bytes[i]) << (8 * (num_bytes - i - 1)));
+  }
+  return result;
+}
+
+RandomPartitioner::Token RandomPartitioner::abs(RandomPartitioner::Token token) {
+  if (token.hi & 0x8000000000000000ULL) {
+    token.hi = ~token.hi;
+    token.lo = ~token.lo;
+
+    uint64_t old_lo = token.lo;
+    ++token.lo;
+    // Carry to "hi" if our "lo" value wrapped
+    if(token.lo < old_lo) {
+      ++token.hi;
+    }
+  }
+  return token;
+}
+
 RandomPartitioner::Token RandomPartitioner::hash(const StringRef& str) {
   Md5 hash;
   hash.update(reinterpret_cast<const uint8_t*>(str.data()), str.size());
+  uint8_t digest[16];
+  hash.final(digest);
   Token token;
-  hash.final(&token.hi, &token.lo);
+
+  // For compatability with Cassandra we interpret the MD5 as a big-endian value:
+  // Reference: https://docs.oracle.com/javase/7/docs/api/java/math/BigInteger.html#BigInteger(byte[])
+  token.hi = encode(digest);
+  token.lo = encode(digest + 8);
+
+  // Then we find the absolute value of the two's complement representation.
+  token = abs(token);
+
   return token;
 }
 

src/token_map_impl.hpp

@@ -121,6 +121,9 @@ struct RandomPartitioner {
     }
   };
 
+  static Token abs(Token token);
+  static uint64_t encode(uint8_t* bytes);
+
   static Token from_string(const StringRef& str);
   static Token hash(const StringRef& str);
   static StringRef name() { return "RandomPartitioner"; }

test/unit_tests/src/uint128.hpp

@@ -95,6 +95,7 @@ class uint128 : boost::operators<uint128>, boost::shiftable<uint128> {
 	uint128(double value)         : lo(static_cast<base_type>(value)), hi(0) {}
 	uint128(const uint128 &value) : lo(value.lo), hi (value.hi)              {}
 	uint128(base_type value)      : lo(value), hi(0)                         {}
+	uint128(base_type lo, base_type hi) : lo(lo), hi(hi)                     {}
 
 	uint128(const std::string &sz) : lo(0), hi(0) {