WIP rollback function hook

src/game_api/script/events.cpp

@@ -471,3 +471,13 @@ void post_event(ON event)
             return true;
         });
 }
+
+void heap_clone_event(ON event, StateMemory* from, StateMemory* to)
+{
+    LuaBackend::for_each_backend(
+        [&](LuaBackend::LockedBackend backend)
+        {
+            backend->on_clone_heap(event, from, to);
+            return true;
+        });
+}

src/game_api/script/events.hpp

@@ -30,6 +30,7 @@ void post_init_layer(LAYER layer);
 void post_unload_layer(LAYER layer);
 void post_room_generation();
 void post_level_generation();
+void heap_clone_event(ON event, StateMemory* from, StateMemory* to);
 
 void on_death_message(STRINGID stringid);
 std::optional<bool> pre_get_feat(FEAT feat);

src/game_api/script/lua_backend.cpp

@@ -1814,3 +1814,24 @@ void LuaBackend::on_post(ON event)
         }
     }
 }
+
+void LuaBackend::on_clone_heap(ON event, StateMemory* from, StateMemory* to)
+{
+    if (!get_enabled())
+        return;
+
+    auto now = get_frame_count();
+    for (auto& [id, callback] : callbacks)
+    {
+        if (is_callback_cleared(id))
+            continue;
+
+        if (callback.screen == event)
+        {
+            set_current_callback(-1, id, CallbackType::Normal);
+            handle_function<void>(this, callback.func, from, to);
+            clear_current_callback();
+            callback.lastRan = now;
+        }
+    }
+}

src/game_api/script/lua_backend.hpp

@@ -135,6 +135,8 @@ enum class ON
     BLOCKED_UPDATE,
     BLOCKED_GAME_LOOP,
     BLOCKED_PROCESS_INPUT,
+    PRE_CLONE_HEAP,
+    POST_CLONE_HEAP,
 };
 
 struct IntOption
@@ -440,6 +442,7 @@ class LuaBackend
     void load_user_data();
     bool on_pre(ON event);
     void on_post(ON event);
+    void on_clone_heap(ON event, StateMemory* from, StateMemory* to);
 };
 
 template <class Inheriting>

src/game_api/script/lua_vm.cpp

@@ -2435,7 +2435,11 @@ end
         "BLOCKED_GAME_LOOP",
         ON::BLOCKED_GAME_LOOP,
         "BLOCKED_PROCESS_INPUT",
