AutoSocket.h

#pragma once
 
#include <xrpl/basics/Log.h>
#include <xrpl/beast/net/IPAddressConversion.h>
#include <xrpl/beast/net/IPEndpoint.h>
 
#include <boost/asio.hpp>
#include <boost/asio/ip/tcp.hpp>
#include <boost/asio/ssl.hpp>
#include <boost/beast/core/bind_handler.hpp>
 
#include <cstddef>
#include <functional>
#include <memory>
#include <string>
#include <utility>
#include <vector>
 
// Socket wrapper that supports both SSL and non-SSL connections.
// Generally, handle it as you would an SSL connection.
// To force a non-SSL connection, just don't call async_handshake.
// To force SSL only inbound, call setSSLOnly.
 
class AutoSocket
{
public:
    using SslSocket = boost::asio::ssl::stream<boost::asio::ip::tcp::socket>;
    using EndpointType = boost::asio::ip::tcp::socket::endpoint_type;
    using SocketPtr = std::unique_ptr<SslSocket>;
    using PlainSocket = SslSocket::next_layer_type;
    // NOLINTNEXTLINE(readability-identifier-naming) -- asio stream layer concept name
    using lowest_layer_type = SslSocket::lowest_layer_type;
    using HandshakeType = SslSocket::handshake_type;
    using ErrorCode = boost::system::error_code;
    using Callback = std::function<void(ErrorCode)>;
 
public:
    AutoSocket(
        boost::asio::io_context& s,
        boost::asio::ssl::context& c,
        bool secureOnly,
        bool plainOnly)
        : secure_(secureOnly)
        , buffer_((plainOnly || secureOnly) ? 0 : 4)
        , j_{beast::Journal::getNullSink()}
    {
        socket_ = std::make_unique<SslSocket>(s, c);
    }
 
    AutoSocket(boost::asio::io_context& s, boost::asio::ssl::context& c)
        : AutoSocket(s, c, false, false)
    {
    }
 
    [[nodiscard]] bool
    isSecure() const
    {
        return secure_;
    }
    SslSocket&
    sslSocket()
    {
        return *socket_;
    }
    PlainSocket&
    plainSocket()
    {
        return socket_->next_layer();
    }
 
    beast::ip::Endpoint
    localEndpoint()
    {
        return beast::ip::fromAsio(lowestLayer().local_endpoint());
    }
 
    beast::ip::Endpoint
    remoteEndpoint()
    {
        return beast::ip::fromAsio(lowestLayer().remote_endpoint());
    }
 
    lowest_layer_type&
    lowestLayer()
    {
        return socket_->lowest_layer();
    }
 
    void
    swap(AutoSocket& s) noexcept
    {
        buffer_.swap(s.buffer_);
        socket_.swap(s.socket_);
        std::swap(secure_, s.secure_);
    }
 
    boost::system::error_code
    cancel(boost::system::error_code& ec)
    {
        return lowestLayer().cancel(ec);
    }
 
    void
    asyncHandshake(HandshakeType type, Callback cbFunc)
    {
        if ((type == SslSocket::client) || (secure_))
        {
            // must be ssl
            secure_ = true;
            socket_->async_handshake(type, cbFunc);
        }
        else if (buffer_.empty())
        {
            // must be plain
            secure_ = false;
            post(socket_->get_executor(), boost::beast::bind_handler(cbFunc, ErrorCode()));
        }
        else
        {
            // autodetect
            socket_->next_layer().async_receive(
                boost::asio::buffer(buffer_),
                boost::asio::socket_base::message_peek,
                [this, cbFunc](ErrorCode const& ec, size_t bytesTransferred) {
                    handleAutodetect(cbFunc, ec, bytesTransferred);
                });
        }
    }
 
    template <typename ShutdownHandler>
    void
    asyncShutdown(ShutdownHandler handler)
    {
        if (isSecure())
        {
            socket_->async_shutdown(handler);
        }
        else
        {
            ErrorCode ec;
            try
            {
                lowestLayer().shutdown(PlainSocket::shutdown_both);
            }
            catch (boost::system::system_error const& e)
            {
                ec = e.code();
            }
            post(socket_->get_executor(), boost::beast::bind_handler(handler, ec));
        }
    }
 
    template <typename Seq, typename Handler>
    void
    asyncReadSome(Seq const& buffers, Handler handler)
    {
        if (isSecure())
        {
            socket_->async_read_some(buffers, handler);
        }
        else
        {
            plainSocket().async_read_some(buffers, handler);
        }
    }
 
