AuroraRuntime/Source/IO/NT/HANDLEPipeServer.cpp

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

    File: HANDLEPipeServer.cpp
    Date: 2024-01-02
    Author: Reece
    Note: Rushed Windows HANDLE sharing of any unnamed kernel object
***/
#include <Source/RuntimeInternal.hpp>
#include <Source/IO/IPC/AuIPCHandle.hpp>

namespace Aurora::IO::NT
{
    static AuMutex gMutex;
    static AuBST<AuUInt32, HANDLE> gHANDLECookieMap;
    static AuThreads::ThreadUnique_t gServerThread {};
    static const auto kDefaultMessageSize = 8192u;

    static auto GetIPCServerNameOfPid(AuUInt32 uPid)
    {
    #if defined(AURORA_PLATFORM_WIN32)
        return fmt::format("\\\\.\\pipe\\_HANDLESOF_{}", uPid);
    #else
        return fmt::format("\\\\.\\pipe\\LOCAL\\_HANDLESOF_{}", uPid);
    #endif
    }
    
    HANDLE GetHandleForToken(AuUInt32 uToken)
    {
        AU_LOCK_GUARD(gMutex);
    
        auto itr = gHANDLECookieMap.find(uToken);
        if (itr != gHANDLECookieMap.end())
        {
            return itr->second;
        }
        else
        {
            return INVALID_HANDLE_VALUE;
        }
    }

    // TODO: use AuLoop
    DWORD WINAPI WaitForCookieThread(LPVOID lpvParam)
    {
        AuUInt32 uTokenID {};
        HANDLE hHandle { INVALID_HANDLE_VALUE }, hPipe { (HANDLE)lpvParam };
        DWORD dwBytesRead {}, dwWritten {};
        BOOL bSuccess { true };

        if (!lpvParam)
        {
            return (DWORD)-1;
        }

        do
        {
            bSuccess = ReadFile(hPipe,
                                &uTokenID,
                                sizeof(uTokenID),
                                &dwBytesRead,
                                NULL);
            if (!bSuccess ||
                dwBytesRead == 0)
            {
                if (GetLastError() != ERROR_BROKEN_PIPE)
                {
                    SysPushErrorIO("IO HANDLE Server Error");
                }
                break;
            }

            hHandle = GetHandleForToken(uTokenID);
            if (hHandle == INVALID_HANDLE_VALUE)
            {
                SysPushErrorIO("Tried to serve bad token: {}", uTokenID);
                break;
            }

            {
                if (!pGetNamedPipeClientProcessId)
                {
                    SysPushErrorIO("UWP Platform Error");
                    break;
                }

                ULONG pid {};
                if (!pGetNamedPipeClientProcessId(hPipe, &pid))
                {
                    SysPushErrorIO("IO HANDLE Server Error");
                    break;
                }

                HANDLE hProcess {};
                if (!(hProcess = OpenProcess(PROCESS_DUP_HANDLE, FALSE, pid)))
                {
                    SysPushErrorIO("IO HANDLE Server Error");
                    break;
                }

                if (!DuplicateHandle(GetCurrentProcess(),
                                     hHandle,
                                     hProcess,
                                     &hHandle,
                                     0,
                                     FALSE,
                                     DUPLICATE_SAME_ACCESS))
                {
                    AuWin32CloseHandle(hProcess);
                    SysPushErrorIO("IO HANDLE Server Error");
                    break;
                }

                AuWin32CloseHandle(hProcess);
            }

            bSuccess = WriteFile(hPipe,        
                                 &hHandle,
                                 sizeof(hHandle),
                                 &dwWritten,
                                 NULL);        

            if (!bSuccess ||
                sizeof(hHandle) != sizeof(HANDLE))
            {
                SysPushErrorIO("IO HANDLE Server Error");
                break;
            }
        }
        while (false);

        AuWin32CloseHandle(hPipe);
        return bSuccess ? 0 : (DWORD)-1;
    }

    static void IPCServer()
    {
        HANDLE hPipe = INVALID_HANDLE_VALUE;
        HANDLE hThread = NULL;

        auto name = GetIPCServerNameOfPid(GetCurrentProcessId());

        while (AuIsThreadRunning())
        {
            // TODO: use AuLoop
            hPipe = CreateNamedPipeA(name.c_str(),          
                                     PIPE_ACCESS_DUPLEX,    
                                     PIPE_TYPE_MESSAGE |    
                                     PIPE_READMODE_MESSAGE |
                                     PIPE_WAIT,             
                                     PIPE_UNLIMITED_INSTANCES,
                                     kDefaultMessageSize,
                                     kDefaultMessageSize,
                                     0,
                                     NULL);

            if (hPipe == INVALID_HANDLE_VALUE)
            {
                SysPushErrorIOResourceFailure("NO IPC Handle Server");
                return;
            }

            if (ConnectNamedPipe(hPipe, NULL) ||
                (GetLastError() == ERROR_PIPE_CONNECTED))
            {
                // TODO: use AuLoop
                hThread = CreateThread(NULL, 
                                       0,
                                       WaitForCookieThread,
                                       (LPVOID)hPipe, 
                                       0,             
                                       NULL);

                if (hThread == NULL)
                {
                    DisconnectNamedPipe(hPipe);
                }
                else
                {
                    AuWin32CloseHandle(hThread);
                }
            }
            else
            {
                AuWin32CloseHandle(hPipe);
            }
        }
    }

