Skip to content

Instantly share code, notes, and snippets.

@microcai
Created September 26, 2018 09:23
Show Gist options
  • Save microcai/d8fe945e16ae6d4bd1f629cf4d87787e to your computer and use it in GitHub Desktop.
Save microcai/d8fe945e16ae6d4bd1f629cf4d87787e to your computer and use it in GitHub Desktop.
#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);
}
};
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment