websocket_stream.hpp

#pragma once

#include <boost/noncopyable.hpp>

#include <boost/beast/websocket.hpp>

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

struct websocket_stream : boost::noncopyable
{
	websocket_stream(websocket_stream&&) = delete;
	websocket_stream& operator= (websocket_stream&&) = delete;

	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){ if (_ws) _ws->~stream(); }, ws_socket);
		ws_socket = ws_socket_ptr(nullptr);
	}

	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)
	{
		create_ws(use_ssl);
	}

	void create_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);
		}
	}

  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, ReadHandler handler)
	{
		return std::visit([=](auto && _ws_socket){ return _ws_socket->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, WriteHandler handler)
	{
		return std::visit([=](auto && _ws_socket){ return _ws_socket->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, BOOST_ASIO_MOVE_ARG(CloseHandler) handler)
	{
		return std::visit([=](auto && _ws){ return _ws->async_close(cr, handler); }, ws_socket);
	}
};