Add secure_random() and secure_random(lower,upper) Presto functions (f…

…acebookincubator#9295) Summary: facebookincubator#5762 prestodb/presto#20530 Pull Request resolved: facebookincubator#9295 Reviewed By: amitkdutta Differential Revision: D58423391 Pulled By: kevinwilfong fbshipit-source-id: 93d30d872b5a1924c86cb1f2c2d4c057ee91d02a
amitkdutta · Jul 1, 2024 · 181c6ae · 181c6ae
1 parent 4f532a0
commit 181c6ae
Show file tree

Hide file tree

Showing 4 changed files with 258 additions and 0 deletions.
diff --git a/velox/docs/functions/presto/math.rst b/velox/docs/functions/presto/math.rst
@@ -139,6 +139,18 @@ 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]
+
+    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
@@ -58,4 +58,92 @@ struct RandFunction {
   }
 };
 
+template <typename T>
+struct SecureRandFunction {
+  static constexpr bool is_deterministic = false;
+
+  template <typename InputType>
+  FOLLY_ALWAYS_INLINE void checkInput(
+      const InputType lower,
+      const InputType upper) {
+    VELOX_USER_CHECK_GT(
+        upper, lower, "upper bound must be greater than lower bound");
+  }
+
+  FOLLY_ALWAYS_INLINE void call(double& out) {
+    out = folly::Random::secureRandDouble01();
+  }
+
+  FOLLY_ALWAYS_INLINE void
+  call(double& out, const double lower, const double upper) {
+    checkInput(lower, upper);
+    if (std::isinf(lower)) {
+      out = std::numeric_limits<double>::quiet_NaN();
+    } else if (std::isinf(upper)) {
+      out = std::numeric_limits<double>::max();
+    } else {
+      out = folly::Random::secureRandDouble(0, upper - lower) + lower;
+    }
+  }
+
+  FOLLY_ALWAYS_INLINE void
+  call(float& out, const float lower, const float upper) {
+    checkInput(lower, upper);
+    if (std::isinf(lower)) {
+      out = std::numeric_limits<float>::quiet_NaN();
+    } else if (std::isinf(upper)) {
+      out = std::numeric_limits<float>::max();
+    } else {
+      out = float(folly::Random::secureRandDouble(0, upper - lower)) + lower;
+    }
+  }
+
+  FOLLY_ALWAYS_INLINE void
+  call(int64_t& out, const int64_t lower, const int64_t upper) {
+    checkInput(lower, upper);
+    if (upper >= 0) {
+      out = int64_t(
+          folly::Random::secureRand64(0, uint64_t(upper) - lower) + lower);
+    } else {
+      out = int64_t(
+          folly::Random::secureRand64(0, uint64_t(upper - lower)) + lower);
+    }
+  }
+
+  FOLLY_ALWAYS_INLINE void
+  call(int32_t& out, const int32_t lower, const int32_t upper) {
+    checkInput(lower, upper);
+    if (upper >= 0) {
+      out = int32_t(
+          folly::Random::secureRand32(0, uint32_t(upper) - lower) + lower);
+    } else {
+      out = int32_t(
+          folly::Random::secureRand32(0, uint32_t(upper - lower)) + lower);
+    }
+  }
+
+  FOLLY_ALWAYS_INLINE void
+  call(int16_t& out, const int16_t lower, const int16_t upper) {
+    checkInput(lower, upper);
+    if (upper >= 0) {
+      out = int16_t(
+          folly::Random::secureRand32(0, uint16_t(upper) - lower) + lower);
+    } else {
+      out = int16_t(
+          folly::Random::secureRand32(0, uint16_t(upper - lower)) + lower);
+    }
+  }
+
+  FOLLY_ALWAYS_INLINE void
+  call(int8_t& out, const int8_t lower, const int8_t upper) {
+    checkInput(lower, upper);
+    if (upper >= 0) {
+      out = int8_t(
+          folly::Random::secureRand32(0, uint8_t(upper) - lower) + lower);
+    } else {
+      out = int8_t(
+          folly::Random::secureRand32(0, uint8_t(upper - lower)) + lower);
+    }
+  }
+};
 } // namespace facebook::velox::functions
