stream.hpp

#pragma once
 
#include "../tcp/stream.hpp"
#include "../type_traits.hpp"
 
#include <boost/beast/core.hpp>
#include <boost/asio/connect.hpp>
#include <boost/asio/ip/tcp.hpp>
#include <boost/beast/websocket/stream.hpp>
#include <boost/beast/websocket/error.hpp>
 
#if MALLOY_FEATURE_TLS
    #include <boost/beast/ssl.hpp>
    #include <boost/asio/ssl/stream.hpp>
#endif
 
#include <variant>
 
namespace malloy::websocket
{
 
    namespace detail
    {
 
#if MALLOY_FEATURE_TLS
        using tls_stream = boost::beast::websocket::stream<
            boost::beast::ssl_stream<malloy::tcp::stream<>>
        >;
#endif
        using websocket_t = std::variant<
#if MALLOY_FEATURE_TLS
            tls_stream,
#endif
            boost::beast::websocket::stream<malloy::tcp::stream<>>
        >;
        template<typename T>
        concept rw_completion_token = boost::asio::completion_token_for<T, void(malloy::error_code, std::size_t)>;
    }

    class stream
    {
        using ws_t = detail::websocket_t;
 
    public:
        explicit
        stream(detail::websocket_t&& ws) :
            m_underlying_conn{ std::move(ws) }
        {
        }
 
        explicit
        stream(boost::beast::websocket::stream<malloy::tcp::stream<>>&& s) :
            m_underlying_conn{ std::move(s) }
        {
        }
 
        explicit
        stream(malloy::tcp::stream<>&& from) :
            stream{boost::beast::websocket::stream<malloy::tcp::stream<>>{std::move(from)}}
        {
        }
 
#if MALLOY_FEATURE_TLS
        explicit
        stream(detail::tls_stream&& ws) :
            m_underlying_conn{std::move(ws)}
        {
        }
 
        explicit
        stream(boost::beast::ssl_stream<malloy::tcp::stream<>>&& from) :
            stream{malloy::websocket::detail::tls_stream{
                boost::beast::websocket::stream<
                    boost::beast::ssl_stream<malloy::tcp::stream<>>
                >{std::move(from)}}}
        {
        }
#endif
 
        template<concepts::const_buffer_sequence Buff, detail::rw_completion_token Callback>
        auto
        async_write(const Buff& buffers, Callback&& done)
        {
            return std::visit(
                [&buffers, done = std::forward<Callback>(done)](auto& stream) mutable {
                    return stream.async_write(buffers, std::forward<Callback>(done));
                },
                m_underlying_conn
            );
        }
 
        template<concepts::const_buffer_sequence Buff>
        auto
        write(const Buff& buffers) -> std::size_t
        {
            return std::visit(
                [&buffers](auto& stream) mutable {
                    return stream.write(buffers);
                },
                m_underlying_conn
            );
        }
 
        template<concepts::dynamic_buffer Buff, detail::rw_completion_token Callback>
        auto
        async_read(Buff& buff, Callback&& done)
        {
            return std::visit(
                [&buff, done = std::forward<Callback>(done)](auto& s) mutable {
                    return s.async_read(buff, std::forward<Callback>(done));
                },
                m_underlying_conn
            );
        }
 
        template<concepts::dynamic_buffer Buff>
        auto
        read(Buff& buff, boost::beast::error_code& ec) -> std::size_t
        {
            return std::visit(
                [&buff, &ec](auto& s) mutable {
                    return s.read(buff, ec);
                },
                m_underlying_conn
            );
        }
 
        template<boost::asio::completion_token_for<void(malloy::error_code)> CompletionToken>
        auto
        async_close(boost::beast::websocket::close_reason why, CompletionToken&& done)
        {
            return std::visit(
                [why, done = std::forward<CompletionToken>(done)](auto& s) mutable {
                    return s.async_close(why, std::forward<CompletionToken>(done));
                },
                m_underlying_conn
            );
        }
 
        void
        set_option(auto&& opt)
        {
            std::visit(
                [opt = std::forward<decltype(opt)>(opt)](auto& s) mutable {
                    s.set_option(std::forward<decltype(opt)>(opt));
                },
                m_underlying_conn
            );
        }
 
        template<typename Body, typename Fields, boost::asio::completion_token_for<void(malloy::error_code)> CompletionHandler>
        auto
        async_accept(const boost::beast::http::request<Body, Fields>& req, CompletionHandler&& done)
        {
            return std::visit(
                [req, done = std::forward<decltype(done)>(done)](auto& s) mutable {
                    return s.async_accept(req, std::forward<decltype(done)>(done));
                },
                m_underlying_conn
            );
        }
 
        template<typename Body, typename Fields>
        auto
        accept(const boost::beast::http::request<Body, Fields>& req) -> boost::beast::error_code
        {
            return std::visit(
                [req](auto& s) {
                    return s.accept(req);
                },
                m_underlying_conn
            );
        }
 
        template<boost::asio::completion_token_for<void(malloy::error_code)> Callback>
        auto
        async_handshake(std::string host, std::string target, Callback&& done)
        {
            return std::visit(
                [host = std::move(host), target = std::move(target), done = std::forward<Callback>(done)](auto& s) mutable {
                    return s.async_handshake(host, target, std::forward<Callback>(done));
                },
                m_underlying_conn
            );
        }

        void
        set_binary(const bool enabled)
        {
            std::visit(
                [enabled](auto& s) mutable {
                    s.binary(enabled);
                },
                m_underlying_conn
            );
        }

        [[nodiscard]]
        bool
        binary() const
        {
            return std::visit(
                [](auto& s) {
                    return s.binary();
                },
                m_underlying_conn
            );
        }

        template<typename Func>
        void
        get_lowest_layer(Func&& visitor)
        {
            std::visit(
                [visitor = std::forward<Func>(visitor)](auto& s) mutable {
                    visitor(boost::beast::get_lowest_layer(s));
                },
                m_underlying_conn
            );
        }

        auto
        get_executor()
        {
            return std::visit(
                [](auto& s) {
                    return s.get_executor();
                },
                m_underlying_conn
            );
        }

        bool
        is_open() const
        {
            return std::visit(
                [](auto& s){
                    return s.is_open();
                },
                m_underlying_conn
            );
        }

        constexpr
        bool
        is_tls() const
        {
#if MALLOY_FEATURE_TLS
            return std::holds_alternative<detail::tls_stream>(m_underlying_conn);
#else 
            return false;
#endif
        }
 
#if MALLOY_FEATURE_TLS
        template<concepts::accept_handler Callback>
        void
        async_handshake_tls(boost::asio::ssl::stream_base::handshake_type type, Callback&& done)
        {
            if (!is_tls())
                throw std::logic_error{"async_handshake_tls called on non-tls stream"};
 
            std::visit(
                [done = std::forward<Callback>(done), type](auto& s) mutable {
                    if constexpr (std::same_as<std::decay_t<decltype(s)>, detail::tls_stream>)
                        s.next_layer().async_handshake(type, std::forward<Callback>(done));
                },
                m_underlying_conn
            );
        }
        #endif
 
    private:
        ws_t m_underlying_conn;
    };
 
}