AuroraRuntime/Source/IO/UNIX/FDIpcServer.cpp

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

    File: FDIpcServer.cpp
    Date: 2022-4-12
    Author: Reece
***/
#include <Source/RuntimeInternal.hpp>
#include "UnixIO.hpp"
#include "FDIpcServer.hpp"
#include <sys/socket.h>
#include <sys/un.h>
#include <Source/IO/FS/FS.hpp>
#include <sys/syscall.h>
#include <fcntl.h>

#if defined(AURORA_IS_LINUX_DERIVED)

#if !defined(PIDFD_NONBLOCK)
    #define PIDFD_NONBLOCK O_NONBLOCK
#endif

#endif

namespace Aurora::IO::UNIX
{
    static AuThreadPrimitives::SpinLock gSpinLock;
    static AuThreads::ThreadUnique_t gUnixIPCHandleServeThread;
    static AuThreadPrimitives::EventUnique_t gReadyEvent;
    static AuBST<AuUInt32, int> gFdCookieMap; 
    static bool gHasProcSyscall  {};

    static int WriteDescripitor(int socket, AuUInt32 cookie, int fd)
    {
        if ((fd == -1) ||
            (socket == -1))
        {
            return -1;
        }

        struct iovec iov 
        {
            .iov_base = &cookie,
            .iov_len = sizeof(cookie)
        };

        struct msghdr msg
        {
            .msg_iov = &iov,
            .msg_iovlen = 1,
            .msg_name = nullptr,
            .msg_namelen = 0
        };

    #if defined(AURORA_IS_LINUX_DERIVED) // smells
        
        union {
            char buf[CMSG_SPACE(sizeof(fd))];
            struct cmsghdr align;
        } u;

        msg.msg_control = u.buf;
        msg.msg_controllen = sizeof(u.buf);

        struct cmsghdr *cmsg = CMSG_FIRSTHDR(&msg);
        *cmsg = (struct cmsghdr){.cmsg_level = SOL_SOCKET,
                                 .cmsg_type = SCM_RIGHTS,
                                 .cmsg_len = CMSG_LEN(sizeof(fd))};

        AuWriteU32(CMSG_DATA(cmsg), 0, fd);
    #else // lain hue
        msg.msg_accrights = (caddr_t) &fd;
        msg.msg_accrightslen = sizeof(fd);
    #endif

        return ::sendmsg(socket, &msg, 0);
    }

    static int ReadDescriptor(int socket, AuUInt32 cookie)
    {
        AuUInt32 foreignCookie {};
        int newFd {-1};
        int ret {};

        if (socket == -1)
        {
            return -1;
        }

        struct iovec iov 
        {
            .iov_base = &foreignCookie,
            .iov_len = sizeof(foreignCookie)
        };

        struct msghdr msg
        {
            .msg_iov = &iov,
            .msg_iovlen = 1,
            .msg_name = nullptr,
            .msg_namelen = 0
        };

    #if defined(AURORA_IS_LINUX_DERIVED)
        union {
            char buf[CMSG_SPACE(sizeof(newFd))];
            struct cmsghdr align;
        } u;

        msg.msg_control = u.buf;
        msg.msg_controllen = sizeof(u.buf);
    #else
        msg.msg_accrights = (caddr_t) &newFd;
        msg.msg_accrightslen = sizeof(newFd);
    #endif

        if ((ret = ::recvmsg(socket, &msg, 0)) <= 0)
        {
            return ret;
        }

    #if defined(AURORA_IS_LINUX_DERIVED)
        auto cmptr = CMSG_FIRSTHDR(&msg);
        if (!cmptr)
        {
            SysPushErrorIO();
            return -1;
        }

        if (cmptr->cmsg_len != CMSG_LEN(sizeof(int)))
        {
            SysPushErrorIO("IPC FD Packet interrupted");
            return -1;
        }

        if (cmptr->cmsg_level != SOL_SOCKET)
        {
            SysPushErrorIO("cmsg_level expected SOL_SOCKET");
            return -1;
        }

        if (cmptr->cmsg_type != SCM_RIGHTS)
        {
            SysPushErrorIO("cmsg_type expected SCM_RIGHTS");
            return -1;
        }

        if (foreignCookie != cookie)
        {
            SysPushErrorIO("invalid IPC FD cookie");
            return -1;
        }

        return AuReadU32(CMSG_DATA(cmptr), 0);

    #else

        if (msg.msg_accrightslen != sizeof(int))
        {
            SysPushErrorIO("IPC FD Packet interrupted");
            return -1;
        }

        if (foreignCookie != cookie)
        {
            SysPushErrorIO("invalid IPC FD cookie");
            return -1;
        }

        return newFd;

    #endif
    }

