NET_PacketSocket.h

Go to the documentation of this file.

/*
 * PROPRIETARY INFORMATION.  This software is proprietary to
 * Side Effects Software Inc., and is not to be reproduced,
 * transmitted, or disclosed in any way without written permission.
 *
 * NAME:        NET_PacketSocket.h
 *
 * COMMENTS:
 *
 *
 */

#ifndef __NET_PACKETSOCKET_H__
#define __NET_PACKETSOCKET_H__

#include "NET_API.h"

#include <UT/UT_BoostAsio.h>
#include <UT/UT_Debug.h>
#include <UT/UT_NonCopyable.h>
#include <UT/UT_SharedPtr.h>
#include <UT/UT_StringHolder.h>
#include <UT/UT_Endian.h>
#include <UT/UT_WorkBuffer.h>
#include <UT/UT_StringArray.h>

NET_API bool
net_uncompressBuffer(UT_WorkBuffer& wbuf, exint comp_size, exint orig_size);

class NET_API NET_PacketSocket
{
private:
    class Impl : public UTenable_shared_from_this<Impl>
    {
    public:
        explicit Impl(ASIO_TCPSocket&& socket)
            : myTimer(socket.get_executor()), myStream(std::move(socket))
        {
        }

        ASIO_TCPSocket& stream() { return myStream; }
        void resetTimer()
        {
            UT_WeakPtr<Impl> me = weak_from_this();
            myTimer.expires_from_now(hboost::posix_time::minutes(10));
            myTimer.async_wait(
                    [me = weak_from_this()](
                            const hboost::system::error_code& ec)
                    {
                        if (ec == hboost::asio::error::operation_aborted)
                            return;
                        if (auto s = me.lock(); s)
                        {
                            s->close();
                        }
                    });
        }
        void cancelTimer()
        {
            myTimer.cancel();
        }

        void close()
        {
            // Closing the socket will propogate
            hboost::system::error_code ignored_ec;
            myStream.shutdown(ASIO_TCPSocket::shutdown_both, ignored_ec);
            myStream.close(ignored_ec);
            myTimer.cancel();
        }

    private:
        /// {@
        /// Used to kill the socket if no activity has occurred within the specified
        /// amount of time.
        ASIO_DeadlineTimer myTimer;
        hboost::posix_time::time_duration myTimeout;
        /// @}
        ASIO_TCPSocket myStream;
    };

public:
    explicit NET_PacketSocket(ASIO_TCPSocket&& socket);
    explicit NET_PacketSocket(ASIO_IOContext& ctx);
    virtual ~NET_PacketSocket() = default;
    UT_NON_COPYABLE(NET_PacketSocket)

    /// Shutdown the socket. After this is called the socket becomes dead
    /// and no other calls should be made.
    void close()
    {
        hboost::asio::dispatch(
                myImpl->stream().get_executor(),
                [me = myImpl->weak_from_this()]()
                {
                    if (auto s = me.lock(); s)
                        s->close();
                });
    }
    /// Access to the underlying socket. This only meant to be used in cases
    /// where information about the connection is required (ie. remote endpoint)
    /// NB: this method should only ever be called prior to the first read/write
    /// of the socket for thread safety.
    ASIO_TCPSocket& nextLayer() { return myImpl->stream(); }

    template <typename DynamicBuffer, typename Handler>
    HBOOST_ASIO_INITFN_RESULT_TYPE(
            Handler,
            void(hboost::system::error_code, std::size_t))
    asyncRead(DynamicBuffer& buffer, Handler&& handler);
    template <typename Handler>
    HBOOST_ASIO_INITFN_RESULT_TYPE(
            Handler,
            void(hboost::system::error_code, std::size_t))
    asyncWrite(UT_StringArray&& buffers, Handler&& handler);
    template <typename Handler>
    HBOOST_ASIO_INITFN_RESULT_TYPE(Handler, void(hboost::system::error_code))
    asyncConnect(const ASIO_TCPSocket::endpoint_type& ep, Handler&& handler)
    {
        return myImpl->stream().async_connect(
                ep, std::forward<Handler>(handler));
    }

private:
    template <class AsyncStream, class DynamicBuffer>
    class AsyncReadImpl;
    template <class AsyncStream>
    class AsyncWriteImpl;

    UT_SharedPtr<Impl> myImpl;
};

// -----------------------------------------------------------------------------
// Packet Read Op
// -----------------------------------------------------------------------------

