AuroraRuntime/Source/Processes/AuProcess.Unix.cpp

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

    File: AuProcess.Unix.cpp
    File: Process.Linux.cpp
    Date: 2021-6-12
    Author: Reece
***/
#include <RuntimeInternal.hpp>
#include "AuProcesses.hpp"
#include "AuProcess.Unix.hpp"

#include <unistd.h>
#include <fcntl.h>
#include <sys/syscall.h>
#include <sys/types.h>
#include <sys/wait.h>

#include <Source/IO/AuIOHandle.Unix.hpp>
#include <Source/IO/CompletionGroup/CompletionGroup.hpp>

#if defined(AURORA_COMPILER_CLANG)
    // warning: enumeration values 'kEnumCount' and 'kEnumInvalid' not handled in switch [-Wswitch
    #pragma clang diagnostic ignored "-Wswitch"
    // Yea, I don't give a shit.
#endif

#if defined(AURORA_IS_LINUX_DERIVED)
    #include <sys/prctl.h>
    #include <setjmp.h>

    #if !defined(CLONE_CLEAR_SIGHAND)
        #define CLONE_CLEAR_SIGHAND 0x100000000ULL
    #endif
#endif


#if defined(AURORA_PLATFORM_BSD)
    #if defined(__FreeBSD__)
        #include <sys/param.h>
    #endif

    #include <sys/cpuset.h>
#endif

#if defined(AURORA_IS_LINUX_DERIVED)
    #include <sched.h>
#endif

namespace Aurora::Processes
{
    static AuRWLock gRWLock;
    static AuHashMap<pid_t, ProcessImpl *> gPidLookupMap; 

    struct ProcessAliveLoopSource : AuLoop::LSEvent
    {
        ProcessAliveLoopSource();
        virtual AuLoop::ELoopSource GetType() override;
    };

    ProcessAliveLoopSource::ProcessAliveLoopSource() : LSEvent(false, false, true)
    {}

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

        for (const auto &arg : this->startup_.args)
        {
            this->cargs_.push_back(arg.c_str());
            this->debug_ += arg + " ";
        }

        this->cargs_.push_back(nullptr);
        if (this->debug_.size())
        {
            this->debug_.resize(this->debug_.size() - 1);
        }

