Subject: [PATCH 8/9] st.c, hash.c: use weak and strong hashes

From: Sokolov Yura aka funny_falcon <[email-masked]>
Date: Mon, 26 Sep 2016 20:26:23 +0300

Idea proposed and implemented first by Vladimir Makarov
<[email-masked]>
- use weaker hash function until ddos attack detected
- switch to stronger (but slower) hash function then.

co-author: Vladimi Makarov <[email-masked]>
---
 bignum.c          |  6 ++++++
 hash.c            | 44 +++++++++++++++++++++-----------------------
 include/ruby/st.h |  5 ++++-
 internal.h        |  1 +
 st.c              | 53 ++++++++++++++++++++++++++++++++++++++++++++++-------
 5 files changed, 78 insertions(+), 31 deletions(-)

diff --git a/bignum.c b/bignum.c
index 18919d0..0789f67 100644
--- a/bignum.c
+++ b/bignum.c
@@ -6648,6 +6648,12 @@ rb_big_hash(VALUE x)
     return ST2FIX(hash);
 }
 
+long
+rb_big_hash_long(VALUE x)
+{
+    return st_hash(BDIGITS(x), sizeof(BDIGIT)*BIGNUM_LEN(x), 0xb19 ^ BIGNUM_SIGN(x));
+}
+
 /*
  * call-seq:
  *   big.coerce(numeric)  ->  array
diff --git a/hash.c b/hash.c
index 3b6c111..67a48a1 100644
--- a/hash.c
+++ b/hash.c
@@ -146,7 +146,7 @@ rb_hash(VALUE obj)
 long rb_objid_hash(st_index_t index);
 
 static st_index_t
-any_hash(VALUE a, st_index_t (*other_func)(VALUE))
+any_hash(VALUE a, st_index_t (*other_func)(VALUE), int strong)
 {
     VALUE hval;
     st_index_t hnum;
@@ -164,14 +164,19 @@ any_hash(VALUE a, st_index_t (*other_func)(VALUE))
 	hnum = rb_objid_hash((st_index_t)a);
     }
     else if (BUILTIN_TYPE(a) == T_STRING) {
-	hnum = rb_str_hash(a);
+	hnum = strong ? rb_str_hash(a) :
+		st_hash(RSTRING_PTR(a), RSTRING_LEN(a), 0xcafe);
     }
     else if (BUILTIN_TYPE(a) == T_SYMBOL) {
 	hnum = RSYMBOL(a)->hashval;
     }
     else if (BUILTIN_TYPE(a) == T_BIGNUM) {
-	hval = rb_big_hash(a);
-	hnum = FIX2LONG(hval);
+	if (strong) {
+	    hval = rb_big_hash(a);
+	    hnum = FIX2LONG(hval);
+	} else {
+	    hnum = rb_big_hash_long(a);
+	}
     }
     else if (BUILTIN_TYPE(a) == T_FLOAT) {
       flt:
@@ -195,7 +200,13 @@ obj_any_hash(VALUE obj)
 static st_index_t
 rb_any_hash(VALUE a)
 {
-    return any_hash(a, obj_any_hash);
+    return any_hash(a, obj_any_hash, 1);
+}
+
+static st_index_t
+rb_any_hash_weak(VALUE a)
+{
+    return any_hash(a, obj_any_hash, 0);
 }
 
 st_index_t
@@ -232,7 +243,7 @@ objid_hash(VALUE obj)
 VALUE
 rb_obj_hash(VALUE obj)
 {
-    st_index_t hnum = any_hash(obj, objid_hash);
+    st_index_t hnum = any_hash(obj, objid_hash, 1);
     return ST2FIX(hnum);
 }
 
@@ -244,6 +255,7 @@ rb_hash_iter_lev(VALUE h)
 
 static const struct st_hash_type objhash = {
     rb_any_cmp,
+    rb_any_hash_weak,
     rb_any_hash,
 };
 
@@ -714,20 +726,6 @@ rb_hash_s_try_convert(VALUE dummy, VALUE hash)
     return rb_check_hash_type(hash);
 }
 
-struct rehash_arg {
-    VALUE hash;
-    st_table *tbl;
-};
-
-static int
-rb_hash_rehash_i(VALUE key, VALUE value, VALUE arg)
-{
-    st_table *tbl = (st_table *)arg;
-
-    st_insert(tbl, (st_data_t)key, (st_data_t)value);
-    return ST_CONTINUE;
-}
-
 /*
  *  call-seq:
  *     hsh.rehash -> hsh
@@ -761,10 +759,9 @@ rb_hash_rehash(VALUE hash)
     if (!RHASH(hash)->ntbl)
         return hash;
     tmp = hash_alloc(0);
-    tbl = st_init_table_with_size(RHASH(hash)->ntbl->type, RHASH(hash)->ntbl->num_entries);
+    tbl = st_copy(RHASH(hash)->ntbl);
     RHASH(tmp)->ntbl = tbl;
-
-    rb_hash_foreach(hash, rb_hash_rehash_i, (VALUE)tbl);
+    st_recalc_hashvals(tbl);
     st_free_table(RHASH(hash)->ntbl);
     RHASH(hash)->ntbl = tbl;
     RHASH(tmp)->ntbl = 0;
@@ -2546,6 +2543,7 @@ rb_hash_assoc(VALUE hash, VALUE key)
 
 	assochash.compare = assoc_cmp;
 	assochash.hash = orighash->hash;
+	assochash.stronghash = orighash->stronghash;
 	table->type = &assochash;
 	args[0] = hash;
 	args[1] = key;
diff --git a/include/ruby/st.h b/include/ruby/st.h
index 8763d9d..10ae61c 100644
--- a/include/ruby/st.h
+++ b/include/ruby/st.h
@@ -63,6 +63,7 @@ typedef uint32_t st_idx_t;
 struct st_hash_type {
     int (*compare)(ANYARGS /*st_data_t, st_data_t*/); /* st_compare_func* */
     st_index_t (*hash)(ANYARGS /*st_data_t*/);        /* st_hash_func* */
