[libc++] Fix potential OOB in poisson_distribution

See details in the original Chromium bug report:
    https://bugs.chromium.org/p/chromium/issues/detail?id=994957

2 files changed, 74 insertions, 1 deletions

diff --git a/libcxx/test/std/numerics/rand/rand.dis/rand.dist.bern/rand.dist.bern.geo/eval.pass.cpp b/libcxx/test/std/numerics/rand/rand.dis/rand.dist.bern/rand.dist.bern.geo/eval.pass.cpp
index bdeaa205404..11eeb97ccb1 100644
--- a/libcxx/test/std/numerics/rand/rand.dis/rand.dist.bern/rand.dist.bern.geo/eval.pass.cpp
+++ b/libcxx/test/std/numerics/rand/rand.dis/rand.dist.bern/rand.dist.bern.geo/eval.pass.cpp
@@ -30,6 +30,16 @@ sqr(T x)
     return x * x;
 }
 
+void test_small_inputs() {
+  std::mt19937 engine;
+  std::geometric_distribution<std::int16_t> distribution(5.45361e-311);
+  typedef std::geometric_distribution<std::int16_t>::result_type result_type;
+  for (int i = 0; i < 1000; ++i) {
+    volatile result_type res = distribution(engine);
+    ((void)res);
+  }
+}
+
 void
 test1()
 {
@@ -296,6 +306,7 @@ int main(int, char**)
     test4();
     test5();
     test6();
+    test_small_inputs();
 
   return 0;
 }
diff --git a/libcxx/test/std/numerics/rand/rand.dis/rand.dist.pois/rand.dist.pois.poisson/eval.pass.cpp b/libcxx/test/std/numerics/rand/rand.dis/rand.dist.pois/rand.dist.pois.poisson/eval.pass.cpp
index 4b07d239364..00a70ac61b5 100644
--- a/libcxx/test/std/numerics/rand/rand.dis/rand.dist.pois/rand.dist.pois.poisson/eval.pass.cpp
+++ b/libcxx/test/std/numerics/rand/rand.dis/rand.dist.pois/rand.dist.pois.poisson/eval.pass.cpp
@@ -30,6 +30,67 @@ sqr(T x)
     return x * x;
 }
 
+void test_bad_ranges() {
+  // Test cases where the mean is around the largest representable integer for
+  // `result_type`. These cases don't generate valid poisson distributions, but
+  // at least they don't blow up.
+  std::mt19937 eng;
+
+  {
+    std::poisson_distribution<std::int16_t> distribution(32710.9);
+    for (int i=0; i < 1000; ++i) {
+      volatile std::int16_t res = distribution(eng);
+      ((void)res);
+    }
+  }
+  {
+    std::poisson_distribution<std::int16_t> distribution(std::numeric_limits<std::int16_t>::max());
+    for (int i=0; i < 1000; ++i) {
+      volatile std::int16_t res = distribution(eng);
+      ((void)res);
+    }
+  }
+  {
+    std::poisson_distribution<std::int16_t> distribution(
+    static_cast<double>(std::numeric_limits<std::int16_t>::max()) + 10);
+    for (int i=0; i < 1000; ++i) {
+      volatile std::int16_t res = distribution(eng);
+      ((void)res);
+    }
+  }
+  {
+    std::poisson_distribution<std::int16_t> distribution(
+      static_cast<double>(std::numeric_limits<std::int16_t>::max()) * 2);
+      for (int i=0; i < 1000; ++i) {
+        volatile std::int16_t res = distribution(eng);
+        ((void)res);
+      }
+  }
+  {
+    // We convert `INF` to `DBL_MAX` otherwise the distribution will hang.
+    std::poisson_distribution<std::int16_t> distribution(std::numeric_limits<double>::infinity());
+    for (int i=0; i < 1000; ++i) {
+      volatile std::int16_t res = distribution(eng);
+      ((void)res);
+    }
+  }
+  {
+    std::poisson_distribution<std::int16_t> distribution(0);
+    for (int i=0; i < 1000; ++i) {
+      volatile std::int16_t res = distribution(eng);
+      ((void)res);
+    }
+  }
+  {
+    // We convert `INF` to `DBL_MAX` otherwise the distribution will hang.
+    std::poisson_distribution<std::int16_t> distribution(-100);
+    for (int i=0; i < 1000; ++i) {
+      volatile std::int16_t res = distribution(eng);
+      ((void)res);
+    }
+  }
+}
+
 int main(int, char**)
 {
     {
@@ -150,5 +211,6 @@ int main(int, char**)
         assert(std::abs((kurtosis - x_kurtosis) / x_kurtosis) < 0.01);
     }
 
-  return 0;
+    test_bad_ranges();
+    return 0;
 }