AuroraRuntime/Source/Processes/AuProcess.Win32.cpp

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

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

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

    void InitWin32()
    {
        JOBOBJECT_EXTENDED_LIMIT_INFORMATION jeli {};

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

        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;
    }

    struct Hack
    {
        HANDLE base;
        AuUInt32 count;
    };
    
    static BOOL CALLBACK TermWinHandleWin32Thread(HWND handle, LPARAM a)
    {
        DWORD windowpid;
        auto context = reinterpret_cast<Hack *>(a);
        auto pid = GetProcessId(context->base);
        
        if (!pGetWindowThreadProcessId || 
            !pSendMessageA)
        {
            return true;
        }

        pGetWindowThreadProcessId(handle, &windowpid);
        if (pid && pid == windowpid)
        {
            pSendMessageA(handle, WM_CLOSE, 0, 0);
            context->count++;
        }
        return true;
    }
    
    static bool SendExtermianteWinloopCloseMessages(HANDLE handle)
    {
        THREADENTRY32 te{};
        HANDLE h{};
        Hack hello;

        hello.count = 0;
        hello.base = handle;

        te.dwSize = sizeof(te);
        te.th32OwnerProcessID = GetProcessId(handle);

        if (!pEnumThreadWindows)
        {
            return false;
        }

        if (!te.th32OwnerProcessID)
        {
            return false;
        }

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

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

        CloseHandle(h);

        if (!hello.count)
        {
            return false;
        }
        
        return true;
    }
    
    static bool SendControlCEquiv(HANDLE handle)
    {
        auto cwnd = GetConsoleWindow();

        if (!(AttachConsole(GetProcessId(handle))))
        {
            return false;
        }

        if (cwnd == GetConsoleWindow())
        {
            FreeConsole();
            return false;
        }

        if (SetConsoleCtrlHandler({}, true))
        {
            GenerateConsoleCtrlEvent(CTRL_C_EVENT, 0);
            SysAssert(FreeConsole(), "Couldn't pop off zombie process console");

            // other people use 2k ms, their code is over a decade or so old tho
            // this is evil; however, since we are calling free, the os should be 
            // smart enough to not terminate us now. id be nice if msft made that
            // defined behaviour. for now, we have to sleep to make sure we dont
            // eat items in the kernels/conhosts/somethings pending work queue
            Sleep(750);
        }
        else
        {
            SysAssert(FreeConsole(), "Couldn't pop off zombie process console");
            return false;
        }
        
        if (!SetConsoleCtrlHandler({}, false))
        {
            SysPushErrorHAL("Couldn't swap control+c handler -> wont accept control termination on exit");
        }
        
        return true;
    }
    
    static bool Wait2500OrUntilClose(HANDLE handle)
    {
        return Threading::YieldPollNs(true, AuTime::SteadyClockNS() + AuMSToNS<AuUInt64>(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);
    }
}