From 782c82638164f70dbb6179e41221443e39c59ada Mon Sep 17 00:00:00 2001
From: heinrich5991 <heinrich5991@gmail.com>
Date: Sat, 13 Mar 2021 16:52:35 +0100
Subject: [PATCH] Add secure_rand_below function

It works by masking the numbers and the using rejection sampling to
ensure a uniform generation.
---
 CMakeLists.txt             |  1 +
 src/base/system.c          | 28 ++++++++++++++++++++++++++++
 src/base/system.h          | 10 ++++++++++
 src/test/secure_random.cpp | 35 +++++++++++++++++++++++++++++++++++
 src/test/test.cpp          |  5 +++++
 5 files changed, 79 insertions(+)
 create mode 100644 src/test/secure_random.cpp

diff --git a/CMakeLists.txt b/CMakeLists.txt
index a1ce29eeb..af3316000 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -2203,6 +2203,7 @@ if(GTEST_FOUND OR DOWNLOAD_GTEST)
     netaddr.cpp
     packer.cpp
     prng.cpp
+    secure_random.cpp
     sorted_array.cpp
     str.cpp
     strip_path_and_extension.cpp
diff --git a/src/base/system.c b/src/base/system.c
index 22cfb398a..b1f06fded 100644
--- a/src/base/system.c
+++ b/src/base/system.c
@@ -3608,6 +3608,34 @@ int secure_rand(void)
 	return (int)(i % RAND_MAX);
 }
 
+// From https://graphics.stanford.edu/~seander/bithacks.html#RoundUpPowerOf2.
+static unsigned int find_next_power_of_two_minus_one(unsigned int n)
+{
+	n--;
+	n |= n >> 1;
+	n |= n >> 2;
+	n |= n >> 4;
+	n |= n >> 4;
+	n |= n >> 16;
+	return n;
+}
+
+int secure_rand_below(int below)
+{
+	unsigned int mask = find_next_power_of_two_minus_one(below);
+	dbg_assert(below > 0, "below must be positive");
+	while(1)
+	{
+		unsigned int n;
+		secure_random_fill(&n, sizeof(n));
+		n &= mask;
+		if((int)n < below)
+		{
+			return n;
+		}
+	}
+}
+
 #if defined(__cplusplus)
 }
 #endif
diff --git a/src/base/system.h b/src/base/system.h
index dd354e181..9a75aca5d 100644
--- a/src/base/system.h
+++ b/src/base/system.h
@@ -2193,6 +2193,16 @@ void secure_random_fill(void *bytes, unsigned length);
 */
 int secure_rand(void);
 
+/*
+	Function: secure_rand_below
+		Returns a random nonnegative integer below the given number,
+		with a uniform distribution.
+
+	Parameters:
+		below - Upper limit (exclusive) of integers to return.
+*/
+int secure_rand_below(int below);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/src/test/secure_random.cpp b/src/test/secure_random.cpp
new file mode 100644
index 000000000..bf8dd2ab5
--- /dev/null
+++ b/src/test/secure_random.cpp
@@ -0,0 +1,35 @@
+#include "test.h"
+#include <gtest/gtest.h>
+
+#include <base/system.h>
+
+TEST(SecureRandom, Fill)
+{
+	unsigned int Bits = 0;
+	while(~Bits)
+	{
+		unsigned int Random;
+		secure_random_fill(&Random, sizeof(Random));
+		Bits |= Random;
+	}
+}
+
+TEST(SecureRandom, Below1)
+{
+	EXPECT_EQ(secure_rand_below(1), 0);
+}
+
+TEST(SecureRandom, Below)
+{
+	int BOUNDS[] = {2, 3, 4, 5, 10, 100, 127, 128, 129};
+	for(int i = 0; i < sizeof(BOUNDS) / sizeof(BOUNDS[0]); i++)
+	{
+		int Below = BOUNDS[i];
+		for(int i = 0; i < 10; i++)
+		{
+			int Random = secure_rand_below(Below);
+			EXPECT_GE(Random, 0);
+			EXPECT_LT(Random, Below);
+		}
+	}
+}
diff --git a/src/test/test.cpp b/src/test/test.cpp
index 761d661f7..d495dfbef 100644
--- a/src/test/test.cpp
+++ b/src/test/test.cpp
@@ -15,5 +15,10 @@ int main(int argc, char **argv)
 {
 	::testing::InitGoogleTest(&argc, argv);
 	net_init();
+	if(secure_random_init())
+	{
+		fprintf(stderr, "random init failed\n");
+		return 1;
+	}
 	return RUN_ALL_TESTS();
 }