    template <typename Seq, typename Condition, typename Handler>
    void
    asyncReadUntil(Seq const& buffers, Condition condition, Handler handler)
    {
        if (isSecure())
        {
            boost::asio::async_read_until(*socket_, buffers, condition, handler);
        }
        else
        {
            boost::asio::async_read_until(plainSocket(), buffers, condition, handler);
        }
    }
 
    template <typename Allocator, typename Handler>
    void
    asyncReadUntil(
        boost::asio::basic_streambuf<Allocator>& buffers,
        std::string const& delim,
        Handler handler)
    {
        if (isSecure())
        {
            boost::asio::async_read_until(*socket_, buffers, delim, handler);
        }
        else
        {
            boost::asio::async_read_until(plainSocket(), buffers, delim, handler);
        }
    }
 
    template <typename Allocator, typename MatchCondition, typename Handler>
    void
    asyncReadUntil(
        boost::asio::basic_streambuf<Allocator>& buffers,
        MatchCondition cond,
        Handler handler)
    {
        if (isSecure())
        {
            boost::asio::async_read_until(*socket_, buffers, cond, handler);
        }
        else
        {
            boost::asio::async_read_until(plainSocket(), buffers, cond, handler);
        }
    }
 
    template <typename Buf, typename Handler>
    void
    asyncWrite(Buf const& buffers, Handler handler)
    {
        if (isSecure())
        {
            boost::asio::async_write(*socket_, buffers, handler);
        }
        else
        {
            boost::asio::async_write(plainSocket(), buffers, handler);
        }
    }
 
    template <typename Allocator, typename Handler>
    void
    asyncWrite(boost::asio::basic_streambuf<Allocator>& buffers, Handler handler)
    {
        if (isSecure())
        {
            boost::asio::async_write(*socket_, buffers, handler);
        }
        else
        {
            boost::asio::async_write(plainSocket(), buffers, handler);
        }
    }
 
    template <typename Buf, typename Condition, typename Handler>
    void
    asyncRead(Buf const& buffers, Condition cond, Handler handler)
    {
        if (isSecure())
        {
            boost::asio::async_read(*socket_, buffers, cond, handler);
        }
        else
        {
            boost::asio::async_read(plainSocket(), buffers, cond, handler);
        }
    }
 
    template <typename Allocator, typename Condition, typename Handler>
    void
    asyncRead(boost::asio::basic_streambuf<Allocator>& buffers, Condition cond, Handler handler)
    {
        if (isSecure())
        {
            boost::asio::async_read(*socket_, buffers, cond, handler);
        }
        else
        {
            boost::asio::async_read(plainSocket(), buffers, cond, handler);
        }
    }
 
    template <typename Buf, typename Handler>
    void
    asyncRead(Buf const& buffers, Handler handler)
    {
        if (isSecure())
        {
            boost::asio::async_read(*socket_, buffers, handler);
        }
        else
        {
            boost::asio::async_read(plainSocket(), buffers, handler);
        }
    }
 
    template <typename Seq, typename Handler>
    void
    asyncWriteSome(Seq const& buffers, Handler handler)
    {
        if (isSecure())
        {
            socket_->async_write_some(buffers, handler);
        }
        else
        {
            plainSocket().async_write_some(buffers, handler);
        }
    }
 
protected:
    void
    handleAutodetect(Callback cbFunc, ErrorCode const& ec, size_t bytesTransferred)
    {
        using namespace xrpl;
 
        if (ec)
        {
            JLOG(j_.warn()) << "Handle autodetect error: " << ec;
            cbFunc(ec);
        }
        else if (
            (buffer_[0] < 127) && (buffer_[0] > 31) &&
            ((bytesTransferred < 2) || ((buffer_[1] < 127) && (buffer_[1] > 31))) &&
            ((bytesTransferred < 3) || ((buffer_[2] < 127) && (buffer_[2] > 31))) &&
            ((bytesTransferred < 4) || ((buffer_[3] < 127) && (buffer_[3] > 31))))
        {
            // not ssl
            JLOG(j_.trace()) << "non-SSL";
            secure_ = false;
            cbFunc(ec);
        }
        else
        {
            // ssl
            JLOG(j_.trace()) << "SSL";
            secure_ = true;
            socket_->async_handshake(SslSocket::server, cbFunc);
        }
    }
 
private:
    SocketPtr socket_;
    bool secure_;
    std::vector<char> buffer_;
    beast::Journal j_;
};