PlainWsSession.hpp

#pragma once
 
#include "util/Taggable.hpp"
#include "web/dosguard/DOSGuardInterface.hpp"
#include "web/impl/WsBase.hpp"
#include "web/interface/ConnectionBase.hpp"
 
#include <boost/asio/ip/tcp.hpp>
#include <boost/beast/core/flat_buffer.hpp>
#include <boost/beast/core/stream_traits.hpp>
#include <boost/beast/core/tcp_stream.hpp>
#include <boost/beast/http/message.hpp>
#include <boost/beast/http/parser.hpp>
#include <boost/beast/http/string_body.hpp>
#include <boost/beast/websocket/stream.hpp>
#include <boost/optional/optional.hpp>
 
#include <chrono>
#include <functional>
#include <memory>
#include <string>
#include <utility>
 
namespace web {

template <SomeServerHandler HandlerType>
class PlainWsSession : public impl::WsBase<PlainWsSession, HandlerType> {
    using StreamType = boost::beast::websocket::stream<boost::beast::tcp_stream>;
    StreamType ws_;
 
public:
    explicit PlainWsSession(
        boost::asio::ip::tcp::socket&& socket,
        std::string ip,
        std::reference_wrapper<util::TagDecoratorFactory const> tagFactory,
        std::reference_wrapper<dosguard::DOSGuardInterface> dosGuard,
        std::shared_ptr<HandlerType> const& handler,
        boost::beast::flat_buffer&& buffer,
        bool isAdmin,
        std::uint32_t maxSendingQueueSize
    )
        : impl::WsBase<PlainWsSession, HandlerType>(
              ip,
              tagFactory,
              dosGuard,
              handler,
              std::move(buffer),
              maxSendingQueueSize
          )
        , ws_(std::move(socket))
    {
        ConnectionBase::isAdmin_ = isAdmin;  // NOLINT(cppcoreguidelines-prefer-member-initializer)
    }
 
    ~PlainWsSession() override = default;

    StreamType&
    ws()
    {
        return ws_;
    }
};

template <SomeServerHandler HandlerType>
class WsUpgrader : public std::enable_shared_from_this<WsUpgrader<HandlerType>> {
    using std::enable_shared_from_this<WsUpgrader<HandlerType>>::shared_from_this;
 
    boost::beast::tcp_stream http_;
    boost::optional<http::request_parser<http::string_body>> parser_;
    boost::beast::flat_buffer buffer_;
    std::reference_wrapper<util::TagDecoratorFactory const> tagFactory_;
    std::reference_wrapper<dosguard::DOSGuardInterface> dosGuard_;
    http::request<http::string_body> req_;
    std::string ip_;
    std::shared_ptr<HandlerType> const handler_;
    bool isAdmin_;
    std::uint32_t maxWsSendingQueueSize_;
 
public:
    WsUpgrader(
        boost::beast::tcp_stream&& stream,
        std::string ip,
        std::reference_wrapper<util::TagDecoratorFactory const> tagFactory,
        std::reference_wrapper<dosguard::DOSGuardInterface> dosGuard,
        std::shared_ptr<HandlerType> const& handler,
        boost::beast::flat_buffer&& buffer,
        http::request<http::string_body> request,
        bool isAdmin,
        std::uint32_t maxWsSendingQueueSize
    )
        : http_(std::move(stream))
        , buffer_(std::move(buffer))
        , tagFactory_(tagFactory)
        , dosGuard_(dosGuard)
        , req_(std::move(request))
        , ip_(std::move(ip))
        , handler_(handler)
        , isAdmin_(isAdmin)
        , maxWsSendingQueueSize_(maxWsSendingQueueSize)
    {
    }

    void
    run()
    {
        boost::asio::dispatch(
            http_.get_executor(),
            boost::beast::bind_front_handler(
                &WsUpgrader<HandlerType>::doUpgrade, shared_from_this()
            )
        );
    }
 
private:
    void
    doUpgrade()
    {
        parser_.emplace();
 
        static constexpr auto kMAX_BODY_SIZE = 10000;
        parser_->body_limit(kMAX_BODY_SIZE);
 
        boost::beast::get_lowest_layer(http_).expires_after(std::chrono::seconds(30));
        onUpgrade();
    }
 
    void
    onUpgrade()
    {
        if (!boost::beast::websocket::is_upgrade(req_))
            return;
 
        // Disable the timeout. The websocket::stream uses its own timeout settings.
        boost::beast::get_lowest_layer(http_).expires_never();
 
        std::make_shared<PlainWsSession<HandlerType>>(
            http_.release_socket(),
            ip_,
            tagFactory_,
            dosGuard_,
            handler_,
            std::move(buffer_),
            isAdmin_,
            maxWsSendingQueueSize_
        )
            ->run(std::move(req_));
    }
};
 
}  // namespace web