2 files changed, 76 insertions(+), 19 deletions(-)

M include/gamecontroller.h
M src/gamecontroller.cpp

@@ 37,11 37,24 @@ private: // types
         std::string data;
     };
 
+    struct BSSSegment
+    {
+        std::string module_name;
+        void       *ptr;
+        size_t      size;
+    };
+
+    struct BSSData : public BSSSegment
+    {
+        BSSData(const BSSSegment &cp) : BSSSegment(cp) {}
+        std::string data;
+    };
+
     struct SavedState
     {
         size_t                   frame_number;
         std::vector<AVSHeapData> heap_data;
-        std::string              bss_data;
+        std::vector<BSSData>     bss_data;
     };
 
 private: // functions


@@ 57,6 70,7 @@ private: // functions
 
     SharedState* _shared_state() const;
     const LDR_DATA_TABLE_ENTRY* _libafp_module();
+    const LDR_DATA_TABLE_ENTRY* _libavs_module();
 
     DWORD _patched_game_tick(void *p);
     void _patched_acio_hbhi_get_csb(void *p);


@@ 70,11 84,12 @@ private: // functions
     void _restore_state();
 
     friend DWORD __fastcall __patched_game_tick(void *p);
-    friend void __cdecl __patched_acio_hbhi_get_csb(void *p);
-    friend DWORD __cdecl __patched_avs_heap_boot(void *ptr, DWORD size, DWORD align);
+    friend void __cdecl     __patched_acio_hbhi_get_csb(void *p);
+    friend DWORD __cdecl    __patched_avs_heap_boot(void *ptr, DWORD size, DWORD align);
 
 private: // variables
     Console                 _console;
+
     GameTickFunc            _orig_game_tick_func_ptr;
     ACIOHBHIGetCSBFunc      _orig_acio_hbhi_get_csb_func_ptr;
     AVSHeapBootFunc         _orig_avs_heap_boot_func_ptr;


@@ 82,10 97,10 @@ private: // variables
     SharedMemory            _shared_memory;
     SavedState              _saved_state;
     std::vector<AVSHeap>    _avs_heaps;
+    std::vector<BSSSegment> _bss_segments;
+    bool                    _bss_segments_initialized;
     bool                    _afp_logging_enabled;
 
     LDR_DATA_TABLE_ENTRY   *_cached_afp_module;
-
-    void                   *_afp_bss_ptr;
-    size_t                  _afp_bss_size;
+    LDR_DATA_TABLE_ENTRY   *_cached_avs_module;
 };

@@ 40,9 40,8 @@ GameController& GameController::instance()
 GameController::GameController()
     : _shared_memory(sizeof(SharedState), SharedMemory::Mode::OPEN, SHMEM_MAPPING_NAME),
       _afp_logging_enabled(false),
-      _cached_afp_module(nullptr),
-      _afp_bss_ptr(nullptr),
-      _afp_bss_size(0)
+      _bss_segments_initialized(false),
+      _cached_afp_module(nullptr)
 {}
 
 GameController::~GameController()


@@ 167,6 166,20 @@ const LDR_DATA_TABLE_ENTRY* GameController::_libafp_module()
     return _cached_afp_module;
 }
 
+const LDR_DATA_TABLE_ENTRY* GameController::_libavs_module()
+{
+    if (!_cached_avs_module) {
+        const ProcessModuleTable modtable;
+        _cached_avs_module = modtable.entry_named(AVS_MODULE_NAME);
+
+        if (!_cached_avs_module) {
+            throw std::runtime_error("Failed to load libavs module.");
+        }
+    }
+
+    return _cached_avs_module;
+}
+
 DWORD GameController::_patched_game_tick(void *p)
 {
     DWORD result = 0;


@@ 357,6 370,7 @@ bool GameController::_patch_avs(const ProcessModuleTable &modtable)
 void GameController::_save_state()
 {
     _saved_state.heap_data.clear();
+    _saved_state.bss_data.clear();
 
     for (const AVSHeap &heap : _avs_heaps) {
         const char *base_addr = reinterpret_cast<const char *>(heap.ptr);


@@ 367,17 381,37 @@ void GameController::_save_state()
         heap_data.data.assign(base_addr, size);
     }
 
-    if (!_afp_bss_ptr) {
-        const LDR_DATA_TABLE_ENTRY *afp_module = _libafp_module();
-        const void *afp_known_bss_addr = ((uint8_t *)afp_module->DllBase + 0x8334c);
+    if (!_bss_segments_initialized) {
+        // load libafp bss segment
+        BSSSegment &afp_bss_segment = _bss_segments.emplace_back();
+        const void *afp_known_bss_addr = ((uint8_t *)_libafp_module()->DllBase + 0x8334c);
         const MEMORY_BASIC_INFORMATION afp_bss_region = mem_region(afp_known_bss_addr);
-
-        _afp_bss_ptr = afp_bss_region.BaseAddress;
-        _afp_bss_size = afp_bss_region.RegionSize;
+        afp_bss_segment.module_name = AFP_MODULE_NAME;
+        afp_bss_segment.ptr = afp_bss_region.BaseAddress;
+        afp_bss_segment.size = afp_bss_region.RegionSize;
+
+        // load libavs bss segment
+        BSSSegment &avs_bss_segment = _bss_segments.emplace_back();
+        const void *avs_known_bss_addr = ((uint8_t *)_libavs_module()->DllBase + 0x597d0);
+        const MEMORY_BASIC_INFORMATION avs_bss_region = mem_region(avs_known_bss_addr);
+        avs_bss_segment.module_name = AVS_MODULE_NAME;
+        avs_bss_segment.ptr = avs_bss_region.BaseAddress;
+        avs_bss_segment.size = avs_bss_region.RegionSize;
+
+        _bss_segments_initialized = true;
     }
 
-    DEBUG_LOG("Saving libafp bss at 0x%p (0x%zx bytes)\n", _afp_bss_ptr, _afp_bss_size);
-    _saved_state.bss_data.assign(reinterpret_cast<const char *>(_afp_bss_ptr), _afp_bss_size);
+    for (const BSSSegment &bss_segment : _bss_segments) {
+        const char *base_ptr = reinterpret_cast<const char *>(bss_segment.ptr);
+        const size_t size = bss_segment.size;
+        DEBUG_LOG("Saving %s bss at 0x%p (0x%zx bytes)\n",
+                  bss_segment.module_name.c_str(),
+                  base_ptr,
+                  size);
+
+        BSSData &bss_data = _saved_state.bss_data.emplace_back(bss_segment);
+        bss_data.data.assign(base_ptr, size);
+    }
 
     _saved_state.frame_number = _shared_state()->frame_count;
 }


@@ 392,8 426,16 @@ void GameController::_restore_state()
         heap_data.data.copy(base_addr, size);
     }
 
-    DEBUG_LOG("Restoring libafp bss to 0x%p (0x%zx bytes)\n", _afp_bss_ptr, _afp_bss_size);
-    _saved_state.bss_data.copy(reinterpret_cast<char *>(_afp_bss_ptr), _afp_bss_size);
+    for (const BSSData &bss_data : _saved_state.bss_data) {
+        char *base_addr = reinterpret_cast<char *>(bss_data.ptr);
+        const size_t size = bss_data.size;
+
+        DEBUG_LOG("Restoring %s bss to 0x%p (0x%zx bytes)\n",
+                  bss_data.module_name.c_str(),
+                  base_addr,
+                  size);
+        bss_data.data.copy(base_addr, size);
+    }
 
     _shared_state()->frame_count = _saved_state.frame_number;
 }