AuroraRuntime/Source/Processes/Process.Win32.cpp

/***
    Copyright (C) 2021-2022 J Reece Wilson (a/k/a "Reece"). All rights reserved.

    File: Process.Win32.cpp
    Date: 2022-1-29
    Author: Reece
***/
#include <Source/RuntimeInternal.hpp>
#include "Process.Win32.hpp"
#include <shellapi.h>
#include <tlhelp32.h>
#include <process.h>

namespace Aurora::Processes
{
    static HANDLE gLeaderJob = INVALID_HANDLE_VALUE;

    void InitWin32()
    {
        gLeaderJob = CreateJobObject(NULL, NULL);
        if (!gLeaderJob)
        {
            SysPushErrorArg("CreateJobObject error");
            return;
        }

        JOBOBJECT_EXTENDED_LIMIT_INFORMATION jeli {};
        jeli.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
        if (!SetInformationJobObject(gLeaderJob, JobObjectExtendedLimitInformation, &jeli, sizeof(jeli)))
        {
            SysPushErrorGen("SetInformationJobObject error");
        }
    }

    void DeinitWin32()
    {
        AuWin32CloseHandle(gLeaderJob);
    }

    void AssignJobWorker(HANDLE handle)
    {
        if (gLeaderJob)
        {
            if (!AssignProcessToJobObject(gLeaderJob, handle))
            {
                SysPushErrorGen("Could not AssignProcessToObject");
            }
        }
    }
    
    static bool HasWin32ProcessExited(HANDLE handle)
    {
        DWORD exitCode;
        if (!GetExitCodeProcess(handle, &exitCode))
        {
            return true;
        }
        
       return exitCode != STILL_ACTIVE;
    }
    
    static BOOL TermWinHandleWin32Thread(HWND handle, LPARAM a)
    {
        SendMessageA(handle, WM_CLOSE, 0, 0);
        (*reinterpret_cast<AuUInt32 *>(a))++;
        return true;
    }
    
    static bool SendExtermianteWinloopCloseMessages(HANDLE handle)
    {
        THREADENTRY32 te{};
        HANDLE h{};

        h = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, GetProcessId(handle));
        if (h == INVALID_HANDLE_VALUE)
        {
            return false;
        }

        te.dwSize = sizeof(te);
        AuUInt32 count{};
        if (Thread32First(h, &te))
        {
            do
            {
                if (te.dwSize >= FIELD_OFFSET(THREADENTRY32, th32OwnerProcessID) +
                    sizeof(te.th32OwnerProcessID))
                {
                    EnumThreadWindows(te.th32ThreadID, TermWinHandleWin32Thread, reinterpret_cast<LPARAM>(&count));
                }
                te.dwSize = sizeof(te);
            } while (Thread32Next(h, &te));
        }

        CloseHandle(h);

        if (!count)
        {
            return false;
        }
        
        return true;
    }
    
    static bool SendControlCEquiv(HANDLE handle)
    {
        if (!(AttachConsole(GetProcessId(handle))))
        {
            return false;
        }

        if (SetConsoleCtrlHandler({}, true))
        {
            GenerateConsoleCtrlEvent(CTRL_C_EVENT, 0);
            SysAssert(FreeConsole(), "Couldn't pop off zombie process console");
            Sleep(750);
        }
        else
        {
            SysAssert(FreeConsole(), "Couldn't pop off zombie process console");
            return false;
        }
        
        if (!SetConsoleCtrlHandler({}, true))
        {
            SysPushErrorHAL("Couldn't swap control+c handler -> wont accept control termination on exit");
        }
        
        return true;
    }
    
    static bool Wait2500OrUntilClose(HANDLE handle)
    {
        return Threading::YieldPoll(true, 2500, [=]()
        {
            return !HasWin32ProcessExited(handle);
        });
    }
    
    bool SendExitSignal(HANDLE handle)
    {
        if (handle == INVALID_HANDLE_VALUE)
        {
            return false;
        }

        if (!(SendControlCEquiv(handle) // send GenerateConsoleCtrlEvent
              || SendExtermianteWinloopCloseMessages(handle))) // send WM_CLOSE
        {
            return false;
        }

        return Wait2500OrUntilClose(handle);
    }
}