NET_WebSocket.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_WebSocket.h
 *
 * COMMENTS:
 *  class used to handle websockets
 */

#ifndef __NET_WEBSOCKET_H__
#define __NET_WEBSOCKET_H__

#include "NET_API.h"

#include "NET_CircularBuffer.h"
#include "NET_Time.h"
#include "NET_WebAPIResponse.h"
#include "NET_WebStream.h"

#include <UT/UT_StringHolder.h>
#include <UT/UT_StringArray.h>
#include <UT/UT_Url.h>
#include <UT/UT_SharedPtr.h>
#include <UT/UT_NetPacket.h>
#include <UT/UT_Lock.h>
#include <UT/UT_Map.h>
#include <UT/UT_ConcurrentQueue.h>
#include <UT/UT_FastRandom.h>

#include <SYS/SYS_AtomicInt.h>

#include <hboost/uuid/uuid.hpp>
#include <hboost/uuid/uuid_generators.hpp>
#include <hboost/uuid/uuid_io.hpp>
#include <hboost/any.hpp>

#include <future>

template <typename T>
using NET_Promise = std::promise<T>;
template <typename T>
using NET_Future = std::future<T>;

/// Class used for web socket operations.
class NET_API NET_WebSocket : public UTenable_shared_from_this<NET_WebSocket>
{
private:
    using OpCode = char;
    static constexpr OpCode OpCode_Continue = 0x0;
    static constexpr OpCode OpCode_Utf8 = 0x01;
    static constexpr OpCode OpCode_Binary = 0x02;
    static constexpr OpCode OpCode_Terminate = 0x08;
    static constexpr OpCode OpCode_Ping = 0x09;
    static constexpr OpCode OpCode_Pong = 0x0a;

public:
    NET_WebSocket();
    virtual ~NET_WebSocket() = default;

    using Mask = uint32;
    /// @brief The currently supported different frame types.
    enum FrameType : OpCode
    {
        Continue = OpCode_Continue,
        Utf8 = OpCode_Utf8,
        Binary = OpCode_Binary,
        Terminate = OpCode_Terminate,
        Ping = OpCode_Ping,
        Pong = OpCode_Pong
    };

    /// Class used to read a websocket frame. This class solves partial reads
    /// and handles bad frames as well.
    class NET_API Buffer
    {
    public:
        Buffer();

        enum class ParseResult
        {
            Error,
            Incomplete,
            Timeout,
            Complete
        };

        ParseResult read(
                NET_WebStream* stream,
                const std::chrono::milliseconds& timeout,
                bool& finished,
                bool& is_client_closed);
        ParseResult readFrame(
                NET_WebStream* stream,
                const std::chrono::milliseconds& timeout,
                bool &finished,
                bool &is_client_closed);

        FrameType frameType() const { return myMessageType; }

        void reset();

        UT_StringHolder data() const
        {
            return UT_StringHolder(myData.data(), myData.entries());
        }

    private:
        enum class ParseState
        {
            UNKNOWN,
            HEADER_1,
            HEADER_2,
            PAYLOAD_EXT_ONE,
            PAYLOAD_EXT_TWO,
            PAYLOAD_MASK_KEY,
            PAYLOAD
        };

        ParseResult parse_(UT_StringView& chunk);
        ParseResult readAndParse_(
                NET_WebStream* stream,
                const std::chrono::milliseconds& timeout,
                bool& is_client_closed);
        void resetForNextFrame_();

        UT_Array<char> myData;
        exint myFrameLength;
        exint myFrameIndex;
        /// @brief The frame type for the full message.
        NET_WebSocket::FrameType myMessageType;
        /// @brief The frame type for this current frame.
        NET_WebSocket::FrameType myFrameType;
        unsigned myIsMasked : 1, myIsFinished : 1;
        ParseState myCurrentState;
        uint8 myMask[4];
    };

    /// @brief All possible error codes that could occur to be sent to the
    /// onError callback.
    enum ErrorCode
    {
        READ_FRAME
    };