template <typename AsyncStream, typename DynamicBuffer>
class NET_PacketSocket::AsyncReadImpl
{
public:
    AsyncReadImpl(
            AsyncStream& stream,
            const UT_SharedPtr<NET_PacketSocket::Impl>& backend,
            DynamicBuffer& buffer)
        : myStream(stream), myImpl(backend), myBuffer(buffer)
    {
    }

    template <typename Self>
    void operator()(
            Self &self,
            hboost::system::error_code ec = {},
            std::size_t bytes_transfered = 0)
    {
        myBytesWritten += bytes_transfered;

        UT_SharedPtr<NET_PacketSocket::Impl> impl = myImpl.lock();
        if (!impl)
        {
            return self.complete(
                    hboost::asio::error::operation_aborted, myBytesWritten);
        }

        if (ec)
        {
            return self.complete(ec, myBytesWritten);
        }

        if (myIsFirst)
        {
            // If this is our first request then kick off a read
            myIsFirst = false;
            impl->resetTimer();
            myStream.async_read_some(myBuffer.prepare(512), std::move(self));
            return;
        }

        // We received more data. Reset the connection dead timer.
        impl->cancelTimer();

        myBuffer.commit(bytes_transfered);

        // If the packet is not compressed and no read size has been read then
        // check the header
        if (!myIsCompressed && myReadPacketSize < 0)
        {
            // Parse the original payload length and compressed flag.
            static constexpr int header_size = 5;
            if (myBuffer.size() >= header_size)
            {
                uint32 len = 0;
                static_assert(sizeof(len) == 4);
                ASIO_MutableBuffer tmp(&len, sizeof(len));
                ASIO_ConstBuffer read_buffer = myBuffer.data();
                // Retrieve the payload length
                hboost::asio::buffer_copy(tmp, read_buffer);
                UTtomips(len);
                myReadPacketSize = len;
                myFinalPacketSize = len;
                // We have a max size of 4096
                if (myReadPacketSize > 4096)
                {
                    return self.complete(
                            hboost::asio::error::message_size, myBytesWritten);
                }

                // Check if the playload is compressed or not.
                myIsCompressed
                        = *((const char*)read_buffer.data() + sizeof(unsigned))
                          == 1;
                if (myIsCompressed)
                    myReadPacketSize = -1;
                myBuffer.consume(header_size);
            }
        }
        // If the packet is compressed read the compressed size
        if (myIsCompressed && myReadPacketSize < 0)
        {
            static constexpr int header_size = 4;
            if (myBuffer.size() >= header_size)
            {
                uint32 len = 0;
                static_assert(sizeof(len) == header_size);
                ASIO_MutableBuffer tmp(&len, sizeof(len));
                ASIO_ConstBuffer read_buffer = myBuffer.data();
                // Retrieve the payload length
                hboost::asio::buffer_copy(tmp, read_buffer);
                UTtomips(len);
                myReadPacketSize = len;
                myBuffer.consume(header_size);

                /// Something went wrong. The uncompressed size should never
                /// be smaller then the compressed size.
                if (myFinalPacketSize < myReadPacketSize)
                {
                    return self.complete(
                            hboost::asio::error::connection_aborted,
                            myBytesWritten);
                }
            }
        }
        // If there is at least enough read data for this packet then handle the
        // request
        if (myReadPacketSize >= 0 && myBuffer.size() >= myReadPacketSize)
        {
            ASIO_ConstBuffer read_buffer = myBuffer.data();
            UT_WorkBuffer data(
                    (const char*)read_buffer.data(), myReadPacketSize);

            if (myIsCompressed)
            {
                if (!net_uncompressBuffer(
                            data, myReadPacketSize, myFinalPacketSize))
                {
                    return self.complete(
                            hboost::system::errc::make_error_code(
                                    hboost::system::errc::bad_message),
                            myBytesWritten);
                }
            }

            // Reset the buffer information so that we can read again
            self.complete({}, myBytesWritten);
            return;
        }

        // Guess our next read size
        std::size_t read_size = 512;
        if (myReadPacketSize > 0)
            read_size = static_cast<std::size_t>(myReadPacketSize);
        const auto size = myBuffer.size();
        const auto limit = myBuffer.max_size() - size;
        std::size_t actual_read_size = std::min(
                std::max(std::size_t(512), myBuffer.capacity() - size),
                std::min(read_size, limit));
        impl->resetTimer();
        // We need more data so read more and then continue parsing.
        myStream.async_read_some(
                myBuffer.prepare(actual_read_size), std::move(self));
    }

private:
    DynamicBuffer& myBuffer;
    bool myIsFirst = true;
    bool myIsCompressed = false;
    exint myReadPacketSize = -1;
    exint myFinalPacketSize = -1;
    std::size_t myBytesWritten = 0;
    AsyncStream& myStream;
    UT_WeakPtr<NET_PacketSocket::Impl> myImpl;
};