-        ON::BLOCKED_PROCESS_INPUT);
+        ON::BLOCKED_PROCESS_INPUT,
+        "PRE_CLONE_HEAP",
+        ON::PRE_CLONE_HEAP,
+        "POST_CLONE_HEAP",
+        ON::POST_CLONE_HEAP);
 
     /* ON
     // LOGO

src/game_api/search.cpp

@@ -1826,6 +1826,15 @@ std::unordered_map<std::string_view, AddressRule> g_address_rules{
             .at_exe()
             .function_start(),
     },
+    {
+        // Put write bp on state.win_state and enter a multiplayer game
+        "heap_clone"sv,
+        PatternCommandBuffer{}
+            .find_inst("4c 8d 05 f4 ca 27 00"_gh)
+            .find_next_inst("eb 27"_gh)
+            .offset(-0xC)
+            .at_exe()
+    },
     {
         // ^ writes to state.pause on state.loading == 3
         "unpause_level"sv,

src/game_api/state.cpp

@@ -302,6 +302,7 @@ State& State::get()
             init_state_update_hook();
             init_process_input_hook();
             init_game_loop_hook();
+            init_state_clone_hook();
 
             auto bucket = Bucket::get();
             if (!bucket->patches_applied)
@@ -694,6 +695,76 @@ void init_state_update_hook()
     }
 }
 
+ExecutableMemory g_heap_clone_redirect;
+using OnHeapClone = void(uint64_t heap_to, uint64_t heap_container_from);
+OnHeapClone* g_heap_clone_trampoline{nullptr};
+void HeapClone(uint64_t heap_to, uint64_t heap_container_from)
+{
+    uint64_t location = memory_read<uint64_t>(State::get().location);
+    StateMemory* state_from = reinterpret_cast<StateMemory*>(memory_read<uint64_t>(heap_container_from + 0x88) + location);
+    StateMemory* state_to = reinterpret_cast<StateMemory*>(heap_to + location);
+    heap_clone_event(ON::PRE_CLONE_HEAP, state_from, state_to);
+    // g_heap_clone_trampoline(heap_container_to, heap_container_from, heap_to);
+    // heap_clone_event(ON::POST_CLONE_HEAP, state_from,state_to);
+}
+
+// Original function params: clone_heap(ThreadStorageContainer to, ThreadStorageContainer from)
+// HeapContainer has heap1 and heap2 variables, and some sort of timer, that just increases constantly, I guess to handle the rollback and multi-threaded stuff
+// The rest of what HeapContainer has is unknown for now
+// After writing to a chosen storage from the content of `from->heap1`, sets `to->heap2` to the newly copied thread storage
+void init_state_clone_hook()
+{
+    auto heap_clone = get_address("heap_clone");
+    // g_heap_clone_trampoline = (OnHeapClone*)(heap_clone+0x65);
+    // Hook the function after it has chosen a thread storage to write to, and pass it to the hook
+    size_t heap_clone_redirect_from_addr = heap_clone+0x65;
+    DEBUG("HEAP_CLONE: {}\n", static_cast<uint64_t>(heap_clone_redirect_from_addr));
+    const std::string redirect_code = fmt::format(
+        "\x51"                         // PUSH       RCX
+        "\x52"                         // PUSH       RDX
+        "\x41\x50"                     // PUSH       R8
+        "\x41\x51"                     // PUSH       R9
+        "\x48\x83\xEC\x28"             // SUB        RSP, 28 // Shadow space + Stack alignment
+        "\x4C\x89\xC9"                 // MOV        RCX, R9 == heap_to
+        "\x48\xb8{}"                   // MOV        RAX, &HeapClone
+        "\xff\xd0"                     // CALL       RAX
+        "\x48\x83\xC4\x28"             // ADD        RSP, 28
+        "\x41\x59"                     // POP        R9
+        "\x41\x58"                     // POP        R8
+        "\x5A"                         // POP        RDX
+        "\x59"                         // POP        RCX
+                                       // Original Code Begin
+        "\x48\x8b\x82\x88\x00\x00\x00" // MOV        RAX, qword ptr [RDX + 0x88]
+        "\x4d\x89\xca"                 // MOV        R10, R9
+        "\x49\x29\xc2"                 // SUB        R10, RAX
+                                       // Original Code End
+        "\x48\xbe{}"                   // MOV        RSI, jump_back_addr
+        "\xff\xe6"sv,                  // JMP        RSI
+        to_le_bytes(&HeapClone),
+        to_le_bytes(heap_clone+0x72));
+
+    g_heap_clone_redirect = ExecutableMemory{redirect_code};
+
+    std::string code = fmt::format(
+        "\x48\xb8{}"  // MOV         RAX, g_heap_clone_redirect.get()
+        "\xff\xe0"    // JMP         RAX
+        "\x90"sv,     // NOP
+        to_le_bytes((size_t)g_heap_clone_redirect.get()));
+
+    write_mem_prot(heap_clone_redirect_from_addr, code, true);
+    // DetourTransactionBegin();
+    // DetourUpdateThread(GetCurrentThread());
+    // DetourAttach((void**)&g_heap_clone_trampoline, &HeapClone);
+
+    // const LONG error = DetourTransactionCommit();
+    // if (error != NO_ERROR)
+    // {
+    //     DEBUG("Failed hooking heap_clone stuff: {}\n", error);
+    // } else {
+    //     write_mem_prot(heap_clone+0x149, "\xFF\x74\x24\x20", true);
+    // }
+}
+
 using OnProcessInput = void(void*);
 OnProcessInput* g_process_input_trampoline{nullptr};
 void ProcessInput(void* s)

src/game_api/state.hpp

@@ -387,17 +387,18 @@ struct State
     SaveData* savedata();
     LiquidPhysicsEngine* get_correct_liquid_engine(ENT_TYPE liquid_type);
 
+    size_t location;
   private:
     State(size_t addr)
         : location(addr){};
 
-    size_t location;
     State(const State&) = delete;
     State& operator=(const State&) = delete;
 };
 void init_state_update_hook();
 void init_process_input_hook();
 void init_game_loop_hook();
+void init_state_clone_hook();
 
 uint8_t enum_to_layer(const LAYER layer, std::pair<float, float>& player_position);
 uint8_t enum_to_layer(const LAYER layer);