    class NET_API Event
    {
    public:
        enum Type
        {
            NONE,
            GENERAL_ERROR,
            PROCESSING,
            MESSAGE,
            MESSAGE_SENT,
            READ_ERROR,
            SEND_ERROR,
            CLOSE,
            PING,
            PONG,
            CONNECTED,
            CONNECT_ERROR
        };
        Type myType;

        UT_StringHolder myMessage;
        bool myIsBinary;
    };

    enum State
    {
        CONNECTING,
        OPEN,
        CLOSING,
        CLOSED
    };

    /// Read an entire message in. Really only useful for testing purposes.
    static bool readAll(NET_WebStream *stream,
                        UT_WorkBuffer &msg,
                        FrameType &type);
    /// Write a websocket message to the stream.
    static bool write(NET_WebStream *stream,
                      FrameType type,
                      const UT_StringRef &message,
                      Mask mask = 0,
                      exint max_frame_size = -1);
    /// Package a frame which is suitable for sending through a web stream.
    static bool makeFrame(UT_WorkBuffer &frame,
                          FrameType,
                          const UT_StringView mview,
                          Mask mask = 0,
                          bool is_finished = true);
    /// Package up the message into the full message suitable for sending on the
    /// socket.
    static bool makeMessage(
            UT_WorkBuffer& full_message,
            FrameType type,
            const UT_StringRef& message,
            Mask mask,
            exint max_fram_size);

    template <typename SocketT = NET_WebSocket, typename... Args>
    static UT_SharedPtr<SocketT> connect(
            const UT_Url& url,
            const std::chrono::milliseconds& timeout,
            Args&&... args)
    {
        UT_SharedPtr<SocketT> ws
                = UTmakeShared<SocketT>(std::forward<Args>(args)...);
        ws->reconnect(url, timeout);

        return ws;
    }

    /// @brief Convert the frame type to string representation of this frame type.
    ///
    /// @param type The frame type to convert.
    ///
    /// @return The string representation of the frame type.
    static constexpr UT_StringLit frameTypeToString(FrameType type)
    {
        switch(type)
        {
        case Continue:
            return UT_StringLit("Continue");
        case Utf8:
            return UT_StringLit("Utf8");
        case Binary:
            return UT_StringLit("Binary");
        case Ping:
            return UT_StringLit("Ping");
        case Pong:
            return UT_StringLit("Pong");
        case Terminate:
            return UT_StringLit("Terminate");
        default:
            return UT_StringLit();
        };
        return UT_StringLit();
    }

    // +++++++
    // NB: The ability to change this is STRICTLY ONLY for test purposes and is
    // not meant to be changed in product under and circumstances!!!
    // +++++++
    static uint64 theMax16FrameSize;

    /// @brief The file descriptor of this websocket.
    ///
    /// @return The file descriptor
    int getFD() const
    {
        if (!myStream)
            return -1;
        return myStream->getFD();
    }

    /// @brief The url this websocket is using.
    ///
    /// @return The url object representing the used url.
    const UT_Url& url() const { return myUrl; }

    /// @brief Send a message to the other end of the websocket. This function
    ///         will send the data immediately and does NOT queue the data.
    ///         This is a blocking function.
    ///
    /// @param data
    bool send(const UT_StringHolder &data)
    {
        UT_ASSERT(isConnected());

        return internalSend(NET_WebSocket::Utf8, data, myMaxFrameSize);
    }
    bool reconnect(const UT_Url& url, const std::chrono::milliseconds& timeout);
    virtual bool disconnect();

    /// @brief Send a ping to the server. Regular use does not need to worry
    /// about this. Other abstractions will do this for you
    void sendPing(const UT_StringRef& ping = UT_StringHolder::theEmptyString)
    {
        UT_ASSERT(isConnected());
        internalSend(NET_WebSocket::Ping, ping, myMaxFrameSize);
        myIsWaitingForPong = true;
    }
    void sendPong(const UT_StringRef& pong = UT_StringHolder::theEmptyString)
    {
        UT_ASSERT(isConnected());

        internalSend(NET_WebSocket::Pong, pong, myMaxFrameSize);
    }

