check integer overflow at Memory class

src/IO/BufferWithOwnMemory.h

@@ -8,6 +8,8 @@
 #include <Common/Exception.h>
 #include <Core/Defines.h>
 
+#include <base/arithmeticOverflow.h>
+
 
 namespace ProfileEvents
 {
@@ -19,6 +21,11 @@ namespace ProfileEvents
 namespace DB
 {
 
+namespace ErrorCodes
+{
+    extern const int ARGUMENT_OUT_OF_BOUND;
+}
+
 
 /** Replacement for std::vector<char> to use in buffers.
   * Differs in that is doesn't do unneeded memset. (And also tries to do as little as possible.)
@@ -38,9 +45,9 @@ struct Memory : boost::noncopyable, Allocator
     Memory() = default;
 
     /// If alignment != 0, then allocate memory aligned to specified value.
-    explicit Memory(size_t size_, size_t alignment_ = 0) : m_capacity(size_), m_size(m_capacity), alignment(alignment_)
+    explicit Memory(size_t size_, size_t alignment_ = 0) : alignment(alignment_)
     {
-        alloc();
+        alloc(size_);
     }
 
     ~Memory()
@@ -75,28 +82,26 @@ struct Memory : boost::noncopyable, Allocator
 
     void resize(size_t new_size)
     {
-        if (0 == m_capacity)
+        if (!m_data)
         {
-            m_size = new_size;
-            m_capacity = new_size;
-            alloc();
+            alloc(new_size);
+            return;
         }
-        else if (new_size <= m_capacity - pad_right)
+
+        if (new_size <= m_capacity - pad_right)
         {
             m_size = new_size;
             return;
         }
-        else
-        {
-            size_t new_capacity = align(new_size, alignment) + pad_right;
 
-            size_t diff = new_capacity - m_capacity;
-            ProfileEvents::increment(ProfileEvents::IOBufferAllocBytes, diff);
+        size_t new_capacity = alignWithPadding(new_size, alignment);
 
-            m_data = static_cast<char *>(Allocator::realloc(m_data, m_capacity, new_capacity, alignment));
-            m_capacity = new_capacity;
-            m_size = m_capacity - pad_right;
-        }
+        size_t diff = new_capacity - m_capacity;
+        ProfileEvents::increment(ProfileEvents::IOBufferAllocBytes, diff);
+
+        m_data = static_cast<char *>(Allocator::realloc(m_data, m_capacity, new_capacity, alignment));
+        m_capacity = new_capacity;
+        m_size = new_size;
     }
 
 private:
@@ -108,24 +113,47 @@ private:
         if (!(value % alignment))
             return value;
 
-        return (value + alignment - 1) / alignment * alignment;
+        // original expression is (value + alignment - 1) / alignment * alignment;
+
+        size_t res = 0;
+
+        if (common::addOverflow<size_t>(value, alignment - 1, res))
+            throw Exception("value is too big to apply alignment", ErrorCodes::ARGUMENT_OUT_OF_BOUND);
+
+        res /= alignment;
+
+        if (common::mulOverflow<size_t>(res, alignment, res))
+            throw Exception("value is too big to apply alignment", ErrorCodes::ARGUMENT_OUT_OF_BOUND);
+
+        return res;
     }
 
-    void alloc()
+    static size_t alignWithPadding(const size_t value, const size_t alignment)
     {
-        if (!m_capacity)
+        size_t res = align(value, alignment);
+
+        if (common::addOverflow<size_t>(res, pad_right, res))
+            throw Exception("value is too big to apply padding 3", ErrorCodes::ARGUMENT_OUT_OF_BOUND);
+
+        return res;
+    }
+
+    void alloc(size_t new_size)
+    {
+        if (!new_size)
         {
             m_data = nullptr;
             return;
         }
 
-        ProfileEvents::increment(ProfileEvents::IOBufferAllocs);
-        ProfileEvents::increment(ProfileEvents::IOBufferAllocBytes, m_capacity);
+        size_t new_capacity = alignWithPadding(new_size, alignment);
+
+        ProfileEvents::increment(ProfileEvents::IOBufferAllocs);
+        ProfileEvents::increment(ProfileEvents::IOBufferAllocBytes, new_capacity);
 
-        size_t new_capacity = align(m_capacity, alignment) + pad_right;
         m_data = static_cast<char *>(Allocator::alloc(new_capacity, alignment));
         m_capacity = new_capacity;
-        m_size = m_capacity - pad_right;
+        m_size = new_size;
     }
 
     void dealloc()

src/IO/tests/gtest_memory_resize.cpp

@@ -0,0 +1,326 @@
+#include <IO/WriteHelpers.h>
+#include <IO/ReadHelpers.h>
+#include <IO/BufferWithOwnMemory.h>
+#include <gtest/gtest.h>
+
+#define EXPECT_THROW_ERROR_CODE(statement, expected_exception, expected_code)    \
+    EXPECT_THROW(                                                                \
+        try                                                                      \
+        {                                                                        \
+            statement;                                                           \
+        }                                                                        \
+        catch (const expected_exception & e)                                     \
+        {                                                                        \
+            EXPECT_EQ(expected_code, e.code());                                  \
+            throw;                                                               \
+        }                                                                        \
+    , expected_exception)
+
+namespace DB
+{
+
+namespace ErrorCodes
+{
+    extern const int ARGUMENT_OUT_OF_BOUND;
+    extern const int LOGICAL_ERROR;
+    extern const int CANNOT_ALLOCATE_MEMORY;
+}
+
+}
+
+using namespace DB;
+
+class DummyAllocator
+{
+    void * DUMMY_ADDRESS = reinterpret_cast<void *>(1);
+
+public:
+    void * alloc(size_t size, size_t /*alignment*/ = 0)
+    {
+        checkSize(size);
+        if (size)
+            return DUMMY_ADDRESS;
+        else
+            return nullptr;
+    }
+
+    void * realloc(void * /*buf*/, size_t /*old_size*/, size_t new_size, size_t /*alignment*/ = 0)
+    {
+        checkSize(new_size);
+        return DUMMY_ADDRESS;
+    }
+
+    void free([[maybe_unused]] void * buf, size_t /*size*/)
+    {
+        assert(buf == DUMMY_ADDRESS);
+    }
+
+    // the same check as in Common/Allocator.h
+    void checkSize(size_t size)
+    {
+        /// More obvious exception in case of possible overflow (instead of just "Cannot mmap").
+        if (size >= 0x8000000000000000ULL)
+            throw DB::Exception(DB::ErrorCodes::LOGICAL_ERROR, "Too large size ({}) passed to allocator. It indicates an error.", size);
+    }
+};
+
+TEST(MemoryResizeTest, SmallInitAndSmallResize)
+{
+    {
+        auto memory = Memory<DummyAllocator>(0);
+        ASSERT_EQ(memory.m_data, nullptr);
+        ASSERT_EQ(memory.m_capacity, 0);
+        ASSERT_EQ(memory.m_size, 0);
+
+        memory.resize(0);
+        ASSERT_EQ(memory.m_data, nullptr);
+        ASSERT_EQ(memory.m_capacity, 0);
+        ASSERT_EQ(memory.m_size, 0);
+
+        memory.resize(1);
+        ASSERT_TRUE(memory.m_data);
+        ASSERT_EQ(memory.m_capacity, 16);
+        ASSERT_EQ(memory.m_size, 1);
+    }
+
+    {
+        auto memory = Memory<DummyAllocator>(1);
+        ASSERT_TRUE(memory.m_data);
+        ASSERT_EQ(memory.m_capacity, 16);
+        ASSERT_EQ(memory.m_size, 1);
+
+        memory.resize(0);
+        ASSERT_TRUE(memory.m_data);
+        ASSERT_EQ(memory.m_capacity, 16);
+        ASSERT_EQ(memory.m_size, 0);
+
+        memory.resize(1);
+        ASSERT_TRUE(memory.m_data);
+        ASSERT_EQ(memory.m_capacity, 16);
+        ASSERT_EQ(memory.m_size, 1);
+    }
+}
+
+TEST(MemoryResizeTest, SmallInitAndBigResizeOverflowWhenPadding)
+{
+    {
+        auto memory = Memory<DummyAllocator>(0);
+        ASSERT_EQ(memory.m_data, nullptr);
+        ASSERT_EQ(memory.m_capacity, 0);
+        ASSERT_EQ(memory.m_size, 0);
+
+        EXPECT_THROW_ERROR_CODE(memory.resize(std::numeric_limits<size_t>::max()), Exception, ErrorCodes::ARGUMENT_OUT_OF_BOUND);
+        ASSERT_EQ(memory.m_data, nullptr); // state is intact after exception
+        ASSERT_EQ(memory.m_size, 0);
+        ASSERT_EQ(memory.m_capacity, 0);
+
+        memory.resize(1);
+        ASSERT_TRUE(memory.m_data);
+        ASSERT_EQ(memory.m_capacity, 16);
+        ASSERT_EQ(memory.m_size, 1);
+
+        memory.resize(2);
+        ASSERT_TRUE(memory.m_data);
+        ASSERT_EQ(memory.m_capacity, 17);
+        ASSERT_EQ(memory.m_size, 2);
+
+        EXPECT_THROW_ERROR_CODE(memory.resize(std::numeric_limits<size_t>::max()), Exception, ErrorCodes::ARGUMENT_OUT_OF_BOUND);
+        ASSERT_TRUE(memory.m_data); // state is intact after exception
+        ASSERT_EQ(memory.m_capacity, 17);
+        ASSERT_EQ(memory.m_size, 2);
+
+        memory.resize(0x8000000000000000ULL-16);
+        ASSERT_TRUE(memory.m_data);
+        ASSERT_EQ(memory.m_capacity, 0x8000000000000000ULL - 1);
+        ASSERT_EQ(memory.m_size, 0x8000000000000000ULL - 16);
+
+#ifndef ABORT_ON_LOGICAL_ERROR
+        EXPECT_THROW_ERROR_CODE(memory.resize(0x8000000000000000ULL-15), Exception, ErrorCodes::LOGICAL_ERROR);
+        ASSERT_TRUE(memory.m_data);  // state is intact after exception
+        ASSERT_EQ(memory.m_capacity, 0x8000000000000000ULL - 1);
+        ASSERT_EQ(memory.m_size, 0x8000000000000000ULL - 16);
+#endif
+    }
+
+    {
+        auto memory = Memory<DummyAllocator>(1);
+        ASSERT_TRUE(memory.m_data);
+        ASSERT_EQ(memory.m_capacity, 16);
+        ASSERT_EQ(memory.m_size, 1);
+
+        EXPECT_THROW_ERROR_CODE(memory.resize(std::numeric_limits<size_t>::max()), Exception, ErrorCodes::ARGUMENT_OUT_OF_BOUND);
+        ASSERT_TRUE(memory.m_data); // state is intact after exception
+        ASSERT_EQ(memory.m_capacity, 16);
+        ASSERT_EQ(memory.m_size, 1);
+
+        memory.resize(1);
+        ASSERT_TRUE(memory.m_data);
+        ASSERT_EQ(memory.m_capacity, 16);
+        ASSERT_EQ(memory.m_size, 1);
+
+#ifndef ABORT_ON_LOGICAL_ERROR
+        EXPECT_THROW_ERROR_CODE(memory.resize(0x8000000000000000ULL-15), Exception, ErrorCodes::LOGICAL_ERROR);
+        ASSERT_TRUE(memory.m_data); // state is intact after exception
+        ASSERT_EQ(memory.m_capacity, 16);
+        ASSERT_EQ(memory.m_size, 1);
+#endif
+    }
+}
+
+
+TEST(MemoryResizeTest, BigInitAndSmallResizeOverflowWhenPadding)
+{
+    {
+        EXPECT_THROW_ERROR_CODE(
+        {
+            auto memory = Memory<DummyAllocator>(std::numeric_limits<size_t>::max());
+        }
+        , Exception
+        , ErrorCodes::ARGUMENT_OUT_OF_BOUND);
+    }
+
+    {
+        EXPECT_THROW_ERROR_CODE(
+        {
+            auto memory = Memory<DummyAllocator>(std::numeric_limits<size_t>::max() - 1);
+        }
+        , Exception
+        , ErrorCodes::ARGUMENT_OUT_OF_BOUND);
+    }
+
+    {
+        EXPECT_THROW_ERROR_CODE(
+        {
+            auto memory = Memory<DummyAllocator>(std::numeric_limits<size_t>::max() - 10);
+        }
+        , Exception
+        , ErrorCodes::ARGUMENT_OUT_OF_BOUND);
+    }
+
+#ifndef ABORT_ON_LOGICAL_ERROR
+    {
+        EXPECT_THROW_ERROR_CODE(
+        {
+            auto memory = Memory<DummyAllocator>(std::numeric_limits<size_t>::max() - 15);
+        }
+        , Exception
+        , ErrorCodes::LOGICAL_ERROR);
+    }
+
+    {
+        EXPECT_THROW_ERROR_CODE(
+        {
+            auto memory = Memory<DummyAllocator>(0x8000000000000000ULL - 15);
+        }
+        , Exception
+        , ErrorCodes::LOGICAL_ERROR);
+    }
+#endif
+
+    {
+       auto memory = Memory<DummyAllocator>(0x8000000000000000ULL - 16);
+       ASSERT_TRUE(memory.m_data);
+       ASSERT_EQ(memory.m_capacity, 0x8000000000000000ULL - 1);
+       ASSERT_EQ(memory.m_size, 0x8000000000000000ULL - 16);
+
+        memory.resize(1);
+        ASSERT_TRUE(memory.m_data);
+        ASSERT_EQ(memory.m_capacity, 0x8000000000000000ULL - 1);
+        ASSERT_EQ(memory.m_size, 1);
+    }
+}
+
+TEST(MemoryResizeTest, AlignmentInRealAllocator)
+{
+    {
+        auto memory = Memory<>(0, 3); // not the power of 2 but less than MALLOC_MIN_ALIGNMENT 8 so user-defined alignment is ignored at Allocator
+        ASSERT_EQ(memory.m_data, nullptr);
+        ASSERT_EQ(memory.m_capacity, 0);
+        ASSERT_EQ(memory.m_size, 0);
+
+        memory.resize(1);
+        ASSERT_TRUE(memory.m_data);
+        ASSERT_EQ(memory.m_capacity, 18);
+        ASSERT_EQ(memory.m_size, 1);
+
+        memory.resize(2);
+        ASSERT_TRUE(memory.m_data);
+        ASSERT_EQ(memory.m_capacity, 18);
+        ASSERT_EQ(memory.m_size, 2);
+
+        memory.resize(3);
+        ASSERT_TRUE(memory.m_data);
+        ASSERT_EQ(memory.m_capacity, 18);
+        ASSERT_EQ(memory.m_size, 3);
+
+        memory.resize(4);
+        ASSERT_TRUE(memory.m_data);
+        ASSERT_EQ(memory.m_capacity, 21);
+        ASSERT_EQ(memory.m_size, 4);
+
+        memory.resize(0);
+        ASSERT_TRUE(memory.m_data);
+        ASSERT_EQ(memory.m_capacity, 21);
+        ASSERT_EQ(memory.m_size, 0);
+
+        memory.resize(1);
+        ASSERT_TRUE(memory.m_data);
+        ASSERT_EQ(memory.m_capacity, 21);
+        ASSERT_EQ(memory.m_size, 1);
+    }
+
+    {
+        auto memory = Memory<>(0, 10); // not the power of 2
+        ASSERT_EQ(memory.m_data, nullptr);
+        ASSERT_EQ(memory.m_capacity, 0);
+        ASSERT_EQ(memory.m_size, 0);
+
+        EXPECT_THROW_ERROR_CODE(memory.resize(1), ErrnoException, ErrorCodes::CANNOT_ALLOCATE_MEMORY);
+        ASSERT_EQ(memory.m_data, nullptr); // state is intact after exception
+        ASSERT_EQ(memory.m_capacity, 0);
+        ASSERT_EQ(memory.m_size, 0);
+    }
+
+    {
+        auto memory = Memory<>(0, 32);
+        ASSERT_EQ(memory.m_data, nullptr);
+        ASSERT_EQ(memory.m_capacity, 0);
+        ASSERT_EQ(memory.m_size, 0);
+
+        memory.resize(1);
+        ASSERT_TRUE(memory.m_data);
+        ASSERT_EQ(memory.m_capacity, 47);
+        ASSERT_EQ(memory.m_size, 1);
+
+        memory.resize(32);
+        ASSERT_TRUE(memory.m_data);
+        ASSERT_EQ(memory.m_capacity, 47);
+        ASSERT_EQ(memory.m_size, 32);
+    }
+}
+
+TEST(MemoryResizeTest, SomeAlignmentOverflowWhenAlignment)
+{
+    {
+        auto memory = Memory<DummyAllocator>(0, 31);
+        ASSERT_EQ(memory.m_data, nullptr);
+        ASSERT_EQ(memory.m_capacity, 0);
+        ASSERT_EQ(memory.m_size, 0);
+
+        memory.resize(0);
+        ASSERT_EQ(memory.m_data, nullptr);
+        ASSERT_EQ(memory.m_capacity, 0);
+        ASSERT_EQ(memory.m_size, 0);
+
+        memory.resize(1);
+        ASSERT_TRUE(memory.m_data);
+        ASSERT_EQ(memory.m_capacity, 46);
+        ASSERT_EQ(memory.m_size, 1);
+
+        EXPECT_THROW_ERROR_CODE(memory.resize(std::numeric_limits<size_t>::max()), Exception, ErrorCodes::ARGUMENT_OUT_OF_BOUND);
+        ASSERT_TRUE(memory.m_data); // state is intact after exception
+        ASSERT_EQ(memory.m_capacity, 46);
+        ASSERT_EQ(memory.m_size, 1);
+    }
+
+}