AuroraRuntime/Source/Processes/AuProcess.NT.cpp

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

    File: AuProcess.NT.cpp
    Date: 2021-6-12
    Author: Reece
***/
#include <Source/RuntimeInternal.hpp>
#include "AuProcess.NT.hpp"
#include "AuProcesses.hpp"
#include <process.h>

#if defined(AURORA_PLATFORM_WIN32)
    #include "AuProcess.Win32.hpp"
#endif

#include "AuArgvQuote.hpp"

#include <Source/IO/FS/FS.hpp>
#include <Source/IO/FS/Async.NT.hpp>
#include <Source/IO/Loop/LSHandle.hpp>
#include <Source/IO/AuIOHandle.hpp>

#include "AuCreatePipeEx.NT.hpp"

namespace Aurora::Processes
{
    struct ProcessAliveLoopSource : AuLoop::LSHandle
    {
        ProcessAliveLoopSource(HANDLE h);
        virtual AuLoop::ELoopSource GetType() override;
    };

    ProcessAliveLoopSource::ProcessAliveLoopSource(HANDLE h) : LSHandle(AuReinterpretCast<AuUInt>(h))
    {
    }

    AuLoop::ELoopSource ProcessAliveLoopSource::GetType()
    {
        return AuLoop::ELoopSource::eProcessDead;
    }

    ProcessImpl::ProcessImpl(StartupParameters &&params) : startup_(AuMove(params))
    {
        AuIOFS::NormalizePath(this->startup_.process, this->startup_.process);
        if (this->startup_.workingDirectory)
        {
            AuString a;
            AuIOFS::NormalizePath(a, this->startup_.workingDirectory.value());
            this->startup_.workingDirectory = a;
        }

        this->startup_.args.insert(startup_.args.begin(), startup_.process);
               
        // ehhhh https://github.com/tritao/WindowsSDK/blob/07983c7ba4f6861d15e23f195744c60c0c249ce0/SDKs/SourceDir/Windows%20Kits/10/Source/10.0.17763.0/ucrt/exec/cenvarg.cpp#L23
        for (const auto &arg : this->startup_.args)
        {
            this->cargs_.push_back(arg.c_str());
            AuString escaped;
            ArgvQuote(arg, escaped, false);
            this->windowsCli_ += escaped + " ";
        }

        this->cargs_.push_back(nullptr);
        if (windowsCli_.size())
        {
            this->windowsCli_.resize(this->windowsCli_.size() - 1);
        }
        
        this->type_ = params.type;
    }

    ProcessImpl::~ProcessImpl()
    {
        if (this->type_ == ESpawnType::eSpawnChildProcessWorker)
        {
            TryKill();
            Terminate();
        }

        SetEvent(this->poke_);

        if (this->thread_)
        {
            this->thread_.reset();
        }

        AuWin32CloseHandle(this->poke_);
        AuWin32CloseHandle(this->process_);
        AuWin32CloseHandle(this->hthread_);

        ShutdownPipes();
    }

    AuUInt ProcessImpl::GetProcessId()
    {
        if (this->process_ == INVALID_HANDLE_VALUE)
        {
            return {};
        }

        return ::GetProcessId(this->process_);
    }

    bool ProcessImpl::TermWinEnumProcesses()
    {
    #if defined(AURORA_PLATFORM_WIN32)
        return SendExitSignal(this->process_);
    #else
        return false;
    #endif
    }

    bool ProcessImpl::TryKill()
    {
        return HasExited() || TermWinEnumProcesses();
    }

    bool ProcessImpl::HasExited()
    {
        DWORD exitCode;

        if (!GetExitCodeProcess(this->process_, &exitCode))
        {
            return true;
        }

        return exitCode != STILL_ACTIVE;
    }

    bool ProcessImpl::Terminate()
    {
        if (this->process_ != INVALID_HANDLE_VALUE)
        {
            return false;
        }

        return TerminateProcess(this->process_, 0) || HasExited();
    }

    AuSPtr<Threading::IWaitable> ProcessImpl::AsWaitable()
    {
        if (!this->thread_) return nullptr;
        return this->thread_->AsWaitable();
    }

    AuSPtr<AuLoop::ILoopSource> ProcessImpl::AsLoopSource()
    {
        return this->loopSource_;
    }

    AuSInt ProcessImpl::GetExitCode()
    {
        return this->exitCode_;
    }

