Fix returning pointer for zero-sized allocation outside of page

The problematic scenario looks like this:

- Allocate new page on arena

- Return pointer for the last piece of the page - so `offset == size`

- Try to allocate zero-sized segment - since `offset + 0 <= size` we return pointer to `page + offset`

- `GetStartOfPage(page + offset)` returns the next page - and it leads to malicious behavior for the next (probably unmapped) page

Now we don't allow to physically allocate zero-sized region and use aligned stub object.
commit_hash:5fa77d6bd78c7f712f35da943fcfe9023f78ec5e

Author: ilezhankin

yql/essentials/minikql/mkql_alloc.cpp

@@ -7,9 +7,9 @@
 
 #include <tuple>
 
-namespace NKikimr {
+namespace NKikimr::NMiniKQL {
 
-namespace NMiniKQL {
+static ui8 ZeroSizeObject alignas(ArrowAlignment)[0];
 
 constexpr ui64 ArrowSizeForArena = (TAllocState::POOL_PAGE_SIZE >> 2);
 
@@ -268,7 +268,14 @@ void TPagedArena::Clear() noexcept {
     }
 }
 
+namespace {
+
 void* MKQLArrowAllocateOnArena(ui64 size) {
+    Y_ENSURE(size);
+    // If size is zero we can get in trouble: when `page->Offset == page->Size`.
+    // The zero size leads to return `ptr` just after the current page.
+    // Then getting start of page for such pointer returns next page - which may be unmapped or unrelevant to `ptr`.
+
     TAllocState* state = TlsAllocState;
     Y_ENSURE(state);
 
@@ -287,8 +294,8 @@ void* MKQLArrowAllocateOnArena(ui64 size) {
 
         page = (TMkqlArrowHeader*)GetAlignedPage();
         NYql::NUdf::SanitizerMakeRegionAccessible(page, sizeof(TMkqlArrowHeader));
-        page->Offset = sizeof(TMkqlArrowHeader);
-        page->Size = pageSize;
+        page->Offset = 0;
+        page->Size = pageSize - sizeof(TMkqlArrowHeader); // for consistency with CleanupArrowList()
         page->UseCount = 1;
 
         if (state->EnableArrowTracking) {
@@ -299,14 +306,20 @@ void* MKQLArrowAllocateOnArena(ui64 size) {
         }
     }
 
-    void* ptr = (ui8*)page + page->Offset;
+    void* ptr = (ui8*)page + page->Offset + sizeof(TMkqlArrowHeader);
     page->Offset += alignedSize;
     ++page->UseCount;
+
+    Y_DEBUG_ABORT_UNLESS(TAllocState::GetPageStart(ptr) == page);
+
     return ptr;
 }
 
-namespace {
 void* MKQLArrowAllocateImpl(ui64 size) {
+    if (Y_UNLIKELY(size == 0)) {
+        return reinterpret_cast<void*>(ZeroSizeObject);
+    }
+
     if (Y_LIKELY(!TAllocState::IsDefaultAllocatorUsed())) {
         if (size <= ArrowSizeForArena) {
             return MKQLArrowAllocateOnArena(size);
@@ -345,28 +358,14 @@ void* MKQLArrowAllocateImpl(ui64 size) {
     header->Size = size;
     return header + 1;
 }
-} // namespace
-
-void* MKQLArrowAllocate(ui64 size) {
-    auto sizeWithRedzones = NYql::NUdf::GetSizeToAlloc(size);
-    void* mem = MKQLArrowAllocateImpl(sizeWithRedzones);
-    return NYql::NUdf::WrapPointerWithRedZones(mem, sizeWithRedzones);
-}
-
-void* MKQLArrowReallocate(const void* mem, ui64 prevSize, ui64 size) {
-    auto res = MKQLArrowAllocate(size);
-    memcpy(res, mem, Min(prevSize, size));
-    MKQLArrowFree(mem, prevSize);
-    return res;
-}
 
 void MKQLArrowFreeOnArena(const void* ptr) {
     auto* page = (TMkqlArrowHeader*)TAllocState::GetPageStart(ptr);
     if (page->UseCount.fetch_sub(1) == 1) {
         if (!page->Entry.IsUnlinked()) {
             TAllocState* state = TlsAllocState;
             Y_ENSURE(state);
-            state->OffloadFree(page->Size);
+            state->OffloadFree(page->Size + sizeof(TMkqlArrowHeader));
             page->Entry.Unlink();
 
             auto it = state->ArrowBuffers.find(page);
@@ -380,8 +379,12 @@ void MKQLArrowFreeOnArena(const void* ptr) {
     return;
 }
 
-namespace {
 void MKQLArrowFreeImpl(const void* mem, ui64 size) {
+    if (Y_UNLIKELY(mem == reinterpret_cast<const void*>(ZeroSizeObject))) {
+        Y_DEBUG_ABORT_UNLESS(size == 0);
+        return;
+    }
+
     if (Y_LIKELY(!TAllocState::IsDefaultAllocatorUsed())) {
         if (size <= ArrowSizeForArena) {
             return MKQLArrowFreeOnArena(mem);
@@ -409,15 +412,34 @@ void MKQLArrowFreeImpl(const void* mem, ui64 size) {
 
     ReleaseAlignedPage(header, fullSize);
 }
+
 } // namespace
 
+void* MKQLArrowAllocate(ui64 size) {
+    auto sizeWithRedzones = NYql::NUdf::GetSizeToAlloc(size);
+    void* mem = MKQLArrowAllocateImpl(sizeWithRedzones);
+    return NYql::NUdf::WrapPointerWithRedZones(mem, sizeWithRedzones);
+}
+
+void* MKQLArrowReallocate(const void* mem, ui64 prevSize, ui64 size) {
+    auto res = MKQLArrowAllocate(size);
+    memcpy(res, mem, Min(prevSize, size));
+    MKQLArrowFree(mem, prevSize);
+    return res;
+}
+
 void MKQLArrowFree(const void* mem, ui64 size) {
     mem = NYql::NUdf::UnwrapPointerWithRedZones(mem, size);
     auto sizeWithRedzones = NYql::NUdf::GetSizeToAlloc(size);
     return MKQLArrowFreeImpl(mem, sizeWithRedzones);
 }
 
 void MKQLArrowUntrack(const void* mem, ui64 size) {
+    if (Y_UNLIKELY(mem == reinterpret_cast<const void*>(ZeroSizeObject))) {
+        Y_DEBUG_ABORT_UNLESS(size == 0);
+        return;
+    }
+
     mem = NYql::NUdf::GetOriginalAllocatedObject(mem, size);
     TAllocState* state = TlsAllocState;
     Y_ENSURE(state);
@@ -437,7 +459,7 @@ void MKQLArrowUntrack(const void* mem, ui64 size) {
             if (!page->Entry.IsUnlinked()) {
                 page->Entry.Unlink(); // unlink page immediately so we don't accidentally free untracked memory within `TAllocState`
                 state->ArrowBuffers.erase(it);
-                state->OffloadFree(page->Size);
+                state->OffloadFree(page->Size + sizeof(TMkqlArrowHeader));
             }
 
             return;
@@ -459,6 +481,4 @@ void MKQLArrowUntrack(const void* mem, ui64 size) {
     }
 }
 
-} // NMiniKQL
-
-} // NKikimr
+} // namespace NKikimr::NMiniKQL

yql/essentials/minikql/mkql_alloc_ut.cpp

@@ -2,8 +2,7 @@
 
 #include <library/cpp/testing/unittest/registar.h>
 
-namespace NKikimr {
-namespace NMiniKQL {
+namespace NKikimr::NMiniKQL {
 
 Y_UNIT_TEST_SUITE(TMiniKQLAllocTest) {
     Y_UNIT_TEST(TestPagedArena) {
@@ -73,7 +72,7 @@ Y_UNIT_TEST_SUITE(TMiniKQLAllocTest) {
         }
         TWithDefaultMiniKQLAlloc::FreeWithSize(p2, 10);
     }
-    Y_UNIT_TEST(InitiallyAcqured) {
+    Y_UNIT_TEST(InitiallyAcquired) {
         {
             TScopedAlloc alloc(__LOCATION__);
             UNIT_ASSERT_VALUES_EQUAL(true, alloc.IsAttached());
@@ -93,7 +92,48 @@ Y_UNIT_TEST_SUITE(TMiniKQLAllocTest) {
             UNIT_ASSERT_VALUES_EQUAL(false, alloc.IsAttached());
         }
     }
-}
+    Y_UNIT_TEST(ArrowAllocateZeroSize) {
+        // Choose small enough pieces to hit arena (using some internal knowledge)
+        const auto pieceSize = AlignUp<size_t>(1, ArrowAlignment);
+        UNIT_ASSERT_EQUAL(0, (TAllocState::POOL_PAGE_SIZE - sizeof(TMkqlArrowHeader)) % pieceSize);
+        const auto pieceCount = (TAllocState::POOL_PAGE_SIZE - sizeof(TMkqlArrowHeader)) / pieceSize;
+        void** ptrs = new void*[pieceCount];
 
+        // Populate the current page on arena to maximum offset
+        TScopedAlloc alloc(__LOCATION__);
+        for (auto i = 0ul; i < pieceCount; ++i) {
+            ptrs[i] = MKQLArrowAllocate(pieceSize);
+        }
+
+        // Check all pieces are on the same page
+        void* pageStart = TAllocState::GetPageStart(ptrs[0]);
+        for (auto i = 1ul; i < pieceCount; ++i) {
+            UNIT_ASSERT_VALUES_EQUAL(pageStart, TAllocState::GetPageStart(ptrs[i]));
+        }
+
+        // Allocate zero-sized piece twice and check it's the same address
+        void* ptrZero1 = MKQLArrowAllocate(0);
+        void* ptrZero2 = MKQLArrowAllocate(0);
+        UNIT_ASSERT_VALUES_EQUAL(ptrZero1, ptrZero2);
+
+        // Allocate one more small piece and check that it's on another page - different from zero-sized piece
+        void* ptrOne = MKQLArrowAllocate(1);
+        UNIT_ASSERT_VALUES_UNEQUAL(pageStart, TAllocState::GetPageStart(ptrOne));
+        UNIT_ASSERT_VALUES_UNEQUAL(TAllocState::GetPageStart(ptrZero1), TAllocState::GetPageStart(ptrOne));
+
+        // Untrack zero-sized piece
+        MKQLArrowUntrack(ptrZero1, 0);
+
+        // Deallocate all the stuff
+        for (auto i = 0ul; i < pieceCount; ++i) {
+            MKQLArrowFree(ptrs[i], pieceSize);
+        }
+        MKQLArrowFree(ptrZero1, 0);
+        MKQLArrowFree(ptrZero2, 0);
+        MKQLArrowFree(ptrOne, 1);
+
+        delete[] ptrs;
+    }
 }
-}
+
+} // namespace NKikimr::NMiniKQL