send_fd.cpp

send_fd.cpp

#include <boost/asio/coroutine.hpp>
#include <boost/asio/local/connect_pair.hpp>
#include <boost/asio/local/stream_protocol.hpp>
#include <iostream>
#include <sys/wait.h>

std::size_t
send_fd(boost::asio::local::stream_protocol::socket& socket,
        int fd,
        boost::system::error_code& ec)
{
    ::msghdr msg{};
    union {
        ::cmsghdr cmsghdr;
        char control[CMSG_SPACE(sizeof(int))];
    } cmsgu;

    ::iovec io = {.iov_base = (void*)"ABC", .iov_len = 3};
    msg.msg_iov = &io;
    msg.msg_iovlen = 1;
    msg.msg_control = cmsgu.control;
    msg.msg_controllen = sizeof(cmsgu.control);

    ::cmsghdr* cmsg = CMSG_FIRSTHDR(&msg);
    cmsg->cmsg_len = CMSG_LEN(sizeof(int));
    cmsg->cmsg_level = SOL_SOCKET;
    cmsg->cmsg_type = SCM_RIGHTS;

    std::memcpy(CMSG_DATA(cmsg), &fd, sizeof(fd));

    if (auto size = ::sendmsg(socket.native_handle(), &msg, 0); size < 0)
    {
        ec = {errno, boost::system::system_category()};
        return 0;
    }
    else
    {
        return size;
    }
}

std::size_t
read_fd(boost::asio::local::stream_protocol::socket& socket,
        int& fd,
        boost::system::error_code& ec)
{
    ::msghdr msg = {0};

    char m_buffer[256];
    ::iovec io = {.iov_base = m_buffer, .iov_len = sizeof(m_buffer)};
    msg.msg_iov = &io;
    msg.msg_iovlen = 1;

    union {
        ::cmsghdr cmsghdr;
        char control[CMSG_SPACE(sizeof(int))];
    } cmsgu;
    msg.msg_control = &cmsgu;
    msg.msg_controllen = sizeof(cmsgu.control);

    if (auto size = ::recvmsg(socket.native_handle(), &msg, 0); size < 0)
    {
        ec = {errno, boost::system::system_category()};
        return 0;
    }
    else
    {
        ::cmsghdr* cmsg = CMSG_FIRSTHDR(&msg);
        unsigned char* data = CMSG_DATA(cmsg);
        std::memcpy(&fd, data, sizeof(fd));
        return size;
    }
}

struct child_op
{
    std::unique_ptr<boost::asio::local::stream_protocol::socket> socket;
    boost::asio::coroutine coro{};

    void operator()(boost::system::error_code ec)
    {
        int fd;
        BOOST_ASIO_CORO_REENTER(coro)
        {
            BOOST_ASIO_CORO_YIELD socket->async_wait(
              boost::asio::socket_base::wait_read, std::move(*this));
            if (ec)
            {
                std::cerr << "Child wait error: " << ec.message() << '\n';
                return;
            }

            read_fd(*socket, fd, ec);
            ::write(fd, "hello", strlen("hello"));
            ::close(fd);
            if (ec)
            {
                std::cerr << "Child read error: " << ec.message() << '\n';
                return;
            }

            std::cout << "Child finished \n";
        }
    }
};

int
open_file(boost::system::error_code& ec)
{
    auto fd = ::open("test.txt", O_APPEND | O_WRONLY | O_CREAT);
    if (fd < 0)
        ec = {errno, boost::system::system_category()};
    else
        ec = {};
    return fd;
}

struct parent_op
{
    std::unique_ptr<boost::asio::local::stream_protocol::socket> socket;
    boost::asio::coroutine coro{};

    void operator()(boost::system::error_code ec)
    {
        int fd;
        BOOST_ASIO_CORO_REENTER(coro)
        {
            BOOST_ASIO_CORO_YIELD socket->async_wait(
              boost::asio::socket_base::wait_write, std::move(*this));
            if (ec)
            {
                std::cerr << "Parent wait_write error: " << ec.message()
                          << '\n';
                return;
            }

            fd = open_file(ec);
            if (ec)
            {
                std::cerr << "Parent open_file error: " << ec.message() << '\n';
                return;
            }

            send_fd(*socket, fd, ec);
            if (ec)
            {
                std::cerr << "Parent send_fd error: " << ec.message() << '\n';
                return;
            }
            ::close(fd);

            BOOST_ASIO_CORO_YIELD socket->async_wait(
              boost::asio::socket_base::wait_read, std::move(*this));
            if (ec)
            {
                std::cerr << "Parent wait_read error: " << ec.message() << '\n';
                return;
            }
        }
    }
};

void
do_fork(boost::asio::io_context& ctx, boost::system::error_code ec)
{
    auto parent =
      std::make_unique<boost::asio::local::stream_protocol::socket>(ctx);
    auto child =
      std::make_unique<boost::asio::local::stream_protocol::socket>(ctx);
    boost::asio::local::connect_pair(*parent, *child);
    ctx.notify_fork(boost::asio::execution_context::fork_prepare);

    if (::pid_t pid = ::fork(); pid == 0)
    {
        ctx.notify_fork(boost::asio::execution_context::fork_child);
        child_op{std::move(child)}(boost::system::error_code{});
    }
    else if (pid > 0)
    {
        ctx.notify_fork(boost::asio::execution_context::fork_parent);
        parent_op{std::move(parent)}(boost::system::error_code{});
    }
    else
    {
        ec = {errno, boost::system::system_category()};
    }
}

void
wait_for_children(boost::system::error_code& ec)
{
    int status = 0;
    for (;;)
    {
        if (::pid_t pid = ::wait(&status); pid < 0 && errno == ECHILD)
        {
            ec = {};
            return;
        }
        else if (pid < 0)
        {
            ec = {errno, boost::system::system_category()};
            return;
        }
    }
}

int
main()
{
    boost::asio::io_context ctx;

    boost::system::error_code ec;
    do_fork(ctx, ec);
    if (ec)
    {
        std::cerr << "Error, fork failed: " << ec.message() << '\n';
        return 1;
    }

    ctx.run();

    wait_for_children(ec);
    if (ec)
    {
        std::cerr << "Error, waiting for children failed: " << ec.message()
                  << '\n';
        return 1;
    }
}