    /// @brief Tell if this websocket is currently in use
    ///
    /// @return True if this websocket is still in use.
    bool isConnected() const { return myState == State::OPEN; }
    bool usableStream() const { return myStream && myStream->isValid(); }

    /// @brief The current max frame size used.
    ///
    /// @return The max frame size
    uint32 maxFrameSize() const { return myMaxFrameSize; }
    /// @brief Set the max number of bytes sent in each frame
    ///
    /// @param size The new max frame size
    void setMaxFrameSize(uint32 size) { myMaxFrameSize = size; }

    /// @brief The current amount of bytes waiting to be sent over the network.
    ///
    /// @return The number of bytes waiting to be sent.
    exint bufferedAmount() const
    {
        UT_AutoLock lock(myMessageLock);
        return myMessageQueue.size();
    }

    bool hasData();

    virtual void onEvent(const Event& ev) {}

protected:
    friend class NET_WebSocketMonitor;

    static bool setFrameType(FrameType& type, OpCode op_code)
    {
        switch (op_code)
        {
        case OpCode_Continue:
            type = FrameType::Continue;
            return true;
        case OpCode_Utf8:
            type = FrameType::Utf8;
            return true;
        case OpCode_Binary:
            type = FrameType::Binary;
            return true;
        case OpCode_Terminate:
            type = FrameType::Terminate;
            return true;
        case OpCode_Ping:
            type = FrameType::Ping;
            return true;
        case OpCode_Pong:
            type = FrameType::Pong;
            return true;
        default:
            return false;
        };
        return false;
    }

    /// @brief Create the entire message and place it on the queue to send.
    ///
    /// @param type The type of message.
    /// @param message The message to include.
    /// @param mask The mask for the message.
    /// @param max_frame_size The max frame size for the message.
    ///
    /// @return True if creating the message and buffering did not fail.
    bool internalSend(
            FrameType type,
            const UT_StringRef& message,
            exint max_frame_size);

    /// @brief Flush the currently queue messages to send over the network.
    void flush_();

    /// @brief Read a single frame from the server. If you need to read the
    /// entire message then loop until this function returns false or
    /// is_finished is true. This function will call the callbacks when a full
    /// message is received.
    ///
    /// @return True if the read frame was a success.
    Event read_(
            const std::chrono::milliseconds& timeout
            = std::chrono::milliseconds(500));

    UT_Url myUrl;

    NET_WebStreamPtr myStream;
    uint32 myMaxFrameSize;

    bool myIsWaitingForPong;

    NET_Time myLastHeartbeat;

    Buffer myReadBuffer;

    SYS_AtomicInt64 myBufferedAmount;

    mutable UT_Lock myMessageLock;
    NET_CircularBuffer<char> myMessageQueue;

    State myState;

    UT_FastRandom myRand;
};

class NET_API NET_WebSocketAPI : public NET_WebSocket
{
public:

    using id_t = hboost::uuids::uuid;

    void onEvent(const Event& ev) override;

    NET_Future<NET_WebAPIResponse> addPromise(id_t id)
    {
        NET_Promise<NET_WebAPIResponse> p;
        NET_Future<NET_WebAPIResponse> f = p.get_future();
        UT_Lock::Scope scope(myLock);
        myPromises.emplace(id, std::move(p));
        return f;
    }
    static hboost::uuids::uuid generateId()
    {
        return hboost::uuids::random_generator()();
    }

private:
    UT_Lock myLock;
    UT_Map<id_t, NET_Promise<NET_WebAPIResponse>> myPromises;
};

using NET_WebSocketSPtr = UT_SharedPtr<NET_WebSocket>;

#endif // __NET_WEBSOCKET_H__