    static bool WriteDescripitorRequest(int socket, AuUInt32 cookie)
    {
        return ::write(socket, &cookie, sizeof(cookie)) == sizeof(cookie);
    }

    static bool ReadDescripitorRequest(int socket, AuUInt32 &cookie)
    {
        timeval tv {};
        tv.tv_usec = 5000;
        ::setsockopt(socket, SOL_SOCKET, SO_RCVTIMEO, AuReinterpretCast<const char *>(&tv), sizeof(tv));
        return ::read(socket, &cookie, sizeof(cookie)) == sizeof(cookie);
    }

    static AuString GetServerPath(pid_t pid)
    {
        AuString path;

#if 0
        if (!AuIOFS::GetUserProgramsFolder(path))
        {
            path = "/opt/";
        }
#endif
        path += "/tmp/UNIX_IPC/";
        path += AuToString(AuUInt32(pid));
        return path;
    }

    static bool GetServerName(sockaddr_un &sockAddr, pid_t pid)
    {
        auto path = GetServerPath(pid);
        if (path.empty())
        {
            return false;
        }

        if (!AuTryResize(path, AuArraySize(sockAddr.sun_path) - 1))
        {
            return false;
        }

        sockAddr.sun_family = AF_UNIX;
        AuMemcpy(sockAddr.sun_path, path.c_str(), path.size());
        return true;
    }

    static void ServerThread()
    {
        sockaddr_un addr {};

        int fd = ::socket(AF_UNIX, SOCK_STREAM, 0);
        if (fd <= 0)
        {
            SysPanic("Couldn't spawn IPC socket");
        }

        if (!GetServerName(addr, getpid()))
        {
            ::close(fd);
            SysPanic("Get IPC socket server name");
        }

        AuIOFS::CreateDirectories(addr.sun_path, true);

        if (::bind(fd, AuReinterpretCast<struct sockaddr *>(&addr), sizeof(addr)) == -1)
        {
            ::close(fd);
            SysPanic("Couldn't bind IPC socket {}", errno);
        }

        ::fcntl(fd, F_SETFL, ::fcntl(fd, F_GETFL, 0) | FD_CLOEXEC);

        if (::listen(fd, 128) == -1)
        {
            ::close(fd);
            SysPanic("Couldn't listen to IPC socket {}", errno);
        }

        gReadyEvent->Set();
        
        while (AuIsThreadRunning())
        {
            AuUInt32 cookie;

            int client = ::accept(fd, NULL, NULL);
            if (client == -1)
            {
                SysPushErrorIO("IPC accept error {}", errno);
                continue;
            }

            if (!ReadDescripitorRequest(client, cookie))
            {
                SysPushErrorIO("IPC request timed out {}", errno);
                goto cont;
            }

            {
                AU_LOCK_GUARD(gSpinLock);

                auto itr = gFdCookieMap.find(cookie);
                if (itr == gFdCookieMap.end())
                {
                    SysPushErrorIO("Missing IPC cookie: {}", cookie);
                    goto cont;
                }

                if (!WriteDescripitor(client, cookie, itr->second))
                {
                    SysPushErrorIO("Couldn't write IPC packet");
                    goto cont;
                }
            }

        cont:
            ::close(client);
        }

        ::close(fd);
    }

