Handle take(-1) better

Summary: It's easy to accidentally pass a negative value to ##take()##, which leads to underflow on conversion to ##size_t##.

Reviewed By: @​rosephilip, @philippv

Differential Revision: D2421459

diff --git a/folly/gen/Base-inl.h b/folly/gen/Base-inl.h
index 23b919735a52ec5c8547aff4c979a12376315bed..7eded731b8e49d3e950ecf839c076ea72d3b4dbb 100644 (file)
--- a/folly/gen/Base-inl.h
+++ b/folly/gen/Base-inl.h
@@ -2308,7 +2308,13 @@ constexpr detail::Indirect indirect{};
 
 constexpr detail::Unwrap unwrap{};
 
-inline detail::Take take(size_t count) { return detail::Take(count); }
+template <class Number>
+inline detail::Take take(Number count) {
+  if (count < 0) {
+    throw std::invalid_argument("Negative value passed to take()");
+  }
+  return detail::Take(static_cast<size_t>(count));
+}
 
 inline detail::Stride stride(size_t s) { return detail::Stride(s); }
 

diff --git a/folly/gen/test/BaseTest.cpp b/folly/gen/test/BaseTest.cpp
index e97cf6c39db9f482f87e64965a03907162aabb38..512dcc9fcffbc4579a1985256f74256f797c733b 100644 (file)
--- a/folly/gen/test/BaseTest.cpp
+++ b/folly/gen/test/BaseTest.cpp
@@ -349,6 +349,11 @@ TEST(Gen, Take) {
       | as<vector>();
     EXPECT_EQ(expected, actual);
   }
+  {
+    int64_t limit = 5;
+    take(limit - 5);
+    EXPECT_THROW(take(limit - 6), std::invalid_argument);
+  }
 }