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/Threading/Primitives/AuSemaphore.Unix.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

namespace Aurora::Processes
{
    static AuThreadPrimitives::RWLockUnique_t 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_);
        }

        ShutdownPipes();
    }

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

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

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

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

    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)
    {
        auto handle = stream == EStandardHandle::eStdError ? 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)
    {
        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)
    {
        switch (fwd)
        {
        case EStreamForward::eCurrentProcess:
            // Intentionally NO-OP
            break;
        case EStreamForward::eAsyncPipe:
            if (::pipe(fds))
            {
                SysPushErrorMem();
                return false;
            }
            break;
        case EStreamForward::eNull:
            fds[0] = ::open("/dev/null", O_RDWR);
            fds[1] = ::open("/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).");
            break;
        }

        return true;
    }

    bool ProcessImpl::Init()
    {
        InitProcessStdHandles(this->startup_.fwdOut, this->pipeStdOut_);
        InitProcessStdHandles(this->startup_.fwdErr, this->pipeStdErr_);
        InitProcessStdHandles(this->startup_.fwdIn, this->pipeStdIn_);

        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], this->pipeStdIn_[1]);
            this->fsStream_->Init(this->fsHandle_);
        }

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

    void ProcessImpl::ForkMain()
    {
        if (this->startup_.fwdIn != EStreamForward::eCurrentProcess)
        {
            ::dup2(pipeStdIn_[0], STDIN_FILENO);
            ::close(pipeStdIn_[0]);
            ::close(pipeStdIn_[1]);
        }
            
        if (this->startup_.fwdErr != EStreamForward::eCurrentProcess)
        {
            ::dup2(pipeStdErr_[1], STDERR_FILENO);
            ::close(pipeStdIn_[0]);
            ::close(pipeStdIn_[1]);
        }
            
        if (this->startup_.fwdOut != EStreamForward::eCurrentProcess)
        {
            ::dup2(pipeStdOut_[1], STDOUT_FILENO);
            ::close(pipeStdIn_[0]);
            ::close(pipeStdIn_[1]);
        }

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

        ::setsid();

    #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
    
        ::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)
        {
            ::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, given {} ({})", this->startup_.process, this->debug_);
            return false;
        }

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

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

                if (pipeStdErr_[1])
                {
                    ::close(AuExchange(pipeStdErr_[1], 0));
                }

                if (pipeStdIn_[0])
                {
                    ::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;
    }
    
    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()
    {
        gRWLock = AuThreadPrimitives::RWLockUnique();

        struct sigaction action =
        {
            .sa_handler = SigChldHandler,
            .sa_flags = SA_ONSTACK
        };

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

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

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

        gRWLock.reset();
    }
}