+    st_index_t (*stronghash)(ANYARGS /*st_data_t*/);        /* st_hash_func* */
 };
 
 #define ST_INDEX_BITS (SIZEOF_ST_INDEX_T * CHAR_BIT)
@@ -86,7 +87,8 @@ struct st_table {
     st_idx_t num_entries;
     st_idx_t first, last;
     unsigned sz : 8;
-    st_idx_t rebuild_num : sizeof(st_idx_t)*8-8;
+    unsigned use_strong: 1;
+    st_idx_t rebuild_num : sizeof(st_idx_t)*8-9;
 };
 
 #define st_is_member(table,key) st_lookup((table),(key),(st_data_t *)0)
@@ -125,6 +127,7 @@ void st_free_table(st_table *);
 void st_cleanup_safe(st_table *, st_data_t);
 void st_clear(st_table *);
 st_table *st_copy(st_table *);
+void st_recalc_hashvals(st_table*);
 CONSTFUNC(int st_numcmp(st_data_t, st_data_t));
 CONSTFUNC(st_index_t st_numhash(st_data_t));
 PUREFUNC(int st_locale_insensitive_strcasecmp(const char *s1, const char *s2));
diff --git a/internal.h b/internal.h
index fad7753..f3d90da 100644
--- a/internal.h
+++ b/internal.h
@@ -894,6 +894,7 @@ extern const char ruby_digitmap[];
 VALUE rb_big_fdiv(VALUE x, VALUE y);
 VALUE rb_big_uminus(VALUE x);
 VALUE rb_big_hash(VALUE);
+long  rb_big_hash_long(VALUE);
 VALUE rb_big_odd_p(VALUE);
 VALUE rb_big_even_p(VALUE);
 size_t rb_big_size(VALUE);
diff --git a/st.c b/st.c
index f48cdf3..550a11c 100644
--- a/st.c
+++ b/st.c
@@ -143,7 +143,8 @@ need_shrink(st_idx_t num_entries, int sz)
 static inline st_index_t
 do_hash(st_data_t key, const st_table * table)
 {
-    st_index_t h = (*table->type->hash)(key);
+    st_index_t h = table->use_strong == 0 ? (*table->type->hash)(key) :
+	    (*table->type->stronghash)(key);
     return h != DELETED ? h : 0x71fe900d;
 }
 
