WsBase.hpp

#pragma once
 
#include "rpc/Errors.hpp"
#include "rpc/common/Types.hpp"
#include "util/Taggable.hpp"
#include "util/log/Logger.hpp"
#include "web/LoadWarning.hpp"
#include "web/SubscriptionContext.hpp"
#include "web/SubscriptionContextInterface.hpp"
#include "web/dosguard/DOSGuardInterface.hpp"
#include "web/interface/Concepts.hpp"
#include "web/interface/ConnectionBase.hpp"
 
#include <boost/asio/buffer.hpp>
#include <boost/asio/error.hpp>
#include <boost/beast/core.hpp>
#include <boost/beast/core/error.hpp>
#include <boost/beast/core/flat_buffer.hpp>
#include <boost/beast/core/role.hpp>
#include <boost/beast/http/field.hpp>
#include <boost/beast/http/message.hpp>
#include <boost/beast/http/status.hpp>
#include <boost/beast/http/string_body.hpp>
#include <boost/beast/version.hpp>
#include <boost/beast/websocket/error.hpp>
#include <boost/beast/websocket/rfc6455.hpp>
#include <boost/beast/websocket/stream_base.hpp>
#include <boost/core/ignore_unused.hpp>
#include <boost/json/array.hpp>
#include <boost/json/parse.hpp>
#include <boost/json/serialize.hpp>
#include <xrpl/protocol/ErrorCodes.h>
 
#include <cstddef>
#include <cstdint>
#include <exception>
#include <functional>
#include <memory>
#include <queue>
#include <string>
#include <utility>
 