        this->type_ = this->startup_.type;
    }

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

       {
            AU_LOCK_GUARD(gRWLock->AsWritable());

            if (this->alive_)
            {
                if (this->type_ == ESpawnType::eSpawnThreadLeader)
                {
                    ::kill(this->pidt_, SIGCONT);
                }
            }

            AuTryRemove(gPidLookupMap, this->pidt_);
        }

       if (auto &pGroup = this->pCompletionGroup_)
       {
           this->bHasExited = true;
           AuStaticCast<IO::CompletionGroup::CompletionGroup>(pGroup)->UnsafeRemoveItem(AuUnsafeRaiiToShared(this));
       }

        ShutdownPipes();
    }

    AuUInt ProcessImpl::GetProcessId()
    {
        return this->pidt_;
    }

    bool ProcessImpl::HasExited()
    { 
        return this->bHasExited;
    }

    void ProcessImpl::ByeLol(AuSInt code)
    {
        this->exitCode_ = code;
        this->bHasExited = true;

        if (this->finished_)
        {
            this->finished_->Set();
        }

        if (this->loopSource_)
        {
            this->loopSource_->Set();
        }

        if (this->fsStream_)
        {
            this->fsStream_->CheckProcess();
        }

        if (this->fsErrorStream_)
        {
            this->fsErrorStream_->CheckProcess();
        }

        if (auto &pCompletionGroup = this->pCompletionGroup_)
        {
            pCompletionGroup->TryTriggerLater();
        }
    }

    void ProcessImpl::ShutdownPipes()
    {
        if (!this->bDontRelOut_)
        {
            if (auto fd = AuExchange(pipeStdOut_[1], {}))
            {
                ::close(fd);
            }
        }

        if (!this->bDontRelErr_)
        {
            if (auto fd = AuExchange(pipeStdErr_[1], {}))
            {
                ::close(fd);
            }
        }

        if (!this->bDontRelIn_)
        {
            if (auto fd = AuExchange(pipeStdIn_[0], {}))
            {
                ::close(fd);
            }
        }
    }

    bool ProcessImpl::TryKill()
    {
        AU_LOCK_GUARD(gRWLock->AsReadable());
        
        if (this->alive_)
        {
            if (::kill(this->pidt_, SIGTERM) == 0)
            {
                return this->finished_->LockMS(500);
            }
        }

        return true;
    }

    bool ProcessImpl::Terminate()
    {
        AU_LOCK_GUARD(gRWLock->AsReadable());

        if (this->alive_)
        {
            return ::kill(this->pidt_, SIGKILL) == 0;
        }
        
        return true;
    }

    AuSPtr<Aurora::Threading::IWaitable> ProcessImpl::AsWaitable()
    {
        return this->finished_;
    }

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

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

    bool ProcessImpl::Read(EStandardHandle stream, const AuMemoryViewStreamWrite &destination, bool nonblock)
    {
        destination.outVariable = 0;

        if (!IO::EStandardStreamIsValid(stream) || (stream == EStandardHandle::eInputStream))
        {
            SysPushErrorArg("Invalid Stream");
            return false;
        }

        if (!destination)
        {
            SysPushErrorArg();
            return false;
        }

        auto handle = stream == EStandardHandle::eErrorStream ? this->pipeStdErr_[0] : this->pipeStdOut_[0];
        if (handle < 0)
        {
            SysPushErrorUninitialized();
            return false;
        }

        auto control = ::fcntl(handle, F_GETFL);
        auto ref = control;

        if (nonblock)
        {
            control |= O_NONBLOCK;
        }
        else
        {
            control &= ~O_NONBLOCK;
        }

        if (ref != control)
        {
            ::fcntl(handle, F_SETFL, control);
        }

        int tmp;

        do
        {
            tmp = ::read(handle, destination.ptr, destination.length);
        }
        while ((tmp == -1 && errno == EINTR));

        if (tmp <= 0)
        {
            if (tmp == 0)
            {
                return nonblock;
            }
        
            SysPushErrorMem();
            return false;
        }

        destination.outVariable = tmp;
        return true;
    }

    bool ProcessImpl::Write(const AuMemoryViewStreamRead &source)
    {
        source.outVariable = 0;

        auto handle = this->pipeStdIn_[1];
        if (!handle)
        {
            return false;
        }

        auto control = ::fcntl(handle, F_GETFL);
        auto ref = control;

        if (/*nonblock*/ true)
        {
            control |= O_NONBLOCK;
        }
        else
        {
            control &= ~O_NONBLOCK;
        }

        if (ref != control)
        {
            ::fcntl(handle, F_SETFL, control);
        }

        return ::write(handle, source.ptr, source.length) == source.length;
    }

    static bool InitProcessStdHandles(EStreamForward fwd, int *fds, AuIO::EStandardStream stream, IO::IIOHandle *pHandle)
    {
        AuOptionalEx<AuUInt64> optHandle;
        bool bIsRead = stream == AuIO::EStandardStream::eInputStream;

        switch (fwd)
        {
        case EStreamForward::eIOHandle:
            optHandle = pHandle->GetOSHandleSafe();

            if (!optHandle)
            {
                return false;
            }

            if (bIsRead)
            {
                fds[0] = dup((int)optHandle.value());
            }
            else
            {
                fds[1] = dup((int)optHandle.value());
            }
            break;
        case EStreamForward::eCurrentProcess:
            if (stream == AuIO::EStandardStream::eInputStream)
            {
                fds[0] = dup(STDIN_FILENO);
                if (fds[0] < 0)
                {
                    return false;
                }
            }
            else if (stream == AuIO::EStandardStream::eErrorStream)
            {
                fds[1] = dup(STDERR_FILENO);
                if (fds[1] < 0)
                {
                    return false;
                }
            }
            else if (stream == AuIO::EStandardStream::eOutputStream)
            {
                fds[1] = dup(STDOUT_FILENO);
                if (fds[1] < 0)
                {
                    return false;
                }
            }
            else
            {
                return false;
            }
            break;
        case EStreamForward::eAsyncPipe:
            if (::pipe(fds))
            {
                SysPushErrorMem();
                return false;
            }
            break;
        case EStreamForward::eNull:
            if (bIsRead)
            {
                fds[0] = PosixOpen("/dev/null", O_RDWR);
            }
            else
            {
                fds[1] = PosixOpen("/dev/null", O_RDWR);
            }
            break;
        case EStreamForward::eNewConsoleWindow:
            SysPushErrorGeneric("AuProcesses is not the right place for PTY support. At least not in this form (this level of abstraction only cares for pipes).");
            return false;
        }

        if (fds[0] > 0)
        {
            if (!AuIO::SetFDShareAccess(fds[0], bIsRead))
            {
                return false;
            }
        }

        if (fds[1] > 0)
        {
            if (!AuIO::SetFDShareAccess(fds[1], !bIsRead))
            {
                return false;
            }
        }

        // Yea, I guess we can't really guard against threads racing to leak close on exec-less fds
        // Let's just not care on UNIX. 
        // Worst case, we leak a daemons input/output FD to a different daemon of a differing permission level
        // ...and somehow you can find that fd
        // ...and probably dup it before ::start()
        // ...and issue the right commands over the stream to do something bad
        // Shouldn't need to worry about his for a while.

        return true;
    }

    bool ProcessImpl::Init()
    {
        InitProcessStdHandles(this->startup_.fwdOut, this->pipeStdOut_, AuIO::EStandardStream::eOutputStream, this->startup_.handleOutStream);
        InitProcessStdHandles(this->startup_.fwdErr, this->pipeStdErr_, AuIO::EStandardStream::eErrorStream, this->startup_.handleErrorStream);
        InitProcessStdHandles(this->startup_.fwdIn, this->pipeStdIn_, AuIO::EStandardStream::eInputStream, this->startup_.handleInputStream);

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

        this->finished_ = AuThreadPrimitives::EventShared(false, false, true);
        
        if (!this->finished_)
        {
            return false;
        }

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

            this->fsStream_ = AuMakeShared<ProcessPipeFileStream>();
            if (!this->fsStream_)
            {
                return false;
            }

            this->fsHandle_->InitFromPairMove(this->pipeStdOut_[0] == -1 ? AuOptionalEx<AuUInt64> {} : AuOptionalEx<AuUInt64> { (AuUInt64)this->pipeStdOut_[0] } , 
                                              this->pipeStdIn_[1] == -1 ? AuOptionalEx<AuUInt64> {} : AuOptionalEx<AuUInt64> { (AuUInt64)this->pipeStdIn_[1] });
            this->fsStream_->Init(this->fsHandle_);
            this->fsStream_->MakeProcess(this);
        }

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

            this->fsErrorStream_ = AuMakeShared<ProcessPipeFileStream>();
            if (!this->fsErrorStream_)
            {
                return false;
            }

            this->fsErrorHandle_->InitFromPairMove(this->pipeStdErr_[0] == -1 ? AuOptionalEx<AuUInt64> {} : AuOptionalEx<AuUInt64> { (AuUInt64)this->pipeStdErr_[0] }, {});
            this->fsErrorStream_->Init(this->fsErrorHandle_);
            this->fsErrorStream_->MakeProcess(this);
        }

        if (auto pCompletionGroup = AuExchange(this->startup_.pAutoJoinCompletionGroup, {}))
        {
            (void)this->TryAttachProcessExitToCompletionGroup(pCompletionGroup);
        }

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

    AuSPtr<IO::IIOHandle> ProcessImpl::GetOutputAndInputHandles()
    {
        return this->fsHandle_;
    }

    AuSPtr<IO::IIOHandle> ProcessImpl::GetErrorStreamHandle()
    {
        return this->fsErrorHandle_;
    }

    AuSPtr<IO::IStreamReader> ProcessImpl::ToStreamReader(EStandardHandle stream)
    {
        if (stream == EStandardHandle::eErrorStream)
        {
            if (auto pThat = this->GetErrorStreamHandle())
            {
                if (auto pThat2 = AuFS::OpenBlockingFileStreamFromHandleShared(pThat))
                {
                    return { pThat2, pThat2->ToStreamReader() };
                }
            }
        }
        else
        {
            if (auto pThat = this->GetOutputAndInputHandles())
            {
                if (auto pThat2 = AuFS::OpenBlockingFileStreamFromHandleShared(pThat))
                {
                    return { pThat2, pThat2->ToStreamReader() };
                }
            }
        }

        return {};
    }

    AuSPtr<IO::IStreamWriter> ProcessImpl::ToStreamWriter()
    {
        if (auto pThat = this->GetErrorStreamHandle())
        {
            if (auto pThat2 = AuFS::OpenBlockingFileStreamFromHandleShared(pThat))
            {
                return { pThat2, pThat2->ToStreamWriter() };
            }
        }

        return {};
    }

    void ProcessImpl::ForkMain()
    {
        PosixDoForkHooks();

        {
            ::dup2(this->pipeStdIn_[0], STDIN_FILENO);
            ::close(this->pipeStdIn_[0]);
        }
            
        {
            ::dup2(this->pipeStdErr_[1], STDERR_FILENO);
            ::close(this->pipeStdErr_[1]);
        }
            
        {
            ::dup2(this->pipeStdOut_[1], STDOUT_FILENO);
            ::close(this->pipeStdOut_[1]);
        }

    #if defined(AURORA_IS_LINUX_DERIVED)
        if (this->type_ == ESpawnType::eSpawnChildProcessWorker)
        {
            ::prctl(PR_SET_PDEATHSIG, SIGTERM);
        }
    #endif

        if (this->type_ != ESpawnType::eSpawnOvermap)
        {
            ::setsid();
        }
        
        if (!this->startup_.bInheritEnvironmentVariables)
        {
            try
            {
                AuList<AuString> keys;
                
                for (const auto &[key, value] : AuProcess::EnvironmentGetAll())
                {
                    keys.push_back(key);
                }

                AuProcess::EnvironmentRemoveMany(keys);
            }
            catch (...)
            {
                SysPanic("Couldn't fork");
            }
        }

        PosixFDYeetus();

        if (this->startup_.environmentVariables.size())
        {
            SysAssert(AuProcess::EnvironmentSetMany(this->startup_.environmentVariables));
        }
    
    #if defined(AURORA_IS_XNU_DERIVED)
        if (this->startup_.workingDirectory)
        {
            ::pthread_chdir_np(this->startup_.workingDirectory.value().c_str());
        }
    #else
        if (this->startup_.workingDirectory)
        {
            ::chroot(this->startup_.workingDirectory.value().c_str());
        }
    #endif

        if (this->startup_.optAffinity)
        {
            auto affinity = this->startup_.optAffinity.Value();

            #if defined(__FreeBSD__)
            
                cpuset_setid(CPU_WHICH_PID, -1, (cpusetid_t *)&affinity.lower); // I guess?
            
            #elif defined(AURORA_IS_LINUX_DERIVED) || (defined(AURORA_IS_POSIX_DERIVED) && defined(_GNU_SOURCE))

                cpu_set_t cpuset;
                CPU_ZERO(&cpuset);
                
                AuUInt8 index {};
                while (affinity.CpuBitScanForward(index, index))
                {
                    CPU_SET(index, &cpuset);
                    index++;
                }

                if (!CPU_COUNT(&cpuset))
                {
                    return;
                }

                #if defined(AURORA_IS_LINUX_DERIVED) 
                    sched_setaffinity(0, sizeof(cpu_set_t), &cpuset);
                #else
                    pthread_setaffinity_np(pthread_self(), cpuset_size(cpuSet), &cpuset);
                #endif

            #elif defined(AURORA_IS_POSIX_DERIVED)
                
                if (auto cpuSet = cpuset_alloc())
                {
                    cpuset_init(cpuSet);

                    AuUInt8 index {};
                    while (affinity.CpuBitScanForward(index, index))
                    {
                        cpuset_set_cpu(cpuSet, index, 1);
                        index++;
                    }

                    pthread_setaffinity_np(pthread_self(), cpuset_size(cpuSet), cpuSet);

                    cpuset_free(cpuSet);
                }

            #else
                
                // Who else?

            #endif
        }
        
        if (this->startup_.posixApplySandboxCOW)
        {
            this->startup_.posixApplySandboxCOW();
        }
    
        ::execv(this->startup_.process.c_str(), (char * const *)this->cargs_.data()); // https://pubs.opengroup.org/onlinepubs/9699919799/functions/exec.html
        
        SysPushErrorGen("execv didn't overwrite the process map. Launch: {} ({})", this->startup_.process, this->debug_);
    }

    bool ProcessImpl::Start()
    {
        this->exitCode_ = 0x10110100;

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

        pid_t pid;
        {
            pid = ::fork();

            if (pid == 0)
            {
                this->ForkMain();
                SysPanic();
                return false;
            }
            else if (pid > 0)
            {
                if (pipeStdOut_[1] &&
                    pipeStdOut_[1] != -1)
                {
                    ::close(AuExchange(pipeStdOut_[1], 0));
                }

                if (pipeStdErr_[1] &&
                    pipeStdErr_[1] != -1)
                {
                    ::close(AuExchange(pipeStdErr_[1], 0));
                }

                if (pipeStdIn_[0] &&
                    pipeStdIn_[0] != -1)
                {
                    ::close(AuExchange(pipeStdIn_[0], 0));
                }

                this->pidt_ = pid;
                this->alive_ = true;

                {
                    AU_LOCK_GUARD(gRWLock->AsWritable());
                    SysAssert(AuTryInsert(gPidLookupMap, pid, this));
                }

                return true;
            }
            else
            {
                return false;
            }
        }

        return true;
    }

    bool ProcessImpl::TryAttachProcessExitToCompletionGroup(const AuSPtr<IO::CompletionGroup::ICompletionGroup> &pCompletionGroup)
    {
        if (this->pCompletionGroup_ ||
            !pCompletionGroup)
        {
            return false;
        }

        this->pCompletionGroup_ = pCompletionGroup;
        pCompletionGroup->AddWorkItem(AuUnsafeRaiiToShared(this));
        pCompletionGroup->TryTriggerLater();
        return true;
    }

    IO::CompletionGroup::ICompletionGroupWorkHandle *ProcessImpl::ToCompletionGroupHandle()
    {
        return this;
    }

    bool ProcessImpl::HasCompletedForGCWI()
    {
        return this->HasExited();
    }

    void ProcessImpl::CleanupForGCWI()
    {
        AuResetMember(this->pCompletionGroup_);
    }

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

    AUKN_SYM IProcess *SpawnNew(StartupParameters &&params)
    {
        auto ret = _new ProcessImpl(AuMove(params));

        if (!ret)
        {
            SysPushErrorMem();
            return {};
        }

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

        return ret;
    }

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

    static void HandleChildTermiantion(pid_t pid, int code)
    {
        AU_LOCK_GUARD(gRWLock->AsReadable());

        auto handler = gPidLookupMap.find(pid);
        if (handler == gPidLookupMap.end())
        {
            return;
        }

        {
            auto process = *handler;
            process.second->ByeLol(code);
        }
    }

    static void SigChldHandler(int)
    {
        int code;
        pid_t pid;

        while (true)
        {
            pid = wait3(&code, WNOHANG, nullptr);

            if ((pid == 0) || 
                (pid == -1))
            {
                break;
            }
            
            HandleChildTermiantion(pid, code);
        }
    }

    void InitUnix()
    {
        struct sigaction action =
        {
            .sa_handler = SigChldHandler,
            .sa_flags = SA_ONSTACK
        };

        ::sigemptyset(&action.sa_mask);
        ::sigaction(SIGCHLD, &action, nullptr);
    }

    void PosixProcessShutdown()
    {
    #if !defined(AURORA_IS_LINUX_DERIVED)
        decltype(gPidLookupMap) cpy;
        {
            AU_LOCK_GUARD(gRWLock->AsReadable());
            cpy = gPidLookupMap;
        }
        for (const auto &[pid, pProcess] : cpy)
        {
            pProcess->HookMainDeath();
        }
    #endif
    }

    void DeinitUnix()
    {
        struct sigaction action =
        {
            .sa_handler = SIG_DFL,
            .sa_flags = SA_ONSTACK | SA_NOCLDWAIT
        };

        ::sigemptyset(&action.sa_mask);
        ::sigaction(SIGCHLD, &action, nullptr);

        PosixProcessShutdown();
    }
}