Implementation of secure_random() and secure_random(lower,upper) func…

…tions
facebookincubator · Sep 17, 2023 · 9b08171 · 9b08171
1 parent aa7cd61
commit 9b08171
Show file tree

Hide file tree

Showing 4 changed files with 215 additions and 0 deletions.
diff --git a/velox/docs/functions/presto/math.rst b/velox/docs/functions/presto/math.rst
@@ -128,6 +128,19 @@ Mathematical Functions
 
     Returns ``x`` rounded to ``d`` decimal places.
 
+.. function:: secure_rand() -> double
+
+    This is an alias for :func:`secure_random()`.
+
+.. function:: secure_random() -> double
+
+    Returns a cryptographically secure random value in the range 0.0 <= x < 1.0.
+
+.. function:: secure_random(lower, upper) → [same as input]
+   :noindex:
+
+    Returns a cryptographically secure random value in the range lower <= x < upper, where lower < upper.
+
 .. function:: sign(x) -> [same as x]
 
     Returns the signum function of ``x``. For both integer and floating point arguments, it returns:

diff --git a/velox/functions/prestosql/Rand.h b/velox/functions/prestosql/Rand.h
@@ -56,4 +56,67 @@ struct RandFunction {
   }
 };
 
+template <typename T>
+struct SecureRandFunction {
+  static constexpr bool is_deterministic = false;
+
+  FOLLY_ALWAYS_INLINE void call(double& result) {
+    result = folly::Random::secureRandDouble01();
+  }
+
+  FOLLY_ALWAYS_INLINE void
+  call(double& out, const double lower, const double upper) {
+    VELOX_USER_CHECK_GE(lower, 0.0, "lower bound must be positive");
+    VELOX_USER_CHECK_GT(upper, 0.0, "upper bound must be positive");
+    VELOX_USER_CHECK_GT(
+        upper, lower, "upper bound must be greater than lower bound");
+    out = folly::Random::secureRandDouble(lower, upper);
+  }
+
+  FOLLY_ALWAYS_INLINE void
+  call(float& out, const float lower, const float upper) {
+    VELOX_USER_CHECK_GE(lower, 0.0, "lower bound must be positive");
+    VELOX_USER_CHECK_GT(upper, 0.0, "upper bound must be positive");
+    VELOX_USER_CHECK_GT(
+        upper, lower, "upper bound must be greater than lower bound");
+    out = float(folly::Random::secureRandDouble(lower, upper));
+  }
+
+  FOLLY_ALWAYS_INLINE void
+  call(int64_t& out, const int64_t lower, const int64_t upper) {
+    VELOX_USER_CHECK_GE(lower, 0, "lower bound must be positive");
+    VELOX_USER_CHECK_GT(upper, 0, "upper bound must be positive");
+    VELOX_USER_CHECK_GT(
+        upper, lower, "upper bound must be greater than lower bound");
+    out = folly::Random::secureRand64(lower, upper);
+  }
+
+  FOLLY_ALWAYS_INLINE void
+  call(int32_t& out, const int32_t lower, const int32_t upper) {
+    VELOX_USER_CHECK_GE(lower, 0, "lower bound must be positive");
+    VELOX_USER_CHECK_GT(upper, 0, "upper bound must be positive");
+    VELOX_USER_CHECK_GT(
+        upper, lower, "upper bound must be greater than lower bound");
+    out = folly::Random::secureRand32(lower, upper);
+  }
+
+  FOLLY_ALWAYS_INLINE void
+  call(int16_t& out, const int16_t lower, const int16_t upper) {
+    VELOX_USER_CHECK_GE(lower, 0, "lower bound must be positive");
+    VELOX_USER_CHECK_GT(upper, 0, "upper bound must be positive");
+    VELOX_USER_CHECK_GT(
+        upper, lower, "upper bound must be greater than lower bound");
+    out = int16_t(folly::Random::secureRand32(lower, upper));
+  }
+
+  FOLLY_ALWAYS_INLINE void
+  call(int8_t& out, const int8_t lower, const int8_t upper) {
+    VELOX_USER_CHECK_GE(lower, 0, "lower bound must be positive");
+    VELOX_USER_CHECK_GT(upper, 0, "upper bound must be positive");
+    VELOX_USER_CHECK_GT(
+        upper, lower, "upper bound must be greater than lower bound");
+    out = int8_t(folly::Random::secureRand32(lower, upper));
+  }
+};
+
 } // namespace facebook::velox::functions