    static bool ReadyServer()
    {   
        if (gUnixIPCHandleServeThread)
        {
            return true;
        }

        gReadyEvent = AuThreadPrimitives::EventUnique(false, false, true);
        if (!gReadyEvent)
        {
            return false;
        } 

        gUnixIPCHandleServeThread = AuThreads::ThreadUnique(AuThreads::ThreadInfo(
            AuMakeShared<AuThreads::IThreadVectorsFunctional>(AuThreads::IThreadVectorsFunctional::OnEntry_t(std::bind(ServerThread)),
                                                              AuThreads::IThreadVectorsFunctional::OnExit_t{}),
            "UNIX IPC File Descriptor Server"
        ));

        if (!gUnixIPCHandleServeThread)
        {
            gReadyEvent.reset();
            return false;
        }

        if (!gUnixIPCHandleServeThread->Run())
        {
            return bool(gUnixIPCHandleServeThread = {});
        }

        gReadyEvent->Lock();
        gReadyEvent.reset();
        return true;
    }

    static bool Ready()
    {
        AU_LOCK_GUARD(gSpinLock);

        if (gHasProcSyscall)
        {
            return true;
        }

        return ReadyServer();
    }

    bool FDServe(int fd, IPC::IPCToken &token)
    {
        if (gHasProcSyscall)
        {
            token.cookie = fd;
            return true;
        }

        if (!Ready())
        {
            return false;
        }

        {
            AU_LOCK_GUARD(gSpinLock);

            do
            {
                token.NewId();
            }
            while (AuExists(gFdCookieMap, token.cookie));

            return AuTryInsert(gFdCookieMap, token.cookie, fd);
        }
    }

    void FDServeEnd(const IPC::IPCToken &handle)
    {
        if (gHasProcSyscall)
        {
            return;
        }

        AU_LOCK_GUARD(gSpinLock);
        AuTryRemove(gFdCookieMap, handle.cookie);
    }

    bool FDAccept(const IPC::IPCToken &handle, int &outFd)
    {
        sockaddr_un addr;
        outFd = -1;

        if (gHasProcSyscall)
        {
        #if defined(AURORA_IS_LINUX_DERIVED)

            int pid = pidfd_open(handle.ToUnixPid(), 0);
            if (pid <= 0)
            {
                SysPushErrorIO("Couldn't open IPC server pid, error: {}", pid);
                return false;
            }

            int result = pidfd_getfd(pid, handle.cookie, 0);
            if (result <= 0)
            {
                SysPushErrorIO("Couldn't get IPC fd, error: {} {}", result, errno);
                ::close(pid);
                return false;
            }

            ::close(pid);
            outFd = result;
            return true;

        #endif
        }

        int fd = ::socket(AF_UNIX, SOCK_STREAM | SOCK_CLOEXEC, 0);
        if (fd <= -1)
        {
            fd = ::socket(AF_UNIX, SOCK_STREAM, 0);
        }

        if (fd <= -1)
        {
            SysPushErrorIO("No Resource?");
            return false;
        }

        if (!GetServerName(addr, handle.ToUnixPid()))
        {
            ::close(fd);
            SysPushErrorIO("Couldn't get remote IPC socket server name");
            return false;
        }

        if (::connect(fd, AuReinterpretCast<struct sockaddr *>(&addr), sizeof(addr)) == -1)
        {
            SysPushErrorIO("Couldn't connect to remote IPC fd server");
            ::close(fd);
            return false;
        }

        if (!WriteDescripitorRequest(fd, handle.cookie))
        {
            SysPushErrorIO("Couldn't write FD request to IPC server");
            ::close(fd);
            return false;
        }

        outFd = ReadDescriptor(fd, handle.cookie);
        if (outFd == -1)
        {
            SysPushErrorNested();
        }

        ::close(fd);
        return true;
    }

    void InitIPCBackend()
    {
    #if defined(AURORA_IS_LINUX_DERIVED)
        auto &platform = AuSwInfo::GetPlatformInfo();

        if ((platform.uKernelMajor > 5) ||
            ((platform.uKernelMajor == 5) && (platform.uKernelMinor >= 6)))
        {
            gHasProcSyscall = gRuntimeConfig.linuxConfig.bIOLinuxHasProcIpcPerms;
        }
    #endif
    }

    void DeinitIPCBackend()
    {
        gUnixIPCHandleServeThread.reset();
        gFdCookieMap.clear();
        gReadyEvent.reset();
    }
}