Merge pull request #5333 from shuujii/fix-heap-buffer-overflow-for-small-Hash-HT-in-Hash-rehash

Fix heap-buffer-overflow for small `Hash` (HT) in `Hash#rehash`

2 files changed, 16 insertions, 1 deletions

diff --git a/src/hash.c b/src/hash.c
index 289f02a91..1fbacc889 100644
--- a/src/hash.c
+++ b/src/hash.c
@@ -718,6 +718,7 @@ ib_bit_for(uint32_t size)
 static uint32_t
 ib_byte_size_for(uint32_t ib_bit)
 {
+  mrb_assert(IB_INIT_BIT <= ib_bit);
   uint32_t ary_size = IB_INIT_BIT == 4 ?
     ib_bit_to_capa(ib_bit) * 2 / IB_TYPE_BIT * ib_bit / 2 :
     ib_bit_to_capa(ib_bit) / IB_TYPE_BIT * ib_bit;
@@ -892,7 +893,13 @@ static void
 ht_rehash(mrb_state *mrb, struct RHash *h)
 {
   /* see comments in `h_rehash` */
-  uint32_t size = ht_size(h), w_size = 0, ea_capa = ht_ea_capa(h);
+  uint32_t size = ht_size(h);
+  if (size <= AR_MAX_SIZE) {
+    ht_to_ar(mrb, h);
+    ar_rehash(mrb, h);
+    return;
+  }
+  uint32_t w_size = 0, ea_capa = ht_ea_capa(h);
   hash_entry *ea = ht_ea(h);
   ht_init(mrb, h, 0, ea, ea_capa, h_ht(h), ib_bit_for(size));
   ht_set_size(h, size);
diff --git a/test/t/hash.rb b/test/t/hash.rb
index c51af03aa..a5e51d83b 100644
--- a/test/t/hash.rb
+++ b/test/t/hash.rb
@@ -944,6 +944,14 @@ assert('Hash#rehash') do
   h = {}
   assert_same(h, h.rehash)
   assert_predicate(h, :empty?)
+
+  h = {}
+  (1..17).each{h[_1] = _1 * 2}
+  (2..16).each{h.delete(_1)}
+  assert_same(h, h.rehash)
+  assert_equal([[1, 2], [17, 34]], h.to_a)
+  assert_equal(2, h.size)
+  [1, 17].each{assert_equal(_1 * 2, h[_1])}
 end
 
 assert('#eql? receiver should be specified key') do