AuroraRuntime/Source/IPC/IPCPipe.NT.cpp

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

    File: IPCPipe.NT.cpp
    Date: 2022-4-15
    Author: Reece
***/
#include <Source/RuntimeInternal.hpp>
#include "IPC.hpp"
#include "IPCHandle.hpp"
#include "IPCPipe.NT.hpp"

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

namespace Aurora::IPC
{
    //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
    // Pipes
    //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
    struct IPCPipeImpl;

    struct IPCHasConnectionEvent : Loop::LSEvent
    {
        IPCHasConnectionEvent(AuSPtr<IPCPipeImpl> parent);

        bool IsSignaled() override;
        Loop::ELoopSource GetType() override;

        bool OnTrigger(AuUInt handle) override;

    private:
        AuWPtr<IPCPipeImpl> parent_;
    };

    struct IPCPipeImpl : IPCPipe, Loop::LSHandle, AuEnableSharedFromThis<IPCPipeImpl>
    {
        IPCPipeImpl(HANDLE clientHandle, HANDLE serverHandle, const IPCHandle &handle);
        ~IPCPipeImpl();

        PROXY_INTERNAL_INTERFACE_(LSHandle::)
        
        virtual AuSPtr<IO::FS::IAsyncTransaction> NewAsyncTransaction() override;

        virtual AuSPtr<Loop::ILoopSource> AsReadChannelIsOpen() override;
        virtual AuSPtr<Loop::ILoopSource> AsReadChannelHasData() override;

        virtual bool Read(const Memory::MemoryViewStreamWrite &write, bool nonblocking) override;
        virtual bool Write(const Memory::MemoryViewStreamRead &read) override;
        virtual AuString ExportToString() override;

        bool IsSignaled() override;
        bool WaitOn(AuUInt32 timeout) override;
        Loop::ELoopSource GetType() override;

        HANDLE GetPipeHandle();

        void TryConnect();
        OVERLAPPED overlapped {};

    private:
        HANDLE serverHandle_ {INVALID_HANDLE_VALUE};
        HANDLE clientHandle_ {INVALID_HANDLE_VALUE};
        IPCHandle ipcHandle_;
        AuSPtr<IO::FS::FileHandle> fsHandle_;
        AuSPtr<IO::FS::NtAsyncFileStream> fsStream_;
        AuSPtr<Loop::ILSEvent> hasClient_;
        AuSPtr<IPCHasConnectionEvent> lshasConnection_;
        bool bFirstTime {true};
    };

    IPCHasConnectionEvent::IPCHasConnectionEvent(AuSPtr<IPCPipeImpl> parent) : parent_(parent), LSEvent(false, false, true)
    {

    }

    bool IPCHasConnectionEvent::IsSignaled()
    {
        return OnTrigger(0);
    }

    Loop::ELoopSource IPCHasConnectionEvent::GetType()
    {
        return Loop::ELoopSource::eSourceIPCHasClient;
    }

    bool IPCHasConnectionEvent::OnTrigger(AuUInt handle)
    {
        auto parent = this->parent_.lock();
        if (!parent)
        {
            SysPushErrorMem("IPC pipe is dead");
            return false;
        }

        parent->TryConnect();
        return WaitForSingleObject(parent->overlapped.hEvent, 0) == WAIT_OBJECT_0;
    }

    IPCPipeImpl::IPCPipeImpl(HANDLE clientHandle, HANDLE serverHandle, const IPCHandle &handle) :
        serverHandle_(serverHandle), clientHandle_(clientHandle), ipcHandle_(handle)
    {
        
        if (serverHandle != INVALID_HANDLE_VALUE)
        {
            this->hasClient_ = Loop::NewLSEvent(false, false, true);
        }

        this->fsHandle_ = AuMakeShared<IO::FS::FileHandle>();
        this->fsStream_ = AuMakeShared<IO::FS::NtAsyncFileStream>();

        this->fsHandle_->Init(this->GetPipeHandle());
        this->fsStream_->Init(this->fsHandle_);

        TryConnect();
    }

    void IPCPipeImpl::TryConnect()
    {
        DWORD idc;

        if (this->serverHandle_ == INVALID_HANDLE_VALUE)
        {
            return;
        }

        this->overlapped.hEvent = (HANDLE)AuStaticCast<Loop::LSEvent>(this->hasClient_)->GetHandle();
     
        bool firstTime = AuExchange(bFirstTime, false);
        if (firstTime || 
            (WaitForSingleObject(this->overlapped.hEvent, 0) == WAIT_OBJECT_0 &&
             GetOverlappedResult(this->serverHandle_, &this->overlapped, &idc, false))
           )
        {
            ResetEvent(this->overlapped.hEvent);

            if (ConnectNamedPipe(this->serverHandle_, &this->overlapped))
            {
                bFirstTime = true;
                TryConnect();
            }
            else if (GetLastError() == ERROR_IO_PENDING)
            {
                // No-op
            }
            else if (GetLastError() == ERROR_PIPE_CONNECTED)
            {
                SetEvent(this->overlapped.hEvent);
            }
        }
        else
        {
            if (WaitForSingleObject(this->overlapped.hEvent, 0) == WAIT_OBJECT_0)
            {
                ResetEvent(this->overlapped.hEvent);
            }
        }
    }

    IPCPipeImpl::~IPCPipeImpl()
    {
    }

    AuSPtr<Loop::ILoopSource> IPCPipeImpl::AsReadChannelIsOpen()
    {
        if (this->serverHandle_ == INVALID_HANDLE_VALUE)
        {
            return {};
        }

        if (!this->lshasConnection_)
        {
            this->lshasConnection_ = AuMakeShared<IPCHasConnectionEvent>(AuSharedFromThis());
        }

        return AuStaticCast<Loop::ILSEvent>(this->lshasConnection_);
    }