diff --git a/velox/functions/prestosql/registration/ArithmeticFunctionsRegistration.cpp b/velox/functions/prestosql/registration/ArithmeticFunctionsRegistration.cpp
@@ -92,6 +92,10 @@ void registerSimpleFunctions(const std::string& prefix) {
   registerFunction<NanFunction, double>({prefix + "nan"});
   registerFunction<RandFunction, double>({prefix + "rand", prefix + "random"});
   registerUnaryIntegral<RandFunction>({prefix + "rand", prefix + "random"});
+  registerFunction<SecureRandFunction, double>(
+      {prefix + "secure_rand", prefix + "secure_random"});
+  registerBinaryNumeric<SecureRandFunction>(
+      {prefix + "secure_rand", prefix + "secure_random"});
   registerFunction<FromBaseFunction, int64_t, Varchar, int64_t>(
       {prefix + "from_base"});
   registerFunction<ToBaseFunction, Varchar, int64_t, int64_t>(

diff --git a/velox/functions/prestosql/tests/RandTest.cpp b/velox/functions/prestosql/tests/RandTest.cpp
@@ -74,5 +74,140 @@ TEST_F(RandTest, nonNullInt8) {
   EXPECT_LT(rand(4), 4);
 }
 
+class SecureRandTest : public functions::test::FunctionBaseTest {
+ protected:
+  template <typename T>
+  std::optional<T> secureRand(std::optional<T> lower, std::optional<T> upper) {
+    return evaluateOnce<T>("secure_random(c0, c1)", lower, upper);
+  };
+};
+
+TEST_F(SecureRandTest, zeroArgTest) {
+  auto result =
+      evaluateOnce<double>("secure_random()", makeRowVector(ROW({}), 1));
+  EXPECT_LT(result, 1.0);
+  EXPECT_GE(result, 0.0);
+
+  result = evaluateOnce<double>("secure_rand()", makeRowVector(ROW({}), 1));
+  EXPECT_LT(result, 1.0);
+  EXPECT_GE(result, 0.0);
+}
+
+TEST_F(SecureRandTest, int64Test) {
+  auto result =
+      secureRand<int64_t>((int64_t)2147532562, (int64_t)4611791058295013614);
+  EXPECT_LT(result, 4611791058295013614);
+  EXPECT_GE(result, 2147532562);
+
+  result = secureRand<int64_t>((int64_t)0, (int64_t)46117910582950136);
+  EXPECT_LT(result, 46117910582950136);
+  EXPECT_GE(result, 0);
+}
+
+TEST_F(SecureRandTest, int32Test) {
+  auto result = secureRand<int32_t>((int32_t)8765432, (int32_t)2145613151);
+  EXPECT_LT(result, 2145613151);
+  EXPECT_GE(result, 8765432);
+
+  result = secureRand<int32_t>((int32_t)0, (int32_t)21456131);
+  EXPECT_LT(result, 21456131);
+  EXPECT_GE(result, 0);
+}
+
+TEST_F(SecureRandTest, int16Test) {
+  auto result = secureRand<int16_t>((int16_t)100, (int16_t)23286);
+  EXPECT_LT(result, 23286);
+  EXPECT_GE(result, 100);
+
+  result = secureRand<int16_t>((int16_t)0, (int16_t)23286);
+  EXPECT_LT(result, 23286);
+  EXPECT_GE(result, 0);
+}
+
+TEST_F(SecureRandTest, int8Test) {
+  auto result = secureRand<int8_t>((int8_t)10, (int8_t)120);
+  EXPECT_LT(result, 120);
+  EXPECT_GE(result, 10);
+
+  result = secureRand<int8_t>((int8_t)0, (int8_t)120);
+  EXPECT_LT(result, 120);
+  EXPECT_GE(result, 0);
+}
+
+TEST_F(SecureRandTest, doubleTest) {
+  auto result = secureRand<double>((double)10.5, (double)120.7895);
+  EXPECT_LT(result, 120.7895);
+  EXPECT_GE(result, 10.5);
+
+  result = secureRand<double>((double)0.0, (double)120.7895);
+  EXPECT_LT(result, 120.7895);
+  EXPECT_GE(result, 0.0);
+}
+
+TEST_F(SecureRandTest, floatTest) {
+  auto result = secureRand<float>((float)10.5, (float)120.7);
+  EXPECT_LT(result, 120.7);
+  EXPECT_GE(result, 10.5);
+
+  result = secureRand<float>((float)0.0, (float)120.7);
+  EXPECT_LT(result, 120.7);
+  EXPECT_GE(result, 0.0);
+}
+
+TEST_F(SecureRandTest, invalidBounds) {
+  VELOX_ASSERT_THROW(
+      secureRand<int64_t>((int64_t)-5, (int64_t)10),
+      "lower bound must be positive");
+  VELOX_ASSERT_THROW(
+      secureRand<int64_t>((int64_t)0, (int64_t)-1),
+      "upper bound must be positive");
+  VELOX_ASSERT_THROW(
+      secureRand<int64_t>((int64_t)15, (int64_t)10),
+      "upper bound must be greater than lower bound");
+  VELOX_ASSERT_THROW(
+      secureRand<int32_t>((int32_t)-5, (int32_t)10),
+      "lower bound must be positive");
+  VELOX_ASSERT_THROW(
+      secureRand<int32_t>((int32_t)0, (int32_t)-1),
+      "upper bound must be positive");
+  VELOX_ASSERT_THROW(
+      secureRand<int32_t>((int32_t)15, (int32_t)10),
+      "upper bound must be greater than lower bound");
+  VELOX_ASSERT_THROW(
+      secureRand<int16_t>((int16_t)-5, (int16_t)10),
+      "lower bound must be positive");
+  VELOX_ASSERT_THROW(
+      secureRand<int16_t>((int16_t)0, (int16_t)-1),
+      "upper bound must be positive");
+  VELOX_ASSERT_THROW(
+      secureRand<int16_t>((int16_t)15, (int16_t)10),
+      "upper bound must be greater than lower bound");
+  VELOX_ASSERT_THROW(
+      secureRand<int8_t>((int8_t)-5, (int8_t)10),
+      "lower bound must be positive");
+  VELOX_ASSERT_THROW(
+      secureRand<int8_t>((int8_t)0, (int8_t)-1),
+      "upper bound must be positive");
+  VELOX_ASSERT_THROW(
+      secureRand<int8_t>((int8_t)15, (int8_t)10),
+      "upper bound must be greater than lower bound");
+  VELOX_ASSERT_THROW(
+      secureRand<double>(-5.7, 10.7), "lower bound must be positive");
+  VELOX_ASSERT_THROW(
+      secureRand<double>(0.0, -1.0), "upper bound must be positive");
+  VELOX_ASSERT_THROW(
+      secureRand<double>(15.6, 10.1),
+      "upper bound must be greater than lower bound");
+  VELOX_ASSERT_THROW(
+      secureRand<float>((float)-5.7, (float)10.7),
+      "lower bound must be positive");
+  VELOX_ASSERT_THROW(
+      secureRand<float>((float)0.0, (float)-1.0),
+      "upper bound must be positive");
+  VELOX_ASSERT_THROW(
+      secureRand<float>((float)15.6, (float)10.1),
+      "upper bound must be greater than lower bound");
+}
+
 } // namespace
 } // namespace facebook::velox::functions