    void ProcessImpl::ShutdownPipes()
    {
        RelOtherHandles();
        //AuWin32CloseHandle(this->pipeStdOutRead_);
        //AuWin32CloseHandle(this->pipeStdErrRead_);
        //AuWin32CloseHandle(this->pipeStdInWrite_);
        this->pipeStdInWrite_ = this->pipeStdOutRead_ = this->pipeStdErrRead_ = INVALID_HANDLE_VALUE;
        this->fsHandle_.reset();
        this->fsStream_.reset();
        this->fsErrorHandle_.reset();
        this->fsErrorStream_.reset();
    }

    bool ProcessImpl::Read(EStandardHandle stream, const AuMemoryViewStreamWrite &destination, bool nonblock)
    {
        DWORD size = destination.length;
        
        if (!EStandardHandleIsValid(stream) || (stream == EStandardHandle::eStdIn))
        {
            SysPushErrorArg("Invalid Stream");
            return {};
        }

        auto handle = stream == EStandardHandle::eStdError ? this->pipeStdErrRead_ : this->pipeStdOutRead_;
        if (handle == INVALID_HANDLE_VALUE)
        {
            return false;
        }

        if (nonblock || !destination.ptr)
        {
            DWORD avail {};
            if (!PeekNamedPipe(handle, NULL, NULL, NULL, &avail, NULL))
            {
                return false;
            }

            if (!avail)
            {
                return true;
            }

            size = AuMin(size, avail);            
        }

        if (!destination.ptr)
        {
            destination.outVariable = size;
            return true;
        }

        auto ret = ReadFile(handle, destination.ptr, size, &size, NULL);
        destination.outVariable = size;
        return ret || nonblock;
    }

    bool ProcessImpl::Write(const AuMemoryViewStreamRead &in)
    {
        if (!in.ptr)
        {
            SysPushErrorArg("Missing pointer to stdin stream section");
            return {};
        }

        DWORD size = in.length;
        if (pipeStdInWrite_ == INVALID_HANDLE_VALUE)
        {
            SysPushErrorUninitialized();
            return false;
        }

        return WriteFile(this->pipeStdInWrite_, in.ptr, size, &size, NULL) && size == in.length;
    }

