ws_endpoint.hpp

Go to the documentation of this file.

// Copyright Takatoshi Kondo 2016
//
// Distributed under the Boost Software License, Version 1.0.
// (See accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)
 
#if !defined(MQTT_WS_ENDPOINT_HPP)
#define MQTT_WS_ENDPOINT_HPP
 
#include <boost/beast/websocket.hpp>
#include <boost/beast/core/flat_buffer.hpp>
#include <boost/asio/bind_executor.hpp>
 
#include <mqtt/namespace.hpp>
#include <mqtt/type_erased_socket.hpp>
#include <mqtt/move.hpp>
#include <mqtt/attributes.hpp>
#include <mqtt/string_view.hpp>
#include <mqtt/error_code.hpp>
 
namespace MQTT_NS {
 
namespace as = boost::asio;
 
template <typename Socket, typename Strand>
class ws_endpoint : public socket {
public:
    template <typename... Args>
    ws_endpoint(as::io_context& ioc, Args&&... args)
        :ws_(ioc, std::forward<Args>(args)...),
         strand_(ioc) {
        ws_.binary(true);
        ws_.set_option(
            boost::beast::websocket::stream_base::decorator(
                [](boost::beast::websocket::request_type& req) {
                    req.set("Sec-WebSocket-Protocol", "mqtt");
                }
            )
        );
    }
 
    MQTT_ALWAYS_INLINE void async_read(
        as::mutable_buffer buffers,
        std::function<void(error_code, std::size_t)> handler
    ) override final {
        auto req_size = as::buffer_size(buffers);
 
        using beast_read_handler_t =
            std::function<void(error_code ec, std::shared_ptr<void>)>;
 
        std::shared_ptr<beast_read_handler_t> beast_read_handler;
        if (req_size <= buffer_.size()) {
            as::buffer_copy(buffers, buffer_.data(), req_size);
            buffer_.consume(req_size);
            handler(boost::system::errc::make_error_code(boost::system::errc::success), req_size);
            return;
        }
 
        beast_read_handler.reset(
            new beast_read_handler_t(
                [this, req_size, buffers, handler = force_move(handler)]
                (error_code ec, std::shared_ptr<void> const& v) mutable {
                    if (ec) {
                        force_move(handler)(ec, 0);
                        return;
                    }
                    if (!ws_.got_binary()) {
                        buffer_.consume(buffer_.size());
                        force_move(handler)
                            (boost::system::errc::make_error_code(boost::system::errc::bad_message), 0);
                        return;
                    }
                    if (req_size > buffer_.size()) {
                        auto beast_read_handler = std::static_pointer_cast<beast_read_handler_t>(v);
                        ws_.async_read(
                            buffer_,
                            as::bind_executor(
                                strand_,
                                [beast_read_handler]
                                (error_code ec, std::size_t) {
                                    (*beast_read_handler)(ec, beast_read_handler);
                                }
                            )
                        );
                        return;
                    }
                    as::buffer_copy(buffers, buffer_.data(), req_size);
                    buffer_.consume(req_size);
                    force_move(handler)(boost::system::errc::make_error_code(boost::system::errc::success), req_size);
                }
            )
        );
        ws_.async_read(
            buffer_,
            as::bind_executor(
                strand_,
                [beast_read_handler]
                (error_code ec, std::size_t) {
                    (*beast_read_handler)(ec, beast_read_handler);
                }
            )
        );
    }
 
    MQTT_ALWAYS_INLINE void async_write(
        std::vector<as::const_buffer> buffers,
        std::function<void(error_code, std::size_t)> handler
    ) override final {
        ws_.async_write(
            buffers,
            as::bind_executor(
                strand_,
                force_move(handler)
            )
        );
    }
 
    MQTT_ALWAYS_INLINE std::size_t write(
        std::vector<as::const_buffer> buffers,
        boost::system::error_code& ec
    ) override final {
        ws_.write(buffers, ec);
        return as::buffer_size(buffers);
    }
 
    MQTT_ALWAYS_INLINE void post(std::function<void()> handler) override final {
        as::post(
            strand_,
            force_move(handler)
        );
    }
 
#if BOOST_VERSION >= 107000
 
    MQTT_ALWAYS_INLINE as::ip::tcp::socket::lowest_layer_type& lowest_layer() override final {
        return boost::beast::get_lowest_layer(ws_);
    }
 
#else  // BOOST_VERSION >= 107000
 
    MQTT_ALWAYS_INLINE as::ip::tcp::socket::lowest_layer_type& lowest_layer() override final {
        return ws_.lowest_layer();
    }
 
#endif // BOOST_VERSION >= 107000
 
    MQTT_ALWAYS_INLINE any native_handle() override final {
        return next_layer().native_handle();
    }
 
    MQTT_ALWAYS_INLINE void close(boost::system::error_code& ec) override final {
        ws_.close(boost::beast::websocket::close_code::normal, ec);
        if (ec) return;
        do {
            boost::beast::flat_buffer buffer;
            ws_.read(buffer, ec);
        } while (!ec);
        if (ec != boost::beast::websocket::error::closed) return;
        ec = boost::system::errc::make_error_code(boost::system::errc::success);
    }
 
#if BOOST_VERSION < 107400 || defined(BOOST_ASIO_USE_TS_EXECUTOR_AS_DEFAULT)
    MQTT_ALWAYS_INLINE as::executor get_executor() override final {
        return lowest_layer().get_executor();
    }
#else  // BOOST_VERSION < 107400 || defined(BOOST_ASIO_USE_TS_EXECUTOR_AS_DEFAULT)
    MQTT_ALWAYS_INLINE as::any_io_executor get_executor() override final {
        return lowest_layer().get_executor();
    }
#endif // BOOST_VERSION < 107400 || defined(BOOST_ASIO_USE_TS_EXECUTOR_AS_DEFAULT)
 
    typename boost::beast::websocket::stream<Socket>::next_layer_type& next_layer() {
        return ws_.next_layer();
    }
 
    template <typename T>
    void set_option(T&& t) {
        ws_.set_option(std::forward<T>(t));
    }
 
    template <typename ConstBufferSequence, typename AcceptHandler>
    void async_accept(
        ConstBufferSequence const& buffers,
        AcceptHandler&& handler) {
        ws_.async_accept(buffers, std::forward<AcceptHandler>(handler));
    }
 
    template<typename ConstBufferSequence, typename ResponseDecorator, typename AcceptHandler>
    void async_accept_ex(
        ConstBufferSequence const& buffers,
        ResponseDecorator const& decorator,
        AcceptHandler&& handler) {
        ws_.async_accept_ex(buffers, decorator, std::forward<AcceptHandler>(handler));
    }
 
    template <typename... Args>
    void async_handshake(Args&& ... args) {
        ws_.async_handshake(std::forward<Args>(args)...);
    }
 
    template <typename... Args>
    void handshake(Args&& ... args) {
        ws_.handshake(std::forward<Args>(args)...);
    }
 
    template <typename ConstBufferSequence>
    std::size_t write(
        ConstBufferSequence const& buffers) {
        ws_.write(buffers);
        return as::buffer_size(buffers);
    }
 
private:
    boost::beast::websocket::stream<Socket> ws_;
    boost::beast::flat_buffer buffer_;
    Strand strand_;
};
 
} // namespace MQTT_NS
 
#endif // MQTT_WS_ENDPOINT_HPP