Fix for PR26812: possible overflow issue in std::allocator::allocate

llvm-svn: 262610

2 files changed, 64 insertions, 2 deletions

diff --git a/libcxx/include/memory b/libcxx/include/memory
index 281c59eb4c3..67f2fc53ad5 100644
--- a/libcxx/include/memory
+++ b/libcxx/include/memory
@@ -1726,7 +1726,15 @@ public:
     _LIBCPP_INLINE_VISIBILITY const_pointer address(const_reference __x) const _NOEXCEPT
         {return _VSTD::addressof(__x);}
     _LIBCPP_INLINE_VISIBILITY pointer allocate(size_type __n, allocator<void>::const_pointer = 0)
-        {return static_cast<pointer>(_VSTD::__allocate(__n * sizeof(_Tp)));}
+        {
+        if (__n > max_size())
+#ifndef _LIBCPP_NO_EXCEPTIONS
+            throw bad_alloc();
+#else
+            assert(!"allocator<T>::allocate::bad_alloc");
+#endif
+        return static_cast<pointer>(_VSTD::__allocate(__n * sizeof(_Tp)));
+        }
     _LIBCPP_INLINE_VISIBILITY void deallocate(pointer __p, size_type) _NOEXCEPT
         {_VSTD::__deallocate((void*)__p);}
     _LIBCPP_INLINE_VISIBILITY size_type max_size() const _NOEXCEPT
@@ -1817,7 +1825,15 @@ public:
     _LIBCPP_INLINE_VISIBILITY const_pointer address(const_reference __x) const _NOEXCEPT
         {return _VSTD::addressof(__x);}
     _LIBCPP_INLINE_VISIBILITY pointer allocate(size_type __n, allocator<void>::const_pointer = 0)
-        {return static_cast<pointer>(_VSTD::__allocate(__n * sizeof(_Tp)));}
+    {
+        if (__n > max_size())
+#ifndef _LIBCPP_NO_EXCEPTIONS
+            throw bad_alloc();
+#else
+            assert(!"allocator<const T>::allocate::bad_alloc");
+#endif
+        return static_cast<pointer>(_VSTD::__allocate(__n * sizeof(_Tp)));
+        }
     _LIBCPP_INLINE_VISIBILITY void deallocate(pointer __p, size_type) _NOEXCEPT
         {_VSTD::__deallocate((void*)__p);}
     _LIBCPP_INLINE_VISIBILITY size_type max_size() const _NOEXCEPT
diff --git a/libcxx/test/std/utilities/memory/default.allocator/allocator.members/allocate.size.pass.cpp b/libcxx/test/std/utilities/memory/default.allocator/allocator.members/allocate.size.pass.cpp
new file mode 100644
index 00000000000..4846b0a496b
--- /dev/null
+++ b/libcxx/test/std/utilities/memory/default.allocator/allocator.members/allocate.size.pass.cpp
@@ -0,0 +1,46 @@
+//===----------------------------------------------------------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is dual licensed under the MIT and the University of Illinois Open
+// Source Licenses. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+// <memory>
+
+// allocator:
+// pointer allocate(size_type n, allocator<void>::const_pointer hint=0);
+
+#include <memory>
+#include <cassert>
+
+template <typename T>
+void test_max(size_t count)
+{
+    std::allocator<T> a;
+    try { a.allocate( count ); }
+    catch ( const std::bad_alloc &) { return ; }
+    assert (false);
+}
+
+int main()
+{
+    {  // Bug 26812 -- allocating too large
+        typedef double T;
+        std::allocator<T> a;
+        test_max<T> (a.max_size() + 1);                // just barely too large
+        test_max<T> (a.max_size() * 2);                // significantly too large
+        test_max<T> (((size_t) -1) / sizeof(T) + 1);   // multiply will overflow
+        test_max<T> ((size_t) -1);                     // way too large
+    }
+
+    {
+        typedef const double T;
+        std::allocator<T> a;
+        test_max<T> (a.max_size() + 1);                // just barely too large
+        test_max<T> (a.max_size() * 2);                // significantly too large
+        test_max<T> (((size_t) -1) / sizeof(T) + 1);   // multiply will overflow
+        test_max<T> ((size_t) -1);                     // way too large
+    }
+}