Luna/libos/include/os/IPC.h

/**
 * @file IPC.h
 * @author apio (cloudapio.eu)
 * @brief Inter-process communication primitives.
 *
 * @copyright Copyright (c) 2023, the Luna authors.
 *
 */

#pragma once
#include <os/Action.h>
#include <os/LocalClient.h>
#include <os/LocalServer.h>

#define IPC_ENUM_SERVER(name) __##name##_SERVER_ERROR = 0
#define IPC_ENUM_CLIENT(name) __##name##_CLIENT_ERROR = 0

namespace os
{
    namespace IPC
    {
        static constexpr usize IPC_STRING_LENGTH = 256;

#define IPC_STRING(name) char name[os::IPC::IPC_STRING_LENGTH];
#define COPY_IPC_STRING(name)                                                                                          \
    TRY(String::from_string_view(StringView::from_fixed_size_cstring(name, os::IPC::IPC_STRING_LENGTH)))
#define SET_IPC_STRING(name, value) strlcpy(name, value, os::IPC::IPC_STRING_LENGTH)

        /**
         * @brief Class used to send and receive IPC messages on the server, using an already established connection to
         * the client.
         */
        class ClientConnection
        {
          public:
            /**
             * @brief Creates a new IPC connection object from an os::LocalServer connection.
             *
             * @param connection The existing connection to use, obtained after calling accept() on an os::LocalServer
             * object.
             * @return Result<OwnedPtr<ClientConnection>> An error, or the IPC connection object.
             */
            static Result<OwnedPtr<ClientConnection>> adopt_connection(LocalServer::Client&& connection);

            /**
             * @brief Check for new messages from the client. The message handler will be called if there is a new
             * message. If the connection is non-blocking, this function returns EAGAIN if there are no new
             * messages, otherwise it blocks until one is received.
             *
             * @return Result<void> Whether the operation succeeded.
             */
            Result<void> check_for_messages();

            /**
             * @brief Set the message handler for this connection. This function is required and will be called every
             * time a new message is received, passing the connection object and the ID of the message type.
             *
             * @param handler The message handler to use.
             * @param arg An arbitrary argument to pass to the handler.
             */
            void set_message_handler(Function<ClientConnection&, u8, void*>&& handler, void* arg)
            {
                m_message_handler = move(handler);
                m_arg = arg;
            }

            /**
             * @brief Read message data from the connection. This function should be called only after the other side of
             * the connection has signaled that it is going to send a message of the specified type, for example inside
             * the message handler.
             *
             * @tparam T The type of the message to read.
             * @param out The variable in which to store the message.
             * @return Result<bool> An error, or whether the message was actually read and stored. If this value is
             * false, this function will have to be called again, either because the connection is non-blocking and the
             * message data has not been sent yet or because a signal interrupted the call.
             */
            template <typename T> Result<bool> read_message(T& out)
            {
                auto rc = m_connection.recv_typed(out);
                if (rc.has_error())
                {
                    if (rc.error() == EAGAIN)
                    {
                        m_ipc_in_progress = true;
                        m_ipc_saved_id = T::ID;
                        return false;
                    }
                    if (rc.error() == EINTR)
                    {
                        m_ipc_in_progress = true;
                        m_ipc_saved_id = T::ID;
                        return false;
                    }
                    else
                        return rc.release_error();
                }
                return true;
            }

            /**
             * @brief Sends an error code, indicating that an operation could not be performed. It is best to send an
             * error back only when the client is expecting it, that is to say, when they are waiting for a reply to a
             * synchronous message.
             *
             * @param error The error code.
             * @return Result<void> Whether the operation succeded.
             */
            Result<void> send_error(int error);

            /**
             * @brief Sends a message without waiting for a reply.
             *
             * @tparam T The type of the message.
             * @param message The message to send.
             * @return Result<void> Whether the operation succeeded.
             */
            template <typename T> Result<void> send_async(const T& message)
            {
                u8 id = T::ID;
                TRY(m_connection.send_typed(id));
                TRY(m_connection.send_typed(message));
                return {};
            }

            /**
             * @brief Closes the connection.
             */
            void disconnect()
            {
                m_connection.disconnect();
            }

            /**
             * @brief Returns the file descriptor associated with this connection.
             *
             * @return int The file descriptor.
             */
            int fd()
            {
                return m_connection.fd();
            }

          private:
            ClientConnection(LocalServer::Client&& connection);

            LocalServer::Client m_connection;
            Function<ClientConnection&, u8, void*> m_message_handler;
            void* m_arg;

            bool m_ipc_in_progress { false };
            u8 m_ipc_saved_id { 0 };
        };

        /**
         * @brief Class used to send and receive IPC messages on the client.
         */
        class Client
        {
          public:
            /**
             * @brief Connect to an IPC server and return a connection object.
             *
             * @param path The path of the socket used by the IPC server.
             * @param blocking Whether the connection should block when waiting for messages.
             * @return Result<OwnedPtr<Client>> An error, or a new connection object.
             */
            static Result<OwnedPtr<Client>> connect(StringView path, bool blocking);

            /**
             * @brief Check for new messages from the server. The message handler will be called if there is a new
             * message. If the connection is non-blocking, this function returns EAGAIN if there are no new
             * messages, otherwise it blocks until one is received.
             *
             * @return Result<void> Whether the operation succeeded.
             */
            Result<void> check_for_messages();

            /**
             * @brief Set the message handler for this connection. This function is required and will be called every
             * time a new message is received, passing the connection object and the ID of the message type.
             *
             * @param handler The message handler to use.
             * @param arg An arbitrary argument to pass to the handler.
             */
            void set_message_handler(Function<Client&, u8, void*>&& handler, void* arg)
            {
                m_message_handler = move(handler);
                m_arg = arg;
            }