// -----------------------------------------------------------------------------
// NET_IPacketSocket::AsyncWriteImpl
// -----------------------------------------------------------------------------

template <typename AsyncStream>
class NET_PacketSocket::AsyncWriteImpl
{
public:
    AsyncWriteImpl(
            AsyncStream& stream,
            const UT_SharedPtr<NET_PacketSocket::Impl>& impl,
            UT_StringArray&& write_queue)
        : myStream(stream), myImpl(impl), myWriteQueue(write_queue)
    {
    }

    template <typename Self>
    void operator()(
            Self& self,
            hboost::system::error_code ec = {},
            std::size_t bytes_transfered = 0)
    {
        myBytesWritten += bytes_transfered;

        UT_SharedPtr<NET_PacketSocket::Impl> impl = myImpl.lock();
        if (!impl)
        {
            return self.complete(
                    hboost::asio::error::operation_aborted, myBytesWritten);
        }

        if (ec)
        {
            return self.complete(ec, myBytesWritten);
        }

        exint total = bytes_transfered;
        while (total > 0 && myWriteQueue.size() > 0)
        {
            exint i = myIndex + total;
            if (i >= myWriteQueue[0].length())
            {
                exint remain = myWriteQueue[0].length() - myIndex;
                total -= remain;
                myIndex = 0;
                myWriteQueue.removeIndex(0);
            }
            else
            {
                myIndex += total;
                break;
            }
        }

        if (myWriteQueue.size() == 0)
        {
            return self.complete({}, myBytesWritten);
        }

        // To improve performance dont use scatter/gather if the messages sent
        // are only 1.
        if (myWriteQueue.size() == 1)
        {
            ASIO_ConstBuffer buffer(
                    myWriteQueue[0].buffer() + myIndex,
                    myWriteQueue[0].length() - myIndex);
            myStream.async_write_some(buffer, std::move(self));
        }
        else
        {
            std::vector<ASIO_ConstBuffer> buffers;
            buffers.reserve(myWriteQueue.size());
            buffers.emplace_back(
                    myWriteQueue[0].buffer() + myIndex,
                    myWriteQueue[0].length() - myIndex);
            for (exint i = 1; i < myWriteQueue.size(); i++)
                buffers.emplace_back(
                        myWriteQueue[i].buffer(), myWriteQueue[i].length());
            myStream.async_write_some(buffers, std::move(self));
        }
    }

private:
    UT_StringArray myWriteQueue;
    exint myIndex = 0;
    std::size_t myBytesWritten = 0;
    AsyncStream& myStream;
    UT_WeakPtr<NET_PacketSocket::Impl> myImpl;
};

template <typename DynamicBuffer, typename Handler>
HBOOST_ASIO_INITFN_RESULT_TYPE(
        Handler,
        void(hboost::system::error_code, std::size_t))
NET_PacketSocket::asyncRead(DynamicBuffer& buffer, Handler&& handler)
{
    static_assert(hboost::asio::is_dynamic_buffer<DynamicBuffer>::value);
    return hboost::asio::async_compose<
            Handler,
            void(hboost::system::error_code, std::size_t)>(
            AsyncReadImpl{myImpl->stream(), myImpl, buffer}, handler,
            myImpl->stream());
}

template <typename Handler>
HBOOST_ASIO_INITFN_RESULT_TYPE(
        Handler,
        void(hboost::system::error_code, std::size_t))
NET_PacketSocket::asyncWrite(UT_StringArray&& queue, Handler&& handler)
{
    // Convert each of the messages to the packet format so that we can
    // send all of them in one go instead of having to send each packet
    // one at a time.
    for (UT_StringHolder& msg : queue)
    {
        UT_WorkBuffer wbuf;
        // Always include the null terminator. I have no clue why this is needed
        // but there are some checks that specifically include the null
        // terminator.
        uint32 orig_len = static_cast<uint32>(msg.length() + 1);
        uint32 len = orig_len;
        UTtomips(len);
        // TODO: add support for compressed responses
        wbuf.append((char*)&len, sizeof(uint32));
        wbuf.append('\0');
        wbuf.append(msg);
        // Make sure to always include the null character.
        wbuf.append('\0');
        msg = std::move(wbuf);
    }

    return hboost::asio::async_compose<
        Handler, void(hboost::system::error_code, std::size_t)>(
        AsyncWriteImpl{myImpl->stream(), myImpl, std::move(queue)},
        handler, myImpl->stream()
    );
}

#endif // __NET_PACKETSOCKET_H__