3 files changed, 56 insertions, 42 deletions

@@ -17,11 +17,10 @@ public:
 };
 
 #define EXE_NAME L"witness64_d3d11.exe"
-#define PROC_ATTACH 0x401
+#define HEARTBEAT 0x401
 #define RANDOMIZE_READY 0x402
 #define RANDOMIZING 0403
 #define RANDOMIZE_DONE 0x404
-#define CHECK_NEWGAME 0x405
 
 // Globals
 HWND g_hwnd;
@@ -40,24 +39,25 @@ LRESULT CALLBACK WndProc(HWND hwnd, UINT message, WPARAM wParam, LPARAM lParam)
                 PostQuitMessage(0);
     } else if (message == WM_COMMAND || message == WM_TIMER) {
         switch (LOWORD(wParam)) {
-            case PROC_ATTACH:
-                if (!g_witnessProc->Initialize(EXE_NAME)) {
-                    SetTimer(g_hwnd, PROC_ATTACH, 1000, NULL); // Retry in 1s
-                } else {
-                    g_randomizer = std::make_shared<Randomizer>(g_witnessProc);
-                    PostMessage(g_hwnd, WM_COMMAND, RANDOMIZE_READY, NULL);
-                    SetTimer(g_hwnd, CHECK_NEWGAME, 1000, NULL); // Start checking for new game
-                }
-                break;
-            case CHECK_NEWGAME:
-                if (g_witnessProc) {
-                    // It shouldn't be possible to pause earlier than this, so subsequent checks will fail.
-                    if (g_witnessProc->GetCurrentFrame() < 80) { // New game
+            case HEARTBEAT:
+                SetTimer(g_hwnd, HEARTBEAT, 1000, NULL);
+                // Potential improvement: Change this call to be part of the HEARTBEAT message.
+                switch (g_witnessProc->Heartbeat(EXE_NAME)) {
+                    case ProcStatus::NotRunning:
+                        // Shut down randomizer, wait for startup
+                        if (g_randomizer) g_randomizer = nullptr;
+                        break;
+                    case ProcStatus::Running:
+                        if (!g_randomizer) {
+                            g_randomizer = std::make_shared<Randomizer>(g_witnessProc);
+                            PostMessage(g_hwnd, WM_COMMAND, RANDOMIZE_READY, NULL);
+                        }
+                        break;
+                    case ProcStatus::NewGame: // This status will fire only once per new game
                         SetWindowText(g_seed, L"");
                         PostMessage(g_hwnd, WM_COMMAND, RANDOMIZE_READY, NULL);
+                        break;
                     }
-                    SetTimer(g_hwnd, CHECK_NEWGAME, 1000, NULL); // Continue checking for new game
-                }
                 break;
             case RANDOMIZE_READY:
                 SetWindowText(g_randomizerStatus, L"Randomize");
@@ -65,6 +65,7 @@ LRESULT CALLBACK WndProc(HWND hwnd, UINT message, WPARAM wParam, LPARAM lParam)
                 break;
             case RANDOMIZING:
                 if (!g_randomizer) break;
+                EnableWindow(g_randomizerStatus, FALSE);
                 SetWindowText(g_randomizerStatus, L"Randomizing...");
                 std::thread([]{
                     g_randomizer->Randomize();
@@ -74,8 +75,6 @@ LRESULT CALLBACK WndProc(HWND hwnd, UINT message, WPARAM wParam, LPARAM lParam)
             case RANDOMIZE_DONE:
                 SetWindowText(g_randomizerStatus, L"Randomized!");
                 break;
-            default:
-                break;
         }
     }
     return DefWindowProc(hwnd, message, wParam, lParam);
@@ -127,9 +126,9 @@ int APIENTRY wWinMain(_In_ HINSTANCE hInstance, _In_opt_ HINSTANCE hPrevInstance
 
     CreateLabel(390, 15, 90, L"Version: " VERSION_STR);
     g_seed = CreateText(10, 10, 100);
-    g_randomizerStatus = CreateButton(120, 10, 100, L"Randomize", RANDOMIZING);
+    g_randomizerStatus = CreateButton(120, 10, 110, L"Randomize", RANDOMIZING);
     EnableWindow(g_randomizerStatus, FALSE);
-    PostMessage(g_hwnd, WM_COMMAND, PROC_ATTACH, NULL);
+    PostMessage(g_hwnd, WM_COMMAND, HEARTBEAT, NULL);
 
         ShowWindow(g_hwnd, nCmdShow);
         UpdateWindow(g_hwnd);
@@ -7,9 +7,6 @@
 #undef PROCESSENTRY32
 #undef Process32Next
 
-Memory::Memory() {
-}
-
 [[nodiscard]]
 bool Memory::Initialize(const std::wstring& processName) {
         // First, get the handle of the process
@@ -48,23 +45,40 @@ bool Memory::Initialize(const std::wstring& processName) {
     return true;
 }
 
-Memory::~Memory() {
-    if (_handle != nullptr) {
-            CloseHandle(_handle);
+ProcStatus Memory::Heartbeat(const std::wstring& processName) {
+    if (!_handle && !Initialize(processName)) {
+        // Couldn't initialize, definitely not running
+        return ProcStatus::NotRunning;
+    }
+
+    DWORD exitCode = 0;
+    GetExitCodeProcess(_handle, &exitCode);
+    if (exitCode != STILL_ACTIVE) {
+        // Process has exited, clean up.
+        _computedAddresses.clear();
+        _handle = NULL;
+        return ProcStatus::NotRunning;
     }
-}
 
-int Memory::GetCurrentFrame() {
-        int SCRIPT_FRAMES;
+    int currentFrame = 0x7FFFFFFF;
         if (GLOBALS == 0x5B28C0) {
-                SCRIPT_FRAMES = 0x5BE3B0;
+        currentFrame = ReadData<int>({0x5BE3B0}, 1)[0];
         } else if (GLOBALS == 0x62D0A0) {
-                SCRIPT_FRAMES = 0x63954C;
+        currentFrame = ReadData<int>({0x63954C}, 1)[0];
         } else {
         assert(false);
-        return 0x7FFFFFFF;
         }
-        return ReadData<int>({SCRIPT_FRAMES}, 1)[0];
+    if (currentFrame < 80) return ProcStatus::NewGame;
+
+    // TODO: Some way to return ProcStatus::Randomized vs ProcStatus::NotRandomized vs ProcStatus::DeRandomized;
+
+    return ProcStatus::Running;
+}
+
+Memory::~Memory() {
+    if (_handle != nullptr) {
+            CloseHandle(_handle);
+    }
 }
 
 void Memory::AddSigScan(const std::vector<byte>& scanBytes, const std::function<void(int index)>& scanFunc)
@@ -7,21 +7,25 @@
 // #define GLOBALS 0x5B28C0
 #define GLOBALS 0x62D0A0
 
+enum class ProcStatus {
+    NotRunning,
+    Running,
+    NewGame
+};
+
 // https://github.com/erayarslan/WriteProcessMemory-Example
 // http://stackoverflow.com/q/32798185
 // http://stackoverflow.com/q/36018838
 // http://stackoverflow.com/q/1387064
 class Memory {
 public:
-        Memory();
-        bool Initialize(const std::wstring& processName);
+        Memory() = default;
+    ProcStatus Heartbeat(const std::wstring& processName);
         ~Memory();
 
         Memory(const Memory& memory) = delete;
         Memory& operator=(const Memory& other) = delete;
 
-        int GetCurrentFrame();
-
         template <class T>
         std::vector<T> ReadArray(int panel, int offset, int size) {
                 return ReadData<T>({GLOBALS, 0x18, panel*8, offset, 0}, size);
@@ -48,9 +52,6 @@ public:
 private:
         template<class T>
         std::vector<T> ReadData(const std::vector<int>& offsets, size_t numItems) {
-        if (GetExitCodeProcess(_process) != STILL_ACTIVE) {
-            // Signal error, somehow
-        }
                 std::vector<T> data;
                 data.resize(numItems);
                 for (int i=0; i<5; i++) {
@@ -73,8 +74,8 @@ private:
                 ThrowError();
         }
 
+        bool Initialize(const std::wstring& processName);
         void ThrowError();
-
         void* ComputeOffset(std::vector<int> offsets);
 
         std::map<uintptr_t, uintptr_t> _computedAddresses;