Harden the mesh decode path against hostile assets

indra/llcommon/alscratcharena.cpp

@@ -37,21 +37,50 @@ namespace
     // size, by which its return is checked to be the last one taken.
     constexpr size_t HEADER = ALScratchArena::ALIGNMENT;
 
+    // The thread's arena. Nothing is done to destroy it, so that it lasts
+    // as long as its thread, past every thread_local's destructor;
+    // ArenaRelease ends it.
     struct Arena
     {
         U8* mBase = nullptr;
         size_t mTop = 0;
-        // The arena could not be made: every block comes from the heap.
+        // The arena could not be made, or its thread has ended: every block
+        // comes from the heap.
         bool mUnavailable = false;
         ALScratchArena::Counters mCounters;
+    };
+
+    thread_local Arena tArena;
 
-        ~Arena()
+    // Frees an arena its thread has ended once no block is out of it.
+    void freeIfEnded(Arena& arena)
+    {
+        if (arena.mUnavailable && arena.mTop == 0)
         {
-            ll_aligned_free<64>(mBase);
+            ll_aligned_free<64>(arena.mBase);
+            arena.mBase = nullptr;
+        }
+    }
+
+    // Ends its thread's arena as the thread ends. It is made with the arena,
+    // so a thread_local made before is destroyed after it: what that takes
+    // then comes from the heap, and a block it holds goes back into the
+    // arena, which is freed once the last is back.
+    struct ArenaRelease
+    {
+        Arena* mArena = nullptr;
+
+        ~ArenaRelease()
+        {
+            if (mArena)
+            {
+                mArena->mUnavailable = true;
+                freeIfEnded(*mArena);
+            }
         }
     };
 
-    thread_local Arena tArena;
+    thread_local ArenaRelease tArenaRelease;
 } // namespace
 
 //static
@@ -62,12 +91,13 @@ void* ALScratchArena::allocate(size_t bytes)
     {
         arena.mBase = static_cast<U8*>(ll_aligned_malloc<64>(CAPACITY));
         arena.mUnavailable = !arena.mBase;
+        tArenaRelease.mArena = &arena;
     }
 
     // Compared before it is rounded, so that a size near the top of size_t,
     // as meshoptimizer asks for when a count overflows, goes to the heap.
     const size_t room = CAPACITY - arena.mTop;