diff --git a/velox/functions/prestosql/registration/MathematicalFunctionsRegistration.cpp b/velox/functions/prestosql/registration/MathematicalFunctionsRegistration.cpp
@@ -100,6 +100,10 @@ void registerMathFunctions(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
@@ -23,6 +23,13 @@ namespace facebook::velox::functions {
 
 namespace {
 
+constexpr double kInf = std::numeric_limits<double>::infinity();
+constexpr double kNan = std::numeric_limits<double>::quiet_NaN();
+constexpr float kInfF = std::numeric_limits<float>::infinity();
+constexpr float kNanF = std::numeric_limits<float>::quiet_NaN();
+constexpr int64_t kLongMax = std::numeric_limits<int64_t>::max();
+constexpr int64_t kLongMin = std::numeric_limits<int64_t>::min();
+
 class RandTest : public functions::test::FunctionBaseTest {
  protected:
   template <typename T>
@@ -44,6 +51,18 @@ class RandTest : public functions::test::FunctionBaseTest {
   std::optional<T> randWithTry(T n) {
     return evaluateOnce<T>("try(rand(c0))", std::make_optional(n));
   }
+
+  template <typename T>
+  std::optional<T> secureRandom(
+      std::optional<T> lower,
+      std::optional<T> upper) {
+    return evaluateOnce<T>("secure_random(c0, c1)", lower, upper);
+  }
+
+  template <typename T>
+  std::optional<T> secureRand(std::optional<T> lower, std::optional<T> upper) {
+    return evaluateOnce<T>("secure_rand(c0, c1)", lower, upper);
+  }
 };
 
 TEST_F(RandTest, zeroArg) {
@@ -74,5 +93,140 @@ TEST_F(RandTest, nonNullInt8) {
   EXPECT_LT(rand(4), 4);
 }
 
+TEST_F(RandTest, secureRandZeroArg) {
+  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(RandTest, secureRandInt64) {
+  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);
+
+  result = secureRand<int64_t>(
+      std::numeric_limits<int64_t>::min(), std::numeric_limits<int64_t>::max());
+  EXPECT_LT(result, std::numeric_limits<int64_t>::max());
+  EXPECT_GE(result, std::numeric_limits<int64_t>::min());
+}
+
+TEST_F(RandTest, secureRandInt32) {
+  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(RandTest, secureRandInt16) {
+  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(RandTest, secureRandInt8) {
+  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(RandTest, secureRandDouble) {
+  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(RandTest, secureRandFloat) {
+  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(RandTest, secureRandInvalid) {
+  VELOX_ASSERT_THROW(
+      secureRand<int64_t>((int64_t)-5, (int64_t)-10),
+      "upper bound must be greater than lower bound");
+  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),
+      "upper bound must be greater than lower bound");
+  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),
+      "upper bound must be greater than lower bound");
+  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),
+      "upper bound must be greater than lower bound");
+  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),
+      "upper bound must be greater than lower bound");
+  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),
+      "upper bound must be greater than lower bound");
+  VELOX_ASSERT_THROW(
+      secureRand<float>((float)15.6, (float)10.1),
+      "upper bound must be greater than lower bound");
+}
+
+TEST_F(RandTest, secureRandSpecialValues) {
+  EXPECT_LT(secureRand<int64_t>(0, kLongMax), kLongMax);
+  EXPECT_GE(secureRand<int64_t>(kLongMin, 0), kLongMin);
+
+  EXPECT_TRUE(std::isnan(secureRand<float>(-kInfF, 0).value()));
+  EXPECT_LE(secureRand<float>(0.0, kInfF), std::numeric_limits<float>::max());
+
+  EXPECT_TRUE(std::isnan(secureRand<double>(-kInf, 0).value()));
+  EXPECT_EQ(secureRand<double>(0.0, kInf), std::numeric_limits<double>::max());
+
+  VELOX_ASSERT_THROW(
+      secureRand<double>(0.0, kNan),
+      "upper bound must be greater than lower bound");
+
+  VELOX_ASSERT_THROW(
+      secureRand<float>(0.0, kNanF),
+      "upper bound must be greater than lower bound");
+}
+
 } // namespace
 } // namespace facebook::velox::functions