    static void StartServerIfNotAlready()
    {
        static AuInitOnce gInitOnce;
        gInitOnce.Call([]()
        {
            gServerThread = AuThreads::Spawn(IPCServer);
            SysAssert(gServerThread, "No NT Handle Server Thread");
            gServerThread->SetName("Microsoft NT HANDLE IPC Server/IPC-Provider");
        });
    }

    bool FDServe(HANDLE hHandle, IPC::IPCToken &outHandle)
    {
        if (!hHandle ||
            hHandle == INVALID_HANDLE_VALUE)
        {
            return false;
        }

        {
            StartServerIfNotAlready();
        }

        {
            AU_LOCK_GUARD(gMutex);

            do
            {
                outHandle.cookie = AuRng::RngU32();
            }
            while (AuExists(gHANDLECookieMap, outHandle.cookie));

            outHandle.word = outHandle.cookie;
            outHandle.pid = GetCurrentProcessId();
            outHandle.path[0] = '\x00'; // msvc removing something it shouldn't have?
            //strcat(outHandle.path, AuToString(outHandle.cookie).c_str()); // ignore warning

            return AuTryInsert(gHANDLECookieMap, AuConstReference(outHandle.cookie), AuConstReference(hHandle));
        }
    }

    void FDServeEnd(const IPC::IPCToken &handle)
    {
        AU_LOCK_GUARD(gMutex);
        AuTryRemove(gHANDLECookieMap, handle.cookie);
    }

    bool FDAccept(const IPC::IPCToken &handle, HANDLE &outHandle)
    {
        outHandle = INVALID_HANDLE_VALUE;

        HANDLE hPipe {};
        HANDLE hHandle {};
        BOOL bSuccess {};
        DWORD dwWritten {};
        auto name = GetIPCServerNameOfPid(handle.pid);

        hPipe = Win32Open(AuLocale::ConvertFromUTF8(name).c_str(),
                          GENERIC_WRITE | GENERIC_READ,
                          0,
                          false,
                          OPEN_ALWAYS,
                          0,
                          FILE_ATTRIBUTE_NORMAL);

        if (hPipe == 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 false;
            }
            else
            {
                SysPushErrorIO("No Server: {}", GetLastError());
                return false;
            }
        }

        bSuccess = WriteFile(hPipe,
                             &handle.cookie,
                             sizeof(AuUInt32),
                             &dwWritten,
                             NULL);

        if (!bSuccess)
        {
            AuWin32CloseHandle(hPipe);
            SysPushErrorIO("No Server");
            return false;
        }

        FlushFileBuffers(hPipe);

        bSuccess = ReadFile(hPipe,
                            &hHandle,
                            sizeof(hHandle),
                            &dwWritten,
                            NULL);
        if (!bSuccess)
        {
            AuWin32CloseHandle(hPipe);
            SysPushErrorIO("No Handle");
            return false;
        }

        outHandle = hHandle;
        return true;
    }

    void ShutdownNTIPCHandleServer()
    {
        gServerThread.reset();
    }

    AUKN_SYM AuString ShareFileDescriptor(HANDLE hHandle)
    {
        IPC::IPCToken token;
        IPC::IPCHandle handle;

        if (!FDServe(hHandle, token))
        {
            return {};
        }

        handle.pid = token.pid;
        handle.PushId(IPC::EIPCHandleType::eIPCSharedFd, token);

        return handle.ToString();
    }

    AUKN_SYM void ShareFileDescriptorStop(const AuROString &handle)
    {
        IPC::IPCHandle handle2;
        if (!handle2.FromString(handle))
        {
            SysPushErrorIO("Invalid handle string");
            return;
        }

        if (auto val = handle2.GetToken(IPC::EIPCHandleType::eIPCSharedFd, 0))
        {
            FDServeEnd(val->token);
        }
        else
        {
            SysPushErrorSyntax("Invalid handle string");
        }
    }

    AUKN_SYM HANDLE ShareFileDescriptorAccept(const AuROString &handle)
    {
        IPC::IPCHandle handle2;
        HANDLE hHandle { INVALID_HANDLE_VALUE };

        if (!handle2.FromString(handle))
        {
            SysPushErrorIO("Invalid handle string");
            return hHandle;
        }

        auto val = handle2.GetToken(IPC::EIPCHandleType::eIPCSharedFd, 0);
        if (!val)
        {
            SysPushErrorSyntax("Invalid handle string");
            return hHandle;
        }

        if (!FDAccept(val->token, hHandle))
        {
            SysPushErrorNested();
            return hHandle;
        }

        return hHandle;
    }
}