    AuSPtr<Loop::ILoopSource> IPCPipeImpl::AsReadChannelHasData()
    {
        // TODO (Hack): we should at least make a shared timer
        return AuUnsafeRaiiToShared(this);
    }
    
    AuSPtr<IO::FS::IAsyncTransaction> IPCPipeImpl::NewAsyncTransaction()
    {
        return this->fsStream_->NewTransaction();
    }

    bool IPCPipeImpl::Read(const Memory::MemoryViewStreamWrite &write, bool nonblocking)
    {
        DWORD size = write.length;

        TryConnect();

        auto h = this->GetPipeHandle();
        if (h == INVALID_HANDLE_VALUE)
        {
            SysPushErrorUninitialized();
            return false;
        }

        if (nonblocking || !write.ptr)
        {
            DWORD avail {};
            if (!PeekNamedPipe(h, NULL, NULL, NULL, &avail, NULL))
            {
                return false;
            }

            if (!avail)
            {
                return true;
            }

            size = AuMin(size, avail);
        }

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

        auto ret = ::ReadFile(h, write.ptr, size, &size, NULL);
        write.outVariable = size;
        return ret;
    }

    bool IPCPipeImpl::Write(const Memory::MemoryViewStreamRead &read)
    {
        auto h = this->GetPipeHandle();
        if (h == INVALID_HANDLE_VALUE)
        {
            SysPushErrorUninitialized();
            return false;
        }

        TryConnect();

        DWORD temp;
        if (!::WriteFile(h, read.ptr, read.length, &temp, nullptr))
        {
            SysPushErrorIO();
            return false;
        }

        read.outVariable = temp;
        return true;
    }

    HANDLE IPCPipeImpl::GetPipeHandle()
    {
        return this->clientHandle_ == INVALID_HANDLE_VALUE ? this->serverHandle_ : this->clientHandle_;
    }

    bool IPCPipeImpl::IsSignaled()
    {
        DWORD avail {};

        TryConnect();

        if (!PeekNamedPipe(this->GetPipeHandle(), NULL, NULL, NULL, &avail, NULL))
        {
            return false;
        }

        return avail;
    }

    bool IPCPipeImpl::WaitOn(AuUInt32 timeout)
    {
        return LSHandle::WaitOn(timeout);
    }

    Loop::ELoopSource IPCPipeImpl::GetType()
    {
        return Loop::ELoopSource::eSourceIPCReadPipe;
    }

    AuString IPCPipeImpl::ExportToString()
    {
        TryConnect();

        return this->clientHandle_ == INVALID_HANDLE_VALUE ?
               this->ipcHandle_.ToString() :
               AuString {};
    }

    AUKN_SYM AuSPtr<IPCPipe>      NewPipe()
    {
        IPCHandle handle;
        handle.NewId();

        auto name = "\\\\.\\pipe\\" + handle.ToNTPath();
        auto maxLength = 16 * AuHwInfo::GetPageSize() ? AuHwInfo::GetPageSize() : 4096;
        auto pipeServer = CreateNamedPipeA(name.c_str(),
                                          PIPE_ACCESS_DUPLEX | FILE_FLAG_FIRST_PIPE_INSTANCE | FILE_FLAG_WRITE_THROUGH | FILE_FLAG_OVERLAPPED,
                                          PIPE_WAIT,
                                          1,
                                          maxLength,
                                          maxLength,
                                          NMPWAIT_WAIT_FOREVER,
                                          nullptr);

        if (pipeServer == INVALID_HANDLE_VALUE)
        {
            SysPushErrorIO("{}", GetLastError());
            return {};
        }

        auto object = AuMakeShared<IPCPipeImpl>(INVALID_HANDLE_VALUE, pipeServer, handle);
        if (!object)
        {
            SysPushErrorMem();
            AuWin32CloseHandle(pipeServer);
            return {};
        }

        return object;
    }

    AUKN_SYM AuSPtr<IPCPipe>      ImportPipe(const AuString &handleString)
    {
        IPCHandle handle;
        HANDLE pipe;

        if (!handle.FromString(handleString))
        {
            SysPushErrorParseError();
            return {};
        }

        auto name = "\\\\.\\pipe\\" + handle.ToNTPath();
        pipe = CreateFileA(name.c_str(),
                            GENERIC_WRITE | GENERIC_READ,
                            0,
                            NULL,
                            OPEN_ALWAYS,
                            FILE_ATTRIBUTE_NORMAL | FILE_FLAG_OVERLAPPED,
                            NULL);

        if (pipe == INVALID_HANDLE_VALUE)
        {
            if (GetLastError() == ERROR_PIPE_BUSY)
            {
                SysPushErrorIO("Pipe is used -> a client has already connected or the nt server is not ready");
                return {};
            }
            SysPushErrorIO("{}", GetLastError());
            return {};
        }

        auto object = AuMakeShared<IPCPipeImpl>(pipe, INVALID_HANDLE_VALUE, handle);
        if (!object)
        {
            SysPushErrorMem();
            AuWin32CloseHandle(pipe);
            return {};
        }

        return object;
    }
}

// > The pipe created by UWP process with name \\.\pipe\Local\PipeName is converted to \\.\pipe\Sessions\<SessionId>\AppContainerNamedObjects\<AppContainerSid>\PipeName. 
// > I can use this to communicate between UWP as server and Win32 as client.
//  https://jike.in/qa/?qa=103904/
//  ...good to know