    bool ProcessImpl::Init()
    {
        SECURITY_ATTRIBUTES saAttr {};

        saAttr.nLength = sizeof(SECURITY_ATTRIBUTES);
        saAttr.bInheritHandle = TRUE;

        if (this->type_ == ESpawnType::eSpawnOvermap)
        {
            return true;
        }

        if (this->process_ != INVALID_HANDLE_VALUE)
        {
            return false;
        }
        
        this->poke_ = CreateEventW(NULL, TRUE, FALSE, NULL);
        if (!this->poke_)
        {
            return false;
        }

        this->exitCode_ = 0x10110100;
        if (this->startup_.fwdOut == EStreamForward::eAsyncPipe)
        {
            if (!CreatePipeEx(&this->pipeStdOutRead_,
                              &this->pipeStdOutWrite_,
                              &saAttr,
                              0,
                              FILE_FLAG_OVERLAPPED,
                              0))
            {
                return false;
            }

            if (!SetHandleInformation(pipeStdOutRead_, HANDLE_FLAG_INHERIT, 0))
            {
                return false;
            }
        }
        else if (this->startup_.fwdOut == EStreamForward::eIOHandle)
        {
            auto optHandle = this->startup_.handleOutStream->GetOSHandleSafe();
            if (!optHandle)
            {
                return {};
            }

            HANDLE hTargetProcess = ::GetCurrentProcess();
            HANDLE hHandle { (HANDLE) optHandle.value() };

            if (!::DuplicateHandle(hTargetProcess,
                                   hHandle,
                                   hTargetProcess,
                                   &this->pipeStdOutWrite_,
                                   GENERIC_READ | GENERIC_WRITE,
                                   TRUE,
                                   0))
            {
                return false;
            }
        }
        else if (this->startup_.fwdOut == EStreamForward::eCurrentProcess)
        {
            HANDLE handle = GetStdHandle(STD_OUTPUT_HANDLE);

            if (handle != INVALID_HANDLE_VALUE && handle)
            {
                SetHandleInformation(handle, HANDLE_FLAG_INHERIT, HANDLE_FLAG_INHERIT);

                this->bDontRelOut_ = true;
                this->pipeStdOutWrite_ = handle;
            }
        }

        if (this->startup_.fwdErr == EStreamForward::eAsyncPipe)
        {
            if (!CreatePipeEx(&this->pipeStdErrRead_,
                              &this->pipeStdErrWrite_,
                              &saAttr,
                              0,
                              FILE_FLAG_OVERLAPPED,
                              0))
            {
                return false;
            }

            if (!SetHandleInformation(this->pipeStdErrRead_, HANDLE_FLAG_INHERIT, 0))
            {
                return false;
            }
        }
        else if (this->startup_.fwdErr == EStreamForward::eIOHandle)
        {
            auto optHandle = this->startup_.handleErrorStream->GetOSHandleSafe();
            if (!optHandle)
            {
                return {};
            }

            HANDLE hTargetProcess = ::GetCurrentProcess();
            HANDLE hHandle { (HANDLE) optHandle.value() };

            if (!::DuplicateHandle(hTargetProcess,
                                   hHandle,
                                   hTargetProcess,
                                   &this->pipeStdErrWrite_,
                                   GENERIC_READ | GENERIC_WRITE,
                                   TRUE,
                                   0))
            {
                return false;
            }
        }
        else if (this->startup_.fwdErr == EStreamForward::eCurrentProcess)
        {
            HANDLE handle = GetStdHandle(STD_ERROR_HANDLE);

            if (handle != INVALID_HANDLE_VALUE && handle)
            {
                SetHandleInformation(handle, HANDLE_FLAG_INHERIT, HANDLE_FLAG_INHERIT);

                this->bDontRelErr_ = true;
                this->pipeStdErrWrite_ = handle;
            }
        }

        if (this->startup_.fwdIn == EStreamForward::eAsyncPipe)
        {
            if (!CreatePipeEx(&this->pipeStdInRead_,
                              &this->pipeStdInWrite_,
                              &saAttr,
                              0,
                              0,
                              FILE_FLAG_OVERLAPPED))
            {
                return false;
            }

            if (!SetHandleInformation(this->pipeStdInWrite_, HANDLE_FLAG_INHERIT, 0))
            {
                return false;
            }
        }
        else if (this->startup_.fwdIn == EStreamForward::eIOHandle)
        {
            auto optHandle = this->startup_.handleInputStream->GetOSHandleSafe();
            if (!optHandle)
            {
                return {};
            }

            HANDLE hTargetProcess = ::GetCurrentProcess();
            HANDLE hHandle { (HANDLE) optHandle.value() };

            if (!::DuplicateHandle(hTargetProcess,
                                   hHandle,
                                   hTargetProcess,
                                   &this->pipeStdInRead_,
                                   GENERIC_READ,
                                   TRUE,
                                   0))
            {
                return false;
            }
        }
        else if (this->startup_.fwdIn == EStreamForward::eCurrentProcess)
        {
            HANDLE handle = GetStdHandle(STD_ERROR_HANDLE);

            if (handle != INVALID_HANDLE_VALUE && handle)
            {
                SetHandleInformation(handle, HANDLE_FLAG_INHERIT, HANDLE_FLAG_INHERIT);

                this->bDontRelIn_ = true;
                this->pipeStdInRead_ = handle;
            }
        }

        {
            HANDLE nulFile;

        #define NEW_NULL_HANDLE                                                                                                                \
            {                                                                                                                                  \
                nulFile = CreateFileA("NUL", GENERIC_READ | GENERIC_WRITE, FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, 0, NULL);  \
                SetHandleInformation(nulFile, HANDLE_FLAG_INHERIT, HANDLE_FLAG_INHERIT);                                                       \
            }

            if (this->pipeStdInRead_ == INVALID_HANDLE_VALUE)
            {
                NEW_NULL_HANDLE;
                this->pipeStdInRead_ = nulFile;
            }

            if (this->pipeStdErrWrite_ == INVALID_HANDLE_VALUE)
            {
                NEW_NULL_HANDLE;
                this->pipeStdErrWrite_ = nulFile;
            }

            if (this->pipeStdOutWrite_ == INVALID_HANDLE_VALUE)
            {
                NEW_NULL_HANDLE;
                this->pipeStdOutWrite_ = nulFile;
            }
        }

        if (this->startup_.fwdIn == EStreamForward::eAsyncPipe || this->startup_.fwdOut == EStreamForward::eAsyncPipe)
        {
            this->fsHandle_ = AuIO::IOHandleShared();
            if (!this->fsHandle_)
            {
                return false;
            }

            this->fsStream_ = AuMakeShared<IO::FS::NtAsyncFileStream>();
            if (!this->fsStream_)
            {
                return false;
            }

            this->fsHandle_->InitFromPairMove((AuUInt64)this->pipeStdOutRead_, (AuUInt64)this->pipeStdInWrite_);
            AuStaticCast<AuIO::AFileHandle>(this->fsHandle_)->bIsAsync = true;
            this->fsStream_->Init(this->fsHandle_);
        }

        if (this->startup_.fwdErr == EStreamForward::eAsyncPipe)
        {
            this->fsErrorHandle_ = AuIO::IOHandleShared();
            if (!this->fsErrorHandle_)
            {
                return false;
            }

            this->fsErrorStream_ = AuMakeShared<IO::FS::NtAsyncFileStream>();
            if (!this->fsErrorStream_)
            {
                return false;
            }

            this->fsErrorHandle_->InitFromPairMove((AuUInt64)this->pipeStdErrRead_, {});
            AuStaticCast<AuIO::AFileHandle>(this->fsErrorHandle_)->bIsAsync = true;
            this->fsErrorStream_->Init(this->fsErrorHandle_);
        }

        return true;
    }