-    if (arena.mBase && room > HEADER && bytes <= room - HEADER)
+    if (!arena.mUnavailable && room > HEADER && bytes <= room - HEADER)
     {
         const size_t size = (bytes + ALIGNMENT - 1) & ~(ALIGNMENT - 1);
         if (size <= room - HEADER)
@@ -93,8 +123,20 @@ void ALScratchArena::deallocate(void* block)
     U8* p = static_cast<U8*>(block);
     if (arena.mBase && p >= arena.mBase && p < arena.mBase + CAPACITY)
     {
-        llassert(p + *reinterpret_cast<const size_t*>(p - HEADER) == arena.mBase + arena.mTop);
-        arena.mTop = static_cast<size_t>(p - HEADER - arena.mBase);
+        // Only the block on top moves the top down. Moving it down to any
+        // other would hand the blocks above, still in use, to the next
+        // taker, so the top stays, and the block's room is lost to the
+        // thread.
+        const size_t start = static_cast<size_t>(p - arena.mBase);
+        if (start >= HEADER && start <= arena.mTop && *reinterpret_cast<const size_t*>(p - HEADER) == arena.mTop - start)
+        {
+            arena.mTop = start - HEADER;
+            freeIfEnded(arena);
+        }
+        else
+        {
+            LL_WARNS_ONCE("ALScratchArena") << "A block was given back out of order or twice" << LL_ENDL;
+        }
         return;
     }
     ::operator delete(block);

indra/llcommon/alscratcharena.h

@@ -41,7 +41,17 @@
 //
 // The order is the caller's to keep: a block goes back before every block
 // taken after it has. Nothing holding a block may yield to a coroutine
-// that takes from the same thread's arena.
+// that takes from the same thread's arena. A block given back out of that
+// order leaves the top where it is, so that no block still in use is
+// handed out again: its room is lost to the thread, and the first such
+// return is logged. A block given back twice is passed over the same way,
+// unless it is the block on top.
+//
+// As its thread ends the arena takes no more blocks: one taken after, as
+// by a thread_local destroyed after it, comes from the heap. The arena is
+// freed once every block taken from it is back, so that a block a
+// thread_local holds stays its holder's until given back. An arena that
+// has lost room to a block given back out of order is never freed.
 class LL_COMMON_API ALScratchArena
 {
 public:

indra/llcommon/tests/alscratcharena_test.cpp

@@ -33,6 +33,7 @@
 #include <cstdint>
 #include <cstring>
 #include <new>
+#include <optional>
 #include <thread>
 
 namespace
@@ -155,6 +156,109 @@ namespace
         ALScratchArena::deallocate(after);
     }
 
+    // A block given back while one taken after it is still out leaves the
+    // top where it is: the next block is taken above the one still out, not
+    // over it. The thread is the case's own. Given back again once it is on
+    // top, the first block brings the top back to the start, so that the
+    // arena is freed as the thread ends.
+    TEST_CASE("alscratcharena: a block given back out of order leaves the blocks above it to their holders", "[alscratcharena]")
+    {
+        void* below = nullptr;
+        void* above = nullptr;
+        void* next = nullptr;
+        void* again = nullptr;
+        std::thread([&]
+                    {
+                        below = ALScratchArena::allocate(64);
+                        above = ALScratchArena::allocate(64);
+                        ALScratchArena::deallocate(below);
+                        next = ALScratchArena::allocate(64);
+                        ALScratchArena::deallocate(next);
+                        ALScratchArena::deallocate(above);
+                        ALScratchArena::deallocate(below);
+                        again = ALScratchArena::allocate(64);
+                        ALScratchArena::deallocate(again);
+                    })
+            .join();
+        CHECK(address(above) > address(below));
+        CHECK(address(next) >= address(above) + 64);
+        CHECK(again == below);
+    }
+
+    // A thread_local made before its thread first takes a block, holding
+    // one as the thread ends, is destroyed after the arena is ended. The
+    // block is still its own, to write and to give back into the arena;
+    // what it takes after comes from the heap.
+    struct LateHolder
+    {
+        void* block = nullptr;
+        // What the arena served the destructor, once it has finished.
+        std::optional<ALScratchArena::Counters>* served = nullptr;
+
+        ~LateHolder()
+        {
+            if (served)
+            {
+                std::memset(block, 0xab, 64);
+                ALScratchArena::deallocate(block);
+                ALScratchArena::resetCounters();
+                ALScratchArena::deallocate(ALScratchArena::allocate(64));
+                *served = ALScratchArena::counters();
+            }
+        }
+    };
+    thread_local LateHolder tLateHolder;
+
+    TEST_CASE("alscratcharena: a block still out as its thread ends goes back into the arena", "[alscratcharena]")
+    {
+        std::optional<ALScratchArena::Counters> late;
+        std::thread([&]
+                    {
+                        tLateHolder.served = &late;
+                        tLateHolder.block = ALScratchArena::allocate(64);
+                    })
+            .join();
+        REQUIRE(late.has_value());
+        CHECK(late->arena_blocks == 0);
+        CHECK(late->heap_blocks == 1);
+    }
+
+    // A thread_local made before its thread first takes a block is
+    // destroyed after the thread's arena is freed, and what it takes then
+    // comes from the heap. A block taken from the freed arena would be
+    // written into memory already given back, and the destructor would not
+    // get as far as saying what it was served.
+    struct LateTaker
+    {
+        // What the arena served the destructor, once it has finished.
+        std::optional<ALScratchArena::Counters>* served = nullptr;
+
+        ~LateTaker()
+        {
+            if (served)
+            {
+                ALScratchArena::resetCounters();
+                ALScratchArena::deallocate(ALScratchArena::allocate(64));
+                *served = ALScratchArena::counters();
+            }
+        }
+    };
+    thread_local LateTaker tLateTaker;
+
+    TEST_CASE("alscratcharena: a block taken after the thread's arena is freed comes from the heap", "[alscratcharena]")
+    {
+        std::optional<ALScratchArena::Counters> late;
+        std::thread([&]
+                    {
+                        tLateTaker.served = &late;
+                        ALScratchArena::deallocate(ALScratchArena::allocate(64));
+                    })
+            .join();
+        REQUIRE(late.has_value());
+        CHECK(late->arena_blocks == 0);
+        CHECK(late->heap_blocks == 1);
+    }
+
     TEST_CASE("alscratcharena: each thread takes from an arena of its own", "[alscratcharena]")
     {
         void* mine = ALScratchArena::allocate(64);