From 2b7c9362bbf0370b4f819d244375844946391c6a Mon Sep 17 00:00:00 2001
From: Daniel Micay <danielmicay@gmail.com>
Date: Sat, 25 Aug 2018 18:32:44 -0400
Subject: [PATCH] consume fewer random bytes for slab randomization

---
 malloc.c |  4 ++--
 random.c | 40 +++++++++++++++++++++++++++++++---------
 random.h |  7 ++++---
 3 files changed, 37 insertions(+), 14 deletions(-)

diff --git a/malloc.c b/malloc.c
index 603d476..71f77a1 100644
--- a/malloc.c
+++ b/malloc.c
@@ -257,7 +257,7 @@ static size_t get_free_slot(struct random_state *rng, size_t slots, struct slab_
     }
 
     // randomize start location for linear search (uniform random choice is too slow)
-    uint64_t random_split = ~(~0UL << get_random_size_uniform(rng, slots));
+    uint64_t random_split = ~(~0UL << get_random_u16_uniform(rng, slots));
 
     size_t slot = ffzl(masked | random_split);
     if (slot) {
@@ -628,7 +628,7 @@ COLD static void init_slow_path(void) {
         random_state_init(&c->rng);
 
         size_t bound = (real_class_region_size - class_region_size) / PAGE_SIZE - 1;
-        size_t gap = (get_random_size_uniform(&rng, bound) + 1) * PAGE_SIZE;
+        size_t gap = (get_random_u64_uniform(&rng, bound) + 1) * PAGE_SIZE;
         c->class_region_start = (char *)ro.slab_region_start + real_class_region_size * i + gap;
 
         size_t size = size_classes[i];
diff --git a/random.c b/random.c
index d3f6f0f..9cee055 100644
--- a/random.c
+++ b/random.c
@@ -16,7 +16,7 @@ static ssize_t getrandom(void *buf, size_t buflen, unsigned int flags) {
 }
 #endif
 
-void get_random_seed(void *buf, size_t size) {
+static void get_random_seed(void *buf, size_t size) {
     while (size > 0) {
         ssize_t r;
 
@@ -56,23 +56,45 @@ void get_random_bytes(struct random_state *state, void *buf, size_t size) {
     }
 }
 
-size_t get_random_size(struct random_state *state) {
-    size_t size;
-    get_random_bytes(state, &size, sizeof(size));
-    return size;
+uint16_t get_random_u16(struct random_state *state) {
+    uint16_t value;
+    get_random_bytes(state, &value, sizeof(value));
+    return value;
 }
 
 // based on OpenBSD arc4random_uniform
-size_t get_random_size_uniform(struct random_state *state, size_t bound) {
+uint16_t get_random_u16_uniform(struct random_state *state, uint16_t bound) {
     if (bound < 2) {
         return 0;
     }
 
-    size_t min = -bound % bound;
+    uint16_t min = (uint16_t)-bound % bound;
 
-    size_t r;
+    uint16_t r;
     do {
-        r = get_random_size(state);
+        r = get_random_u16(state);
+    } while (r < min);
+
+    return r % bound;
+}
+
+uint64_t get_random_u64(struct random_state *state) {
+    uint64_t value;
+    get_random_bytes(state, &value, sizeof(value));
+    return value;
+}
+
+// based on OpenBSD arc4random_uniform
+uint64_t get_random_u64_uniform(struct random_state *state, uint64_t bound) {
+    if (bound < 2) {
+        return 0;
+    }
+
+    uint64_t min = -bound % bound;
+
+    uint64_t r;
+    do {
+        r = get_random_u64(state);
     } while (r < min);
 
     return r % bound;
diff --git a/random.h b/random.h
index 40f97ee..4568401 100644
--- a/random.h
+++ b/random.h
@@ -10,10 +10,11 @@ struct random_state {
     uint8_t cache[RANDOM_CACHE_SIZE];
 };
 
-void get_random_seed(void *buf, size_t size);
 void random_state_init(struct random_state *state);
 void get_random_bytes(struct random_state *state, void *buf, size_t size);
-size_t get_random_size(struct random_state *state);
-size_t get_random_size_uniform(struct random_state *state, size_t bound);
+uint16_t get_random_u16(struct random_state *state);
+uint16_t get_random_u16_uniform(struct random_state *state, uint16_t bound);
+uint64_t get_random_u64(struct random_state *state);
+uint64_t get_random_u64_uniform(struct random_state *state, uint64_t bound);
 
 #endif