    AuSPtr<AuIO::IAsyncTransaction> ProcessImpl::NewAsyncTransaction()
    {
        return this->fsStream_ ? this->fsStream_->NewTransaction() : AuSPtr<AuIO::IAsyncTransaction> {};
    }

    AuSPtr<AuIO::IAsyncTransaction> ProcessImpl::NewErrorStreamAsyncTransaction()
    {
        return this->fsErrorStream_ ? this->fsErrorStream_->NewTransaction() : AuSPtr<AuIO::IAsyncTransaction> {};
    }

    bool ProcessImpl::Start()
    {
        if (this->process_ != INVALID_HANDLE_VALUE)
        {
            return false;
        }

        if (this->type_ == ESpawnType::eSpawnOvermap)
        {
            _spawnv(_P_OVERLAY, this->startup_.process.c_str(), this->cargs_.data());
            SysPushErrorGen("_spawnv didn't overwrite the process map, given {} ({})", this->startup_.process, this->windowsCli_);
            return false;
        }

        PROCESS_INFORMATION processInfo = { 0 };
        {
            STARTUPINFOW startupInfo = { 0 };
            startupInfo.cb = sizeof(startupInfo);
            
            startupInfo.hStdInput  = this->pipeStdInRead_;
            startupInfo.hStdError  = this->pipeStdErrWrite_;
            startupInfo.hStdOutput = this->pipeStdOutWrite_;
            startupInfo.dwFlags |= STARTF_USESTDHANDLES;

            auto cwd = this->startup_.workingDirectory;
            std::wstring wcwd;
            if (cwd)
            {
                wcwd = Locale::ConvertFromUTF8(cwd.value());
                if (!wcwd.size())
                {
                    SysPushErrorMem();
                    return false;
                }
            }

            AuUInt32 uCreateFlags {};

            if (this->startup_.type == ESpawnType::eSpawnThreadLeader)
            {
                uCreateFlags |= DETACHED_PROCESS | CREATE_NEW_PROCESS_GROUP;
            }
            else
            {
                if (((this->startup_.fwdIn == EStreamForward::eNewConsoleWindow) ||
                    (this->startup_.fwdErr == EStreamForward::eNewConsoleWindow) ||
                    (this->startup_.fwdOut == EStreamForward::eNewConsoleWindow)) &&
                    (!this->startup_.bNoShowConsole))
                {
                    uCreateFlags |= CREATE_NEW_CONSOLE;
                }
            }

            if (this->startup_.bInDebugMode ||
                this->startup_.optAffinity)
            {
                uCreateFlags |= CREATE_SUSPENDED;
            }

            if (this->startup_.bNoShowConsole)
            {
                 // yea we can keep CREATE_NO_WINDOW on for non-console apps. its legal -> https://docs.microsoft.com/en-us/windows/win32/procthread/process-creation-flags
                uCreateFlags |= CREATE_NO_WINDOW;
            }

            AuList<AuPair<AuString, AuString>> envVars;
            if (this->startup_.bInheritEnvironmentVariables)
            {
                if (this->startup_.environmentVariables.size())
                {
                    envVars = AuProcess::EnvironmentGetAll();

                    for (const auto pair : this->startup_.environmentVariables)
                    {
                        envVars.push_back(pair);
                    }
                }
            }
            else
            {
                envVars = this->startup_.environmentVariables;
            }
            std::wstring envVarBlock;

            if (envVars.size() || !this->startup_.bInheritEnvironmentVariables)
            {
                for (const auto &[key, value] : envVars)
                {
                    envVarBlock += AuLocale::ConvertFromUTF8(key);
                    envVarBlock += L'=';
                    envVarBlock += AuLocale::ConvertFromUTF8(value);
                    envVarBlock += L'\x00';
                }

                envVarBlock += L'\x00';

                uCreateFlags |= CREATE_UNICODE_ENVIRONMENT;
            }

            auto result = CreateProcessW(Locale::ConvertFromUTF8(this->startup_.process).c_str(),
                                         Locale::ConvertFromUTF8(this->windowsCli_).data(),
                                         NULL, NULL, true,
                                         uCreateFlags,
                                         envVarBlock.size() ? envVarBlock.data() : nullptr, 
                                         wcwd.size() ? wcwd.data() : nullptr, &startupInfo, &processInfo);

            if (!result)
            {
                DWORD wr = GetLastError();
                SysPushErrorGen("CreateProcess failed");
                return false;
            }

            this->process_ = processInfo.hProcess;
            this->hthread_ = processInfo.hThread;

            RelOtherHandles();

            if (this->startup_.optAffinity)
            {
                {
                // This is good enough for now
                // Minimum supported client: Windows XP [desktop apps | UWP apps]
                    auto primary = this->startup_.optAffinity.Value();
                    auto secondary = this->startup_.optAffinity.Value();
                    secondary.lower = 0;
                    if ((!bool(secondary)) &&
                        (bool(primary)))
                    {
                        ::SetProcessAffinityMask(this->process_, primary.lower);
                        ::SetThreadAffinityMask(this->hthread_, primary.lower);
                    }
                    else
                    {
                        auto cpuSets = primary.ToCpuSets();
                        // TODO:
                    }
                }

                if (!this->startup_.bInDebugMode)
                {
                    ::ResumeThread(this->hthread_);
                }
            }

            if (this->type_ == ESpawnType::eSpawnChildProcessWorker)
            {
            #if defined(AURORA_PLATFORM_WIN32)
                AssignJobWorker(processInfo.hProcess);
            #endif
            }
        }

        // TODO: delegate to a singular worker thread / SetThreadPoolWait
        auto a = [=]()
        {
            HANDLE b[] =
            {
                this->poke_, processInfo.hProcess
            };

            WaitForMultipleObjects(2, b, FALSE, INFINITE);

            DWORD exitCode;
            auto result = GetExitCodeProcess(processInfo.hProcess, &exitCode);
            if (result)
            {
                this->exitCode_ = exitCode;
            }
        };


        this->thread_ = AuThreads::ThreadUnique(AuThreads::ThreadInfo(
            AuMakeShared<AuThreads::IThreadVectorsFunctional>(AuThreads::IThreadVectorsFunctional::OnEntry_t(std::bind(a)),
                                                              AuThreads::IThreadVectorsFunctional::OnExit_t{})
        ));
        
        if (!this->thread_)
        {
            return false;
        }

        this->loopSource_ = AuMakeShared<ProcessAliveLoopSource>(this->process_);
        if (!this->loopSource_)
        {
            return {};
        }

        this->thread_->Run();
        return true;
    }


    void ProcessImpl::RelOtherHandles()
    {
        if (!this->bDontRelOut_)
        {
            AuWin32CloseHandle(this->pipeStdOutWrite_);
        }

        if (!this->bDontRelErr_)
        {
            AuWin32CloseHandle(this->pipeStdErrWrite_);
        }

        if (!this->bDontRelIn_)
        {
            AuWin32CloseHandle(this->pipeStdInRead_);
        }
    }

    AUKN_SYM IProcess *SpawnNew(StartupParameters &&params)
    {
        try
        {
            auto hi = _new ProcessImpl(AuMove(params));
            if (!hi)
            {
                return {};
            }

            if (!hi->Init())
            {
                delete hi;
                return {};
            }

            return hi;
        }
        catch (...)
        {
            return {};
        }
    }

    AUKN_SYM void SpawnRelease(IProcess *process)
    {
        AuSafeDelete<ProcessImpl *>(process);
    }
}