@@ -249,6 +250,7 @@ stat_col(void)
     fprintf(f, "collision: %d / %d (%6.2f)\n", collision.all, collision.total,
 	    ((double)collision.all / (collision.total)) * 100);
     fprintf(f, "num: %d, str: %d, strcase: %d\n", collision.num, collision.str, collision.strcase);
+    fprintf(f, "max: %d\n", collision.max);
     fclose(f);
 }
 #endif
@@ -495,36 +497,52 @@ find_exact_key(const st_table *table, st_data_t key, st_index_t hash_val)
     }
 }
 
-static inline void
+/* Switch to strong hash if collision chain is greater than ST_STRONG_THRESHOLD.
+ * ST_STRONG_THRESHOLD is chosen to never switch during `make check`
+ * (and, probably, in real life).
+ * So, additional test run with lower threshold should be run occasionally.
+ */
+#define ST_STRONG_THRESHOLD 40
+static inline int
 insert_entry_simple(st_table* table, st_index_t hash_val, st_idx_t idx)
 {
     const void* bins = table->as.bins;
     int sz = table->sz;
+    int cnt = 0;
     st_idx_t bin_pos;
     FL_INIT(hash_val);
-    for (;;FL_INC()) {
+    for (;;FL_INC(), cnt++) {
 	bin_pos = FL_POS(table);
 	if (bin_get(bins, sz, bin_pos) == IDX_NULL)
 	    break;
     }
     bin_set(bins, sz, bin_pos, idx+1);
+#if defined(HASH_LOG) && collision_check
+    if (table->type->stronghash && cnt > collision.max) collision.max = cnt;
+#endif
+    return cnt >= ST_STRONG_THRESHOLD;
 }
 
-static inline void
+static inline int
 insert_entry_reuse(st_table* table, st_index_t hash_val, st_idx_t idx)
 {
     const void* bins = table->as.bins;
     const st_table_entry* entries = table->as.entries;
     int sz = table->sz;
+    int cnt = 0;
     st_idx_t bin_pos, old;
     FL_INIT(hash_val);
-    for (;;FL_INC()) {
+    for (;;FL_INC(), cnt++) {
 	bin_pos = FL_POS(table);
 	old = bin_get(bins, sz, bin_pos);
 	if (old == IDX_NULL || entries[RI(old)].hash == DELETED)
 	    break;
     }
     bin_set(bins, sz, bin_pos, idx+1);
+#if defined(HASH_LOG) && collision_check
+    if (table->type->stronghash && cnt > collision.max) collision.max = cnt;
+#endif
+    return cnt >= ST_STRONG_THRESHOLD;
 }
 
 int
@@ -569,10 +587,15 @@ add_direct(st_table *table, st_data_t key, st_data_t value,
     entry->key = key;
     entry->record = value;
     if (st_sz[table->sz].nbins != 0) {
+	int attack;
 	if (table->num_entries == table->last) {
-	    insert_entry_simple(table, hash_val, en_idx);
+	    attack = insert_entry_simple(table, hash_val, en_idx);
 	} else {
-	    insert_entry_reuse(table, hash_val, en_idx);
+	    attack = insert_entry_reuse(table, hash_val, en_idx);
+	}
+	if (UNLIKELY(attack) && !table->use_strong && table->type->stronghash != NULL) {
+	    table->use_strong = 1;
+	    st_recalc_hashvals(table);
 	}
     }
 }
@@ -725,6 +748,22 @@ st_rehash(register st_table *table)
     }
 }
 
+void
+st_recalc_hashvals(st_table *table)
+{
+    st_idx_t i;
+    st_table_entry *ptr;
+
+    ptr = table->as.entries + table->first;
+    for (i = table->first; i < table->last; i++, ptr++) {
+	if (ptr->hash != DELETED) {
+	    ptr->hash = do_hash(ptr->key, table);
+	}
+    }
+    st_fix_bins(table);
+    table->rebuild_num++;
+}
+
 st_table*
 st_copy(st_table *old_table)
 {
-- 
2.7.4