namespace web::impl {

template <template <typename> typename Derived, SomeServerHandler HandlerType>
class WsBase : public ConnectionBase,
               public std::enable_shared_from_this<WsBase<Derived, HandlerType>> {
    using std::enable_shared_from_this<WsBase<Derived, HandlerType>>::shared_from_this;
 
    boost::beast::flat_buffer buffer_;
    std::reference_wrapper<dosguard::DOSGuardInterface> dosGuard_;
    bool sending_ = false;
    std::queue<std::shared_ptr<std::string>> messages_;
    std::shared_ptr<HandlerType> const handler_;
 
    SubscriptionContextPtr subscriptionContext_;
    std::uint32_t maxSendingQueueSize_;
 
protected:
    util::Logger log_{"WebServer"};
    util::Logger perfLog_{"Performance"};
 
    void
    wsFail(boost::beast::error_code ec, char const* what)
    {
        // Don't log if the WebSocket stream was gracefully closed at both endpoints
        if (ec != boost::beast::websocket::error::closed) {
            LOG(log_.error()) << tag() << ": " << what << ": " << ec.message() << ": "
                              << ec.value();
        }
 
        if (!ec_ && ec != boost::asio::error::operation_aborted) {
            ec_ = ec;
            boost::beast::get_lowest_layer(derived().ws()).socket().close(ec);
        }
    }
 
public:
    explicit WsBase(
        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,
        std::uint32_t maxSendingQueueSize
    )
        : ConnectionBase(tagFactory, ip)
        , buffer_(std::move(buffer))
        , dosGuard_(dosGuard)
        , handler_(handler)
        , maxSendingQueueSize_(maxSendingQueueSize)
    {
        upgraded = true;  // NOLINT (cppcoreguidelines-pro-type-member-init)
 
        LOG(perfLog_.debug()) << tag() << "session created";
    }
 
    ~WsBase() override
    {
        LOG(perfLog_.debug()) << tag() << "session closed";
        dosGuard_.get().decrement(clientIp_);
    }
 
    Derived<HandlerType>&
    derived()
    {
        return static_cast<Derived<HandlerType>&>(*this);
    }
 
    void
    doWrite()
    {
        sending_ = true;
        derived().ws().async_write(
            boost::asio::buffer(messages_.front()->data(), messages_.front()->size()),
            boost::beast::bind_front_handler(&WsBase::onWrite, derived().shared_from_this())
        );
    }
 
    void
    onWrite(boost::system::error_code ec, std::size_t)
    {
        messages_.pop();
        sending_ = false;
        if (ec) {
            wsFail(ec, "Failed to write");
        } else {
            maybeSendNext();
        }
    }
 
    void
    maybeSendNext()
    {
        if (ec_ || sending_ || messages_.empty())
            return;
 
        doWrite();
    }
 
    void
    sendSlowDown(std::string const& request) override
    {
        sendError(rpc::XrpldError::RpcSlowDown, request);
    }

    void
    send(std::shared_ptr<std::string> msg) override
    {
        // Note: post used instead of dispatch to guarantee async behavior of wsFail and
        // maybeSendNext
        boost::asio::post(
            derived().ws().get_executor(),
            [this, self = derived().shared_from_this(), msg = std::move(msg)]() {
                if (messages_.size() > maxSendingQueueSize_) {
                    wsFail(boost::asio::error::timed_out, "Client is too slow");
                    return;
                }
 
                messages_.push(msg);
                maybeSendNext();
            }
        );
    }

    SubscriptionContextPtr
    makeSubscriptionContext(util::TagDecoratorFactory const& factory) override
    {
        if (subscriptionContext_ == nullptr) {
            subscriptionContext_ =
                std::make_shared<SubscriptionContext>(factory, shared_from_this());
        }
        return subscriptionContext_;
    }

    void
    send(std::string&& msg, http::status) override
    {
        if (!dosGuard_.get().add(clientIp_, msg.size())) {
            if (auto const warned = withLoadWarning(msg); warned.has_value()) {
                // Reserialize when we need to include this warning
                msg = boost::json::serialize(*warned);
            }
        }
        auto sharedMsg = std::make_shared<std::string>(std::move(msg));
        send(std::move(sharedMsg));
    }

    void
    run(http::request<http::string_body> req)
    {
        using namespace boost::beast;
 
        derived().ws().set_option(websocket::stream_base::timeout::suggested(role_type::server));
 
        // Set a decorator to change the Server of the handshake
        derived().ws().set_option(
            websocket::stream_base::decorator([](websocket::response_type& res) {
                res.set(
                    http::field::server,
                    std::string(BOOST_BEAST_VERSION_STRING) + " websocket-server-async"
                );
            })
        );
 
        derived().ws().async_accept(
            req, bind_front_handler(&WsBase::onAccept, this->shared_from_this())
        );
    }
 
    void
    onAccept(boost::beast::error_code ec)
    {
        if (ec)
            return wsFail(ec, "accept");
 
        LOG(perfLog_.info()) << tag() << "accepting new connection";
 
        doRead();
    }
 
    void
    doRead()
    {
        if (dead())
            return;
 
        // Note: use entirely new buffer so previously used, potentially large, capacity is
        // deallocated
        buffer_ = boost::beast::flat_buffer{};
 
        derived().ws().async_read(
            buffer_, boost::beast::bind_front_handler(&WsBase::onRead, this->shared_from_this())
        );
    }
 
    void
    onRead(boost::beast::error_code ec, std::size_t bytesTransferred)
    {
        boost::ignore_unused(bytesTransferred);
 
        if (ec)
            return wsFail(ec, "read");
 
        LOG(perfLog_.info()) << tag() << "Received request from ip = " << clientIp_;
 
        std::string requestStr{static_cast<char const*>(buffer_.data().data()), buffer_.size()};
 
        try {
            (*handler_)(requestStr, shared_from_this());
        } catch (std::exception const&) {
            sendError(rpc::XrpldError::RpcInternal, std::move(requestStr));
        }
 
        doRead();
    }
 
private:
    void
    sendError(rpc::XrpldError error, std::string requestStr)
    {
        auto e = rpc::makeError(error);
 
        try {
            auto request = boost::json::parse(requestStr);
            if (request.is_object() && request.as_object().contains("id"))
                e["id"] = request.as_object().at("id");
            e["request"] = std::move(request);
        } catch (std::exception const&) {
            e["request"] = requestStr;
        }
 
        this->send(std::make_shared<std::string>(boost::json::serialize(e)));
    }
};
}  // namespace web::impl