            /**
             * @brief Read message data from the connection. This function should be called only after the other side of
             * the connection has signaled that it is going to send a message of the specified type, for example inside
             * the message handler.
             *
             * @tparam T The type of the message to read.
             * @param out The variable in which to store the message.
             * @return Result<bool> An error, or whether the message was actually read and stored. If this value is
             * false, this function will have to be called again, either because the connection is non-blocking and the
             * message data has not been sent yet or because a signal interrupted the call.
             */
            template <typename T> Result<bool> read_message(T& out)
            {
                auto rc = m_connection->recv_typed(out);
                if (rc.has_error())
                {
                    if (rc.error() == EAGAIN)
                    {
                        m_ipc_in_progress = true;
                        m_ipc_saved_id = T::ID;
                        return false;
                    }
                    if (rc.error() == EINTR)
                    {
                        m_ipc_in_progress = true;
                        m_ipc_saved_id = T::ID;
                        return false;
                    }
                    else
                        return rc.release_error();
                }
                return true;
            }

            /**
             * @brief Sends a message without waiting for a reply.
             *
             * @tparam T The type of the message.
             * @param message The message to send.
             * @return Result<void> Whether the operation succeeded.
             */
            template <typename T> Result<void> send_async(const T& message)
            {
                u8 id = T::ID;
                TRY(m_connection->send_typed(id));
                TRY(m_connection->send_typed(message));
                return {};
            }

            /**
             * @brief Sends a message and waits for a reply.
             *
             * @tparam ResponseType The type of the response.
             * @tparam T The type of the message.
             * @param message The message to send.
             * @return Result<ResponseType> An error, or the response.
             */
            template <typename ResponseType, typename T> Result<ResponseType> send_sync(const T& message)
            {
                u8 id = T::ID;
                TRY(m_connection->send_typed(id));
                TRY(m_connection->send_typed(message));

                // We allow receiving 5 messages of different types, but if those have passed and we still don't have a
                // reply, fail with ENOMSG.
                int max_other_messages = 5;

                while (max_other_messages)
                {
                    u8 response_id;
                    auto rc = m_connection->recv_typed(response_id);
                    if (rc.has_error() && (rc.error() == EAGAIN || rc.error() == EINTR)) continue;

                    if (response_id == 0) // Error result
                    {
                        while (1)
                        {
                            int code;
                            rc = m_connection->recv_typed(code);
                            if (rc.has_error() && (rc.error() == EAGAIN || rc.error() == EINTR)) continue;
                            return err(code);
                        }
                    }

                    if (response_id != ResponseType::ID)
                    {
                        m_message_handler(*this, response_id, m_arg);
                        max_other_messages--;
                        continue;
                    }

                    while (1)
                    {
                        ResponseType response;
                        rc = m_connection->recv_typed(response);
                        if (rc.has_error() && (rc.error() == EAGAIN || rc.error() == EINTR)) continue;
                        return response;
                    }
                }

                return err(ENOMSG);
            }

            /**
             * @brief Closes the connection.
             */
            void disconnect()
            {
                m_connection->disconnect();
            }

            /**
             * @brief Returns the file descriptor associated with this connection.
             *
             * @return int The file descriptor.
             */
            int fd()
            {
                return m_connection->fd();
            }

          private:
            Client(OwnedPtr<LocalClient>&& connection);

            OwnedPtr<LocalClient> m_connection;
            Function<Client&, u8, void*> m_message_handler;
            void* m_arg;

            bool m_ipc_in_progress { false };
            u8 m_ipc_saved_id { 0 };
        };

        /**
         * @brief API used to notify a parent process when a child process finishes initialization. The Notifier struct
         * is the parent part of the API.
         *
         */
        struct Notifier
        {
            int pfds[2];

            /**
             * @brief Create a new Notifier.
             *
             * This function will create a pipe for the parent and child to communicate.
             *
             * @return Notifier The new Notifier object.
             */
            static Notifier create();

            /**
             * @brief Hook the Notifier into any child process started afterwards.
             *
             * This will set an environment variable, which if detected by a child process, will use it to notify the
             * parent whenever it's ready.
             *
             * The recommended order to call this API is:
             * hook()
             * fork+exec
             * unhook()
             */
            void hook();

            /**
             * @brief Remove the previously created environment variable, so that any future child processes will not
             * notify this Notifier.
             *
             */
            void unhook();

            /**
             * @brief Wait for a child process to be ready. If several child processes are hooked by the hook() method,
             * this method will only catch the first one that notifies the parent.
             *
             * @param timeout If positive, specifies the timeout after which the function fails if no notification is
             * received.
             * @return true The child is ready.
             * @return false The method timed out.
             */
            bool wait(int timeout = -1);

            /**
             * @brief Combines hook(), unhook() and wait() into one single method. This method takes a function, and
             * executes it in a "hooked" context, so that any child process started by this function will automatically
             * detect the parent when it's ready (if it supports the notification API). Then, the function waits for the
             * child to be ready.
             *
             * @param action The function to run.
             * @param timeout If positive, specifies the timeout after which the function fails if no notification is
             * received.
             * @return true The child is ready.
             * @return false The method timed out.
             */
            static bool run_and_wait(os::Action&& action, int timeout = -1);
        };

        /**
         * @brief Use this function to notify a parent process whenever your program finishes initialization.
         *
         * If the parent has not used the Notifier API to request a notification from the child process, this function
         * does nothing.
         *
         */
        void notify_parent();
    }
}