Correct some migratability issues with lyieldable yield state

ARES.bt

@@ -34,6 +34,15 @@ const uchar UTAG_LLEVENTS = 29;
 const uchar UTAG_LLTIMERS = 30;
 const uchar UTAG_STRBUF = 31;
 const uchar UTAG_OPAQUE_BUFFER = 32;
+const uchar UTAG_YIELD_STATE = 33;
+
+/* One SlotManager's region of a UTAG_YIELD_STATE buffer. The innermost flag
+ * says whether this was the deepest active manager. */
+typedef struct {
+    uint16 length;           /* Slots plus the innermost flag */
+    uchar slots[length - 1];
+    uint8_t innermost;
+} YieldRegion;
 
 /* Mirrors enum AresType in VM/src/ares.cpp. These are wire values and are
  * assigned explicitly there - do not derive them from Luau's lua_Type. */
@@ -800,6 +809,22 @@ typedef struct {
         uchar data[used.v];
         break;
     }
+    case UTAG_YIELD_STATE:
+    {
+        UVarint length;
+        local int64 end = FTell() + length.v;
+        uint8_t abi_version; /* The frame's continuation must answer with at least this */
+        local int innermost = 0;
+        while (!innermost)
+        {
+            YieldRegion region;
+            innermost = region.innermost;
+        }
+        /* A buffer reused from an earlier yield can be larger than the chain */
+        if (FTell() < end)
+            uchar slack[end - FTell()];
+        break;
+    }
     case UTAG_OPAQUE_BUFFER:
     default:
     {

VM/include/llsl.h

@@ -32,6 +32,12 @@ enum class YieldableStatus : uint8_t {
 #define UTAG_LLTIMERS 30
 #define UTAG_STRBUF 31
 #define UTAG_OPAQUE_BUFFER 32
+// A yieldable C function's slot buffer, see lyieldable.h
+#define UTAG_YIELD_STATE 33
+
+// Passed as the status to a yieldable's continuation to get its ABI version
+// back instead of a resume
+#define LUA_YIELDABLE_ABI_QUERY (-1)
 
 // Internal global names for event/timer managers (hidden from user code)
 #define LLEVENTS_GLOBAL_NAME "/$ LLEvents"

VM/include/lua.h

@@ -746,7 +746,7 @@ static void populateperms(lua_State *L, bool forUnpersist)
     eris_persist_static(strlib, gmatch_aux)
 #endif
 #if defined(eris_c) || defined(lyieldstrlib_c)
-    eris_persist_cont(lyieldstrlib, yieldable_gmatch_aux_v0, yieldable_gmatch_aux_v0_k)
+    eris_persist_cont(lyieldstrlib, yieldable_gmatch_aux, yieldable_gmatch_aux_k)
 #endif
 #if defined(eris_c) || defined(lutf8lib_c)
     eris_persist_static(utf8lib, iter_aux)
@@ -762,12 +762,12 @@ static void populateperms(lua_State *L, bool forUnpersist)
     eris_persist_static_cont(corolib, auxwrapy, auxwrapcont)
 #endif
 #if defined(eris_c) || defined(lllevents_c)
-    eris_persist_static_cont(llevents, llevents_handle_event_v0, llevents_handle_event_v0_k)
+    eris_persist_static_cont(llevents, llevents_handle_event, llevents_handle_event_k)
     eris_persist_static_cont(llevents, llevents_once_wrapper, llevents_once_wrapper_cont)
     eris_persist_static(llevents, timer_wrapper_guard)
 #endif
 #if defined(eris_c) || defined(llltimers_c)
-    eris_persist_static_cont(llltimers, lltimers_tick_v0, lltimers_tick_v0_k)
+    eris_persist_static_cont(llltimers, lltimers_tick, lltimers_tick_k)
     eris_persist_cont(llltimers, timer_event_wrapper, timer_event_wrapper_cont)
 #endif
 #if defined(eris_c) || defined(llprim_c)

VM/include/lyieldable.h

@@ -2,7 +2,7 @@
 //
 // Slots are C++ primitives cached locally via PrimitiveSlot<T>. No stack
 // interaction on init. On yield, destructors pack cached values into a
-// an opaque tagged userdata (UTAG_OPAQUE_BUFFER) at position 1.
+// a tagged userdata (UTAG_YIELD_STATE) at position 1.
 //
 // This takes a lot of inspiration from UThreadInjector's instrumentation style,
 // but makes it explicit via macros.
@@ -15,6 +15,7 @@
 #pragma once
 
 #include "lua.h"
+#include "llsl.h"
 #include "Luau/Common.h"
 
 #include <cstddef>
@@ -27,6 +28,10 @@ static_assert(sizeof(bool) == 1, "DEFINE_SLOT(bool) requires sizeof(bool) == 1")
 namespace Luau
 {
 
+// Buffer layout: [u8 abi version][u16 region length][slots...][u8 innermost],
+// then each child region as [u16 region length][slots...][u8 innermost].
+inline constexpr size_t kYieldStateHeader = sizeof(uint8_t) + sizeof(uint16_t);
+
 // Forward declaration — slot destructors access SlotManager members.
 class SlotManager;
 
@@ -72,15 +77,22 @@ class SlotManager
 
     ptrdiff_t bufferStackOffset;
     bool initMode;
+    uint8_t abiVersion = 0;
     bool finalized = false;
 
     SlotManager* parent = nullptr;
     SlotManager* currentChild = nullptr;
-    // Byte 0 is reserved for the version tag; bytes 1-2 for the region length.
-    // Slot data starts at offset 3 for the root manager.
-    size_t baseOffset = sizeof(uint8_t) + sizeof(uint16_t);
+    // Where the base of this frame starts in the buffer
+    size_t baseOffset = kYieldStateHeader;
     uint16_t regionLengthOffset = sizeof(uint8_t);
     uint16_t innermostOffset = 0;
+    // Where this manager's region sits in the buffer being resumed from,
+    // which need not match where it will be written
+    size_t readBaseOffset = kYieldStateHeader;
+    uint16_t storedLength = 0;
+
+    // Resume only: reads and bounds-checks this manager's stored region length
+    void readStoredLength();
 
 public:
     // Set by flushForYield(). Slot destructors check this to decide
@@ -98,7 +110,7 @@ class SlotManager
     lua_Callbacks* callbacks = nullptr;
 
     // Root constructor: pushes nil at position 1 (init) or reads buffer (resume).
-    SlotManager(lua_State* L, bool is_init);
+    SlotManager(lua_State* L, bool is_init, uint8_t abi_version);
 
     // Child constructor: chains to parent, shares buffer.
     // Force-inlined: called per inner-loop iteration in match helpers.
@@ -114,6 +126,9 @@ class SlotManager
 
     bool isInit() const { return initMode; }
 
+    LUA_NORETURN void corrupt() const;
+    LUA_NORETURN void tooLarge() const;
+
     // Lock the slot layout. Must be called after all slots are allocated
     // and before YIELD_DISPATCH_BEGIN.
     // Force-inlined: called per inner-loop iteration in match helpers.
@@ -161,14 +176,14 @@ inline PrimitiveSlot<T> SlotManager::allocSlot(T* storage)
     static_assert(std::is_trivially_copyable_v<T>, "PrimitiveSlot<T> requires a trivially copyable type");
     LUAU_ASSERT(!finalized);
     auto off = static_cast<uint16_t>(baseOffset + requiredSize);
-    requiredSize += sizeof(T);
 
-    if (!initMode)
-    {
-        LUAU_ASSERT(off + sizeof(T) <= bufferSize);
-        memcpy(storage, bufferData + off, sizeof(T));
-    }
+    // If there's not enough data left in the buffer, we're best off just leaving the
+    // default there. This is intentionally a non-fatal error, intended to support appending
+    // optional slots without ABI version bumps.
+    if (!initMode && requiredSize + sizeof(T) < storedLength)
+        memcpy(storage, bufferData + readBaseOffset + requiredSize, sizeof(T));
 
+    requiredSize += sizeof(T);
     return PrimitiveSlot<T>{storage, off, this};
 }
 
@@ -179,13 +194,24 @@ inline PrimitiveSlot<T>::~PrimitiveSlot()
         memcpy(mgr->bufferData + offset, data, sizeof(T));
 }
 
+inline void SlotManager::readStoredLength()
+{
+    size_t length_at = readBaseOffset - sizeof(uint16_t);
+    if (readBaseOffset > bufferSize)
+        corrupt();
+    memcpy(&storedLength, bufferData + length_at, sizeof(uint16_t));
+    if (storedLength < 1 || readBaseOffset + storedLength > bufferSize)
+        corrupt();
+}
+
 // Child constructor. Chains to parent, sharing the same buffer region.
 // initMode is determined by parent.innermost: if parent is innermost
 // (no child data in buffer), this is a fresh call; otherwise, resume.
 LUAU_FORCEINLINE SlotManager::SlotManager(SlotManager& parent)
     : L(parent.L)
     , bufferStackOffset(parent.bufferStackOffset)
     , initMode(parent.innermost)
+    , abiVersion(parent.abiVersion)
     , parent(&parent)
     , baseOffset(parent.baseOffset + parent.requiredSize + sizeof(uint16_t))
     , regionLengthOffset(static_cast<uint16_t>(parent.baseOffset + parent.requiredSize))
@@ -200,6 +226,8 @@ LUAU_FORCEINLINE SlotManager::SlotManager(SlotManager& parent)
     {
         bufferData = parent.bufferData;
         bufferSize = parent.bufferSize;
+        readBaseOffset = parent.readBaseOffset + parent.storedLength + sizeof(uint16_t);
+        readStoredLength();
     }
 }
 
@@ -232,23 +260,25 @@ LUAU_FORCEINLINE void SlotManager::finalize()
     innermostOffset = static_cast<uint16_t>(baseOffset + requiredSize);
     requiredSize += sizeof(uint8_t);
     LUAU_ASSERT(requiredSize <= UINT16_MAX);
+    // Every offset in this region, including the ones allocSlot already
+    // truncated, sits below the region's end, so this is the one place the
+    // uint16 layout has to be checked. Nothing has been written through them
+    // yet: finalize precedes every yield point.
+    if (baseOffset + requiredSize > UINT16_MAX)
+        tooLarge();
     finalized = true;
 
     if (!initMode)
     {
-        // Validate that the stored region length matches what we expect.
-        uint16_t storedLength;
-        memcpy(&storedLength, bufferData + regionLengthOffset, sizeof(uint16_t));
-        LUAU_ASSERT(storedLength == requiredSize);
-
-        // Read innermost flag from buffer.
-        LUAU_ASSERT(innermostOffset + sizeof(uint8_t) <= bufferSize);
-        memcpy(&innermost, bufferData + innermostOffset, sizeof(uint8_t));
+        // The stored region was bounds-checked when its length was read
+        memcpy(&innermost, bufferData + readBaseOffset + storedLength - 1, sizeof(uint8_t));
+        // make extra sure we didn't muck up the read by validating `innermost`
+        if (innermost > 1)
+            corrupt();
 
-        // Only the innermost manager validates total size and clears buffer.
+        // Only the innermost manager clears the buffer.
         if (innermost)
         {
-            LUAU_ASSERT(baseOffset + requiredSize <= bufferSize);
             bufferData = nullptr;
             bufferSize = 0;
             // Leave buffer at position 1 for reuse on next yield.

VM/include/lyieldablemacros.h

@@ -24,19 +24,21 @@ using Luau::YieldGuard;
 // DEFINE_YIELDABLE_IMPL: base macro parameterized by linkage specifier.
 // Generates the init function (name), continuation (name_k), and body
 // (name_body) from a single function definition. The body receives
-// lua_State* L and bool is_init. Body is always static.
-#define DEFINE_YIELDABLE_IMPL(linkage, name, version)                \
-    static int name##_v##version##_body(lua_State* L, bool is_init);     \
-    linkage int name##_v##version(lua_State* L)                          \
-    {                                                       \
-        return name##_v##version##_body(L, true);                        \
-    }                                                       \
-    linkage int name##_v##version##_k(lua_State* L, int status)          \
-    {                                                       \
-        lua_checkstack(L, LUA_MINSTACK);                    \
-        return name##_v##version##_body(L, false);                       \
-    }                                                       \
-    static int name##_v##version##_body(lua_State* L, bool is_init)
+// lua_State* L, bool is_init and the ABI version. Body is always static.
+#define DEFINE_YIELDABLE_IMPL(linkage, name, version)                       \
+    static int name##_body(lua_State* L, bool is_init, uint8_t abi_version); \
+    linkage int name(lua_State* L)                                           \
+    {                                                                        \
+        return name##_body(L, true, (version));                              \
+    }                                                                        \
+    linkage int name##_k(lua_State* L, int status)                           \
+    {                                                                        \
+        if (status == LUA_YIELDABLE_ABI_QUERY)                               \
+            return (version);                                                \
+        lua_checkstack(L, LUA_MINSTACK);                                     \
+        return name##_body(L, false, (version));                             \
+    }                                                                        \
+    static int name##_body(lua_State* L, bool is_init, uint8_t abi_version)
 
 // DEFINE_YIELDABLE: static linkage (single translation unit).
 #define DEFINE_YIELDABLE(name, version) DEFINE_YIELDABLE_IMPL(static, name, version)
@@ -83,10 +85,10 @@ using Luau::YieldGuard;
 #define YIELD_DISPATCH(phase_name) \
     case Phase::phase_name: goto _yieldable_label_##phase_name
 
-// Closes the dispatch switch.
+// Closes the dispatch switch. A phase this build doesn't know is a bad buffer.
 #define YIELD_DISPATCH_END()                                                        \
     default:                                                                        \
-        LUAU_ASSERT(!"Unhandled yieldable phase");                                  \
+        _yieldable_slots.corrupt();                                                 \
     } (void)0
 
 // Fires the interrupt handler and yields if the VM requests it.

VM/src/ares.cpp

@@ -1774,6 +1774,7 @@ static void p_userdata(Info *info) {                               /* ... udata
     case UTAG_PROXY:
     case UTAG_QUATERNION:
     case UTAG_OPAQUE_BUFFER:
+    case UTAG_YIELD_STATE:
       WRITE_VALUE(size, ares_size_t);
       WRITE_RAW(value, size);
       break;
@@ -1842,6 +1843,7 @@ static void u_userdata(Info *info) {                                   /* ... */
     uint8_t utag = READ_VALUE(uint8_t);
     switch(utag) {
       case UTAG_OPAQUE_BUFFER:
+      case UTAG_YIELD_STATE:
       case UTAG_PROXY:
       {
           size_t size = READ_VALUE(ares_size_t);
@@ -3135,6 +3137,25 @@ u_thread(Info *info) {                                                 /* ... */
       }
       // We don't actually use the function for anything, just checking!
       lua_pop(info->L, 1);                                    /* ... thread */
+
+      // Okay, we may be resuming an lyieldable function. Make sure we can actually load the
+      // version provided by comparing it against what our impl actually supports
+      if (thread->status == LUA_OK || thread->status == LUA_YIELD || thread->status == LUA_BREAK) {
+        StkId base = thread->ci->base;
+        if (base < thread->top && ttisuserdata(base) && uvalue(base)->tag == UTAG_YIELD_STATE) {
+          if (!func_cl->c.cont) {
+            eris_error(info, "malformed data: yield state under a function without a continuation");
+          }
+          // Silly hack: Abuse the continuation function by defining a custom status, so we can ask
+          // what max version this func supports.
+          int version = func_cl->c.cont(thread, LUA_YIELDABLE_ABI_QUERY);
+          int found = (uint8_t)uvalue(base)->data[0];
+          if (version < 0 || version > UINT8_MAX || found > version) {
+            const char *debugname = cclosure_debugname(func_cl);
+            eris_error(info, "state saved by a newer version of %s", debugname ? debugname : "<unknown>");
+          }
+        }
+      }
     } else {
       if (ci_kind != ERIS_CI_KIND_NONE) {
         eris_error(info, "malformed data: invalid call info kind");