websocket_stream.hpp

#pragma once

#include <boost/beast/websocket.hpp>

#include <boost/beast/websocket/ssl.hpp>

struct websocket_stream : boost::noncopyable
{
	using wss_socket_t = boost::beast::websocket::stream<boost::asio::ssl::stream<boost::asio::ip::tcp::socket&>>;
	using ws_socket_t = boost::beast::websocket::stream<boost::asio::ip::tcp::socket&>;

	typedef ws_socket_t* ws_socket_ptr;
	typedef wss_socket_t* wss_socket_ptr;

	boost::asio::ssl::context ssl_ctx;
	boost::asio::ip::tcp::socket tcp_socket;

	std::variant<wss_socket_ptr, ws_socket_ptr> ws_socket;

	char real_socket_object_storage[std::max(sizeof(ws_socket_t), sizeof(wss_socket_t)) ];

	~websocket_stream()
	{
		std::visit([](auto&& _ws){ _ws->~stream(); }, ws_socket);
	}

	websocket_stream(boost::asio::io_context& io)
		: ssl_ctx(boost::asio::ssl::context::tls_client)
		, tcp_socket(io)
	{
	}

	websocket_stream(boost::asio::io_context& io, bool use_ssl)
		: ssl_ctx(boost::asio::ssl::context::tls_client)
		, tcp_socket(io)
	{
		recreate_ws(use_ssl);
	}

	void recreate_ws(bool use_ssl)
	{
		if (use_ssl)
		{
			static_assert(sizeof(real_socket_object_storage) >= sizeof(wss_socket_t), "allocate more storage please" );
			ws_socket = new (real_socket_object_storage) wss_socket_t(tcp_socket, ssl_ctx);
		}
		else
		{
			static_assert(sizeof(real_socket_object_storage) >= sizeof(ws_socket_t), "allocate more storage please" );
			ws_socket = new (real_socket_object_storage) ws_socket_t(tcp_socket);
		}
	}

	void recreate_ws()
	{
		std::visit([this](auto && _ws_socket)
		{
			_ws_socket->~stream();

			if constexpr (std::is_same_v<std::decay_t<decltype(_ws_socket)>, wss_socket_ptr>)
			{
				ws_socket = new (real_socket_object_storage) wss_socket_t(tcp_socket, ssl_ctx);
			}
			else
			{
				ws_socket = new (real_socket_object_storage) ws_socket_t(tcp_socket);
			}
		}, ws_socket);
	}

	template <typename MutableBufferSequence, typename ReadHandler>
	BOOST_ASIO_INITFN_RESULT_TYPE(ReadHandler,
		void (boost::system::error_code, std::size_t))
	async_read_some(const MutableBufferSequence& buffers,
		BOOST_ASIO_MOVE_ARG(ReadHandler) handler)
	{
		return std::visit([&](auto && stream){ return stream->async_read_some(buffers, handler); }, ws_socket);
	}

	template <typename ConstBufferSequence, typename WriteHandler>
	BOOST_ASIO_INITFN_RESULT_TYPE(WriteHandler,
		void (boost::system::error_code, std::size_t))
	async_write(const ConstBufferSequence& buffers,
		BOOST_ASIO_MOVE_ARG(WriteHandler) handler)
	{
		return std::visit([&](auto && stream){ return stream->async_write(buffers, handler); }, ws_socket);
	}

	template<class HandshakeHandler>
	BOOST_ASIO_INITFN_RESULT_TYPE(HandshakeHandler,
		void(boost::system::error_code))
	async_handshake(boost::string_view host, boost::string_view target, BOOST_ASIO_MOVE_ARG(HandshakeHandler) handler)
	{
		return std::visit([&](auto&& _ws_socket){ return _ws_socket->async_handshake(host, target, handler); }, ws_socket);
	}

	bool is_message_done() const
	{
		return std::visit([](auto&& ws){ return ws->is_message_done();}, ws_socket);
	}

    template<class CloseHandler>
    BOOST_ASIO_INITFN_RESULT_TYPE(
        CloseHandler, void(boost::system::error_code))
    async_close(boost::beast::websocket::close_reason const& cr, CloseHandler&& handler)
	{
		return std::visit([=](auto && _ws){ return _ws->async_close(cr, handler); }, ws_socket);
	}
};