websocketserverendpoint.h

/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/. */
 
#pragma once
 
// Local Includes
#include "iwebsocketserverendpoint.h"
#include "websocketmessage.h"
#include "websocketserver.h"
 
// External Includes
#include <memory.h>
#include <future>
#include <mutex>
#include <unordered_set>
#include <mathutils.h>
#include <rapidjson/document.h>
#include <rapidjson/istreamwrapper.h>
#include <rapidjson/prettywriter.h>
#include <rapidjson/error/en.h>
#include <nap/logger.h>
 
namespace nap
{
    template<typename config>
    class WebSocketServerEndPointSetup : public IWebSocketServerEndPoint
    {
        RTTI_ENABLE(IWebSocketServerEndPoint)
    public:
        // Stop endpoint on destruction
        ~WebSocketServerEndPointSetup() override;
 
        bool init(utility::ErrorState& errorState) override;
 
        bool start(nap::utility::ErrorState& error) override;
 
        bool isOpen() const override;
 
        void stop() override;
 
        bool send(const WebSocketConnection& connection, const std::string& message, EWebSocketOPCode code, nap::utility::ErrorState& error) override;
 
        bool send(const WebSocketConnection& connection, void const* payload, int length, EWebSocketOPCode code, nap::utility::ErrorState& error) override;
 
        bool broadcast(const std::string& message, EWebSocketOPCode code, nap::utility::ErrorState& error) override;
 
        bool broadcast(void const* payload, int length, EWebSocketOPCode code, nap::utility::ErrorState& error) override;
 
        std::string getHostName(const WebSocketConnection& connection) override;
 
        void getHostNames(std::vector<std::string>& outHosts) override;
 
        int getConnectionCount() override;
 
         bool acceptsNewConnections() override;
 
    protected:
        void registerListener(IWebSocketServer& server) override;
 
        void unregisterListener(IWebSocketServer& server) override;
 
        void run();
 
        void onConnectionOpened(wspp::ConnectionHandle connection);
 
        void onConnectionClosed(wspp::ConnectionHandle connection);
 
        void onConnectionFailed(wspp::ConnectionHandle connection);
 
        void onMessageReceived(wspp::ConnectionHandle con, wspp::MessagePtr msg);
 
        void onHTTP(wspp::ConnectionHandle con);
 
        bool onValidate(wspp::ConnectionHandle con);
 
        bool onPing(wspp::ConnectionHandle con, std::string msg);
 
        bool disconnect(nap::utility::ErrorState& error);
 
        bool mRunning = false;                                                  
        std::mutex mConnectionMutex;                                            
        std::unordered_set<WebSocketTicketHash> mClientHashes;                  
        websocketpp::server<config> mEndPoint;                                  
        uint32 mLogLevel = 0;                                                   
        uint32 mAccessLogLevel = 0;                                             
        std::future<void> mServerTask;                                          
        std::vector<wspp::ConnectionHandle> mConnections;                       
        std::mutex mListenerMutex;                                              
        std::vector<IWebSocketServer*> mListeners;                              
    };
 
 
    // Not secured web-socket server end point
    using WebSocketServerEndPoint = WebSocketServerEndPointSetup<wspp::Config>;
 
 
    class NAPAPI SecureWebSocketServerEndPoint : public WebSocketServerEndPointSetup<wspp::ConfigTLS>
    {
        RTTI_ENABLE(WebSocketServerEndPointSetup<wspp::ConfigTLS>)
    public:
        enum class ETLSMode : nap::uint8
        {
            Intermediate = 0,       
            Modern                  
        };
 
        ETLSMode mMode = ETLSMode::Intermediate;                
        std::string mCertificateFile;                           
        std::string mPrivateKeyFile;                            
        std::string mPassphrase;                                
 
        bool init(utility::ErrorState& errorState) override;
 
        bool start(nap::utility::ErrorState &error) override;
 
    private:
        std::shared_ptr<asio::ssl::context> onTLSInit(wspp::ConnectionHandle con);
    };
 
 
    // Template Definitions
 
    template<typename config>
    bool WebSocketServerEndPointSetup<config>::start(nap::utility::ErrorState& error)
    {
        // Run until stopped
        assert(!mRunning);
 
        // Listen to messages on this specific port and ip address (if given)
        std::error_code stdec;
        if (mIPAddress.empty())
        {
            mEndPoint.listen(static_cast<uint16>(mPort), stdec);
        }
        else
        {
            mEndPoint.listen(mIPAddress, utility::stringFormat("%d", static_cast<uint16>(mPort)), stdec);
        }
 
        // Contains the error when opening port fails.
        if (stdec)
        {
            error.fail(stdec.message());
            return false;
        }
 
        // Queues a connection accept operation
        mEndPoint.start_accept(stdec);
        if (stdec)
        {
            error.fail(stdec.message());
            return false;
        }
 
        // Run until stopped
        mServerTask = std::async(std::launch::async, std::bind(&WebSocketServerEndPointSetup<config>::run, this));
        mRunning = true;
 
        return true;
    }
 
 
    template<typename config>
    bool WebSocketServerEndPointSetup<config>::isOpen() const
    {
        return mRunning;
    }
 
 
    template<typename config>
    bool WebSocketServerEndPointSetup<config>::send(const WebSocketConnection& connection, const std::string& message, EWebSocketOPCode code, nap::utility::ErrorState& error)
    {
        std::error_code stdec;
        mEndPoint.send(connection.mConnection, message, static_cast<wspp::OpCode>(code), stdec);
        if (stdec)
        {
            error.fail(stdec.message());
            return false;
        }
        return true;
    }
 
 
    template<typename config>
    bool WebSocketServerEndPointSetup<config>::send(const WebSocketConnection& connection, void const* payload, int length, EWebSocketOPCode code, nap::utility::ErrorState& error)
    {
        std::error_code stdec;
        mEndPoint.send(connection.mConnection, payload, length, static_cast<wspp::OpCode>(code), stdec);
        if (stdec)
        {
            error.fail(stdec.message());
            return false;
        }
        return true;
    }
 
 
    template<typename config>
    bool WebSocketServerEndPointSetup<config>::broadcast(const std::string& message, EWebSocketOPCode code, nap::utility::ErrorState& error)
    {
        bool success(true);
        std::lock_guard<std::mutex> lock(mConnectionMutex);
        for (auto& connection : mConnections)
        {
            if (!send(connection, message, code, error))
                success = false;
        }
        return success;
    }
 
 
    template<typename config>
    bool WebSocketServerEndPointSetup<config>::broadcast(void const* payload, int length, EWebSocketOPCode code, nap::utility::ErrorState& error)
    {
        bool success(true);
        std::lock_guard<std::mutex> lock(mConnectionMutex);
        for (auto& connection : mConnections)
        {
            if (!send(connection, payload, length, code, error))
                success = false;
        }
        return success;
    }
 
 
    template<typename config>
    std::string WebSocketServerEndPointSetup<config>::getHostName(const WebSocketConnection& connection)
    {
        std::error_code stdec;
        auto cptr = mEndPoint.get_con_from_hdl(connection.mConnection, stdec);
        return stdec ? "" : cptr->get_host();
    }
 
    template<typename config>
    void WebSocketServerEndPointSetup<config>::getHostNames(std::vector<std::string>& outHosts)
    {
        outHosts.clear();
        std::error_code stdec;
        std::lock_guard<std::mutex> lock(mConnectionMutex);
        outHosts.reserve(mConnections.size());
        for (auto& connection : mConnections)
        {
            auto cptr = mEndPoint.get_con_from_hdl(connection, stdec);
            if (!stdec)
                outHosts.emplace_back(cptr->get_host());
        }
    }
 
 
    template<typename config>
    int WebSocketServerEndPointSetup<config>::getConnectionCount()
    {
        std::lock_guard<std::mutex> lock(mConnectionMutex);
        return mConnections.size();
    }
 
    template<typename config>
    bool WebSocketServerEndPointSetup<config>::acceptsNewConnections()
    {
        if (mConnectionLimit < 0)
            return true;
 
        std::lock_guard<std::mutex> lock(mConnectionMutex);
        return mConnections.size() < mConnectionLimit;
    }
 
 
    template<typename config>
    void WebSocketServerEndPointSetup<config>::run()
    {
        // Start running until stopped
        mEndPoint.run();
    }
 
 
    template<typename config>
    void WebSocketServerEndPointSetup<config>::stop()
    {
        if (mRunning)
        {
            // Stop listening for new connections
            std::error_code stdec;
            mEndPoint.stop_listening(stdec);
            if (stdec)
            {
                assert(false);
                nap::Logger::error("%s: %s", mID.c_str(), stdec.message().c_str());
            }
 
            // Close all client connections
            utility::ErrorState napec;
            if (!disconnect(napec))
            {
                assert(false);
                nap::Logger::error("%s: %s", mID.c_str(), napec.toString().c_str());
            }
 
            // Explicitly stop
            mEndPoint.stop();
 
            // Wait for thread to finish
            assert(mServerTask.valid());
            mServerTask.wait();
            mRunning = false;
        }
    }
 
 
    template<typename config>
    WebSocketServerEndPointSetup<config>::~WebSocketServerEndPointSetup()
    {
        stop();
    }
 
 
    template<typename config>
    bool WebSocketServerEndPointSetup<config>::init(utility::ErrorState& errorState)
    {
        // Convert log levels
        mLogLevel = computeWebSocketLogLevel(mLibraryLogLevel);
        mAccessLogLevel = mLogConnectionUpdates ? websocketpp::log::alevel::all ^ websocketpp::log::alevel::frame_payload
            : websocketpp::log::alevel::fail;
 
        // Initiate logging
        mEndPoint.clear_error_channels(websocketpp::log::elevel::all);
        mEndPoint.set_error_channels(mLogLevel);
 
        mEndPoint.clear_access_channels(websocketpp::log::alevel::all);
        mEndPoint.set_access_channels(mAccessLogLevel);
 
        // If the endpoint can be re-used by other processes
        mEndPoint.set_reuse_addr(mAllowPortReuse);
 
        // Init asio
        std::error_code stdec;
        mEndPoint.init_asio(stdec);
        if (stdec)
        {
            errorState.fail(stdec.message());
            return false;
        }
 
        // Install connection open / closed handlers
        mEndPoint.set_http_handler(std::bind(&WebSocketServerEndPointSetup<config>::onHTTP, this, std::placeholders::_1));
        mEndPoint.set_open_handler(std::bind(&WebSocketServerEndPointSetup<config>::onConnectionOpened, this, std::placeholders::_1));
        mEndPoint.set_close_handler(std::bind(&WebSocketServerEndPointSetup<config>::onConnectionClosed, this, std::placeholders::_1));
        mEndPoint.set_fail_handler(std::bind(&WebSocketServerEndPointSetup<config>::onConnectionFailed, this, std::placeholders::_1));
        mEndPoint.set_validate_handler(std::bind(&WebSocketServerEndPointSetup<config>::onValidate, this, std::placeholders::_1));
        mEndPoint.set_ping_handler(std::bind(&WebSocketServerEndPointSetup<config>::onPing, this, std::placeholders::_1, std::placeholders::_2));
 
        // Install message handler
        mEndPoint.set_message_handler(std::bind(
            &WebSocketServerEndPointSetup<config>::onMessageReceived, this,
            std::placeholders::_1, std::placeholders::_2
        ));
 
        // Create unique hashes out of tickets
        // Server side tickets are required to have a password and username
        for (const auto& ticket : mClients)
            mClientHashes.emplace(ticket->toHash());
        return true;
    }
 
 
    template<typename config>
    void WebSocketServerEndPointSetup<config>::onConnectionOpened(wspp::ConnectionHandle connection)
    {
        std::error_code stdec;
        auto cptr = mEndPoint.get_con_from_hdl(connection, stdec);
        if (stdec)
        {
            nap::Logger::error(stdec.message());
            return;
        }
 
        // Add to list of actively managed connections
        {
            std::lock_guard<std::mutex> lock(mConnectionMutex);
            mConnections.emplace_back(cptr);
        }
 
        for (IWebSocketServer* listener : mListeners)
            listener->onConnectionOpened(WebSocketConnection(connection));
    }
 
 
    template<typename config>
    void WebSocketServerEndPointSetup<config>::onConnectionClosed(wspp::ConnectionHandle connection)
    {
        std::error_code stdec;
        auto cptr = mEndPoint.get_con_from_hdl(connection, stdec);
        if (stdec)
        {
            nap::Logger::error(stdec.message());
            return;
        }
 
        // Signal that it closed
        for (IWebSocketServer* listener : mListeners)
            listener->onConnectionClosed(WebSocketConnection(connection), cptr->get_ec().value(), cptr->get_ec().message());
 
        // Remove from internal list of connections
        {
            std::lock_guard<std::mutex> lock(mConnectionMutex);
            auto found_it = std::find_if(mConnections.begin(), mConnections.end(), [&](const auto& it)
                {
                    auto client_ptr = mEndPoint.get_con_from_hdl(it, stdec);
                    if (stdec)
                    {
                        nap::Logger::error(stdec.message());
                        return false;
                    }
                    return cptr == client_ptr;
                });
 
            // Having no connection to remove is a serious error and should never occur.
            if (found_it == mConnections.end())
            {
                assert(false);
                return;
            }
            mConnections.erase(found_it);
        }
    }
 
 
    template<typename config>
    void WebSocketServerEndPointSetup<config>::onConnectionFailed(wspp::ConnectionHandle connection)
    {
        std::error_code stdec;
        auto cptr = mEndPoint.get_con_from_hdl(connection, stdec);
        if (stdec)
        {
            nap::Logger::error(stdec.message());
            return;
        }
 
        for (IWebSocketServer* listener : mListeners)
            listener->onConnectionFailed(WebSocketConnection(connection), cptr->get_ec().value(), cptr->get_ec().message());
    }
 
 
    template<typename config>
    void WebSocketServerEndPointSetup<config>::onMessageReceived(wspp::ConnectionHandle con, wspp::MessagePtr msg)
    {
        WebSocketMessage message(msg);
        for (IWebSocketServer* listener : mListeners)
            listener->onMessageReceived(WebSocketConnection(con), message);
    }
 
 
    template<typename config>
    void WebSocketServerEndPointSetup<config>::onHTTP(wspp::ConnectionHandle con)
    {
        // Get handle to connection
        std::error_code stdec;
        auto conp = mEndPoint.get_con_from_hdl(con, stdec);
        if (stdec)
        {
            nap::Logger::error(stdec.message());
            conp->set_status(websocketpp::http::status_code::internal_server_error);
            return;
        }
 
        // Set whether the response can be shared with requesting code from the given origin.
        conp->append_header("Access-Control-Allow-Origin", mAccessAllowControlOrigin);
 
        // Get request method
        std::string method = conp->get_request().get_method();
 
        // Handle CORS preflight request
        if (method.compare("OPTIONS") == 0)
        {
            conp->set_status(websocketpp::http::status_code::no_content);
            conp->append_header("Access-Control-Allow-Methods", "OPTIONS, POST");
            conp->append_header("Access-Control-Allow-Headers", "Content-Type");
            return;
        }
 
        // Reject methods other than OPTIONS and POST
        if (method.compare("POST") != 0)
        {
            conp->set_status(websocketpp::http::status_code::method_not_allowed,
                "only OPTIONS and POST requests are allowed");
            return;
        }
 
        // When there is no access policy the server doesn't generate tickets
        if (mMode == EAccessMode::EveryOne)
        {
            conp->set_status(websocketpp::http::status_code::conflict,
                "unable to generate ticket, no access policy set");
            return;
        }
 
        // Get request body
        std::string body = conp->get_request_body();
 
        // Try to parse as JSON
        rapidjson::Document document;
        rapidjson::ParseResult parse_result = document.Parse(body.c_str());
        if (!parse_result || !document.IsObject())
        {
            conp->set_status(websocketpp::http::status_code::bad_request,
                "unable to parse as JSON");
            return;
        }
 
        // Extract user information, this field is required
        if (!document.HasMember("user"))
        {
            conp->append_header("WWW-Authenticate", "NAPUserPass");
            conp->set_status(websocketpp::http::status_code::unauthorized,
                "missing required member: 'user");
            return;
        }
 
        // Extract pass information, this field is required
        if (!document.HasMember("pass"))
        {
            conp->append_header("WWW-Authenticate", "NAPUserPass");
            conp->set_status(websocketpp::http::status_code::unauthorized,
                "missing required member: 'pass");
            return;
        }
 
        // Create ticket
        nap::WebSocketTicket ticket;
 
        // Populate and initialize
        ticket.mID = math::generateUUID();
        ticket.mUsername = document["user"].GetString();
        ticket.mPassword = document["pass"].GetString();
        utility::ErrorState error;
        if (!ticket.init(error))
        {
            conp->append_header("WWW-Authenticate", "NAPUserPass");
            conp->set_status(websocketpp::http::status_code::unauthorized,
                utility::stringFormat("invalid username or password: %s", error.toString().c_str()));
            return;
        }
 
        // If we only allow specific clients we extract the password here.
        // The pass is used in combination with the username to create a validation hash.
        if (mMode == EAccessMode::Reserved)
        {
            // Locate ticket, if there is no tick that matches the request
            // The username or password is wrong and the request invalid
            if (mClientHashes.find(ticket.toHash()) == mClientHashes.end())
            {
                conp->append_header("WWW-Authenticate", "NAPUserPass");
                conp->set_status(websocketpp::http::status_code::unauthorized,
                    "invalid username or password");
                return;
            }
        }
 
        // Convert ticket to binary blob
        // This ticket can now be used to validate the request
        std::string ticket_str;
        if (!ticket.toBinaryString(ticket_str, error))
        {
            nap::Logger::error(error.toString());
            conp->set_status(websocketpp::http::status_code::internal_server_error);
            return;
        }
 
        // Set ticket as body
        conp->set_body(ticket_str);
        conp->set_status(websocketpp::http::status_code::ok);
    }
 
 
    template<typename config>
    bool WebSocketServerEndPointSetup<config>::onValidate(wspp::ConnectionHandle con)
    {
        // Get connection handle
        std::error_code stdec;
        auto conp = mEndPoint.get_con_from_hdl(con, stdec);
        if (stdec)
        {
            nap::Logger::error(stdec.message());
            conp->set_status(websocketpp::http::status_code::internal_server_error);
            return false;
        }
 
        // Make sure we accept new connections
        if (!acceptsNewConnections())
        {
            conp->set_status(websocketpp::http::status_code::forbidden,
                "client connection count exceeded");
            return false;
        }
 
        // Get sub-protocol field
        const std::vector<std::string>& sub_protocol = conp->get_requested_subprotocols();
 
        // When every client connection is allowed no sub-protocol field must be defined.
        if (mMode == EAccessMode::EveryOne)
        {
            if (sub_protocol.empty())
                return true;
 
            // Sub-protocol requested but can't authenticate.
            conp->set_status(websocketpp::http::status_code::not_found, "invalid sub-protocol");
            return false;
        }
 
        // We have tickets and therefore specific clients we accept.
        // Use the sub_protocol to extract client information
        if (sub_protocol.empty())
        {
            conp->set_status(websocketpp::http::status_code::forbidden,
                "unable to extract ticket");
            return false;
        }
 
        // Extract ticket
        conp->select_subprotocol(sub_protocol[0]);
 
        // Convert from binary into string
        utility::ErrorState error;
        rtti::DeserializeResult result;
        WebSocketTicket* client_ticket = WebSocketTicket::fromBinaryString(sub_protocol[0], result, error);
        if (client_ticket == nullptr)
        {
            conp->set_status(websocketpp::http::status_code::forbidden,
                "first sub-protocol argument is not a valid ticket object");
            return false;
        }
 
        // Valid ticket was extracted, allowed access
        if (mMode == EAccessMode::Ticket)
            return true;
 
        // Locate ticket
        if (mClientHashes.find(client_ticket->toHash()) == mClientHashes.end())
        {
            conp->set_status(websocketpp::http::status_code::forbidden,
                "not a valid ticket");
            return false;
        }
 
        // Welcome!
        return true;
    }
 
 
    template<typename config>
    bool WebSocketServerEndPointSetup<config>::onPing(wspp::ConnectionHandle con, std::string msg)
    {
        return true;
    }
 
 
    template<typename config>
    bool WebSocketServerEndPointSetup<config>::disconnect(nap::utility::ErrorState& error)
    {
        std::lock_guard<std::mutex> lock(mConnectionMutex);
        bool success = true;
        for (auto& connection : mConnections)
        {
            std::error_code stdec;
            mEndPoint.close(connection, websocketpp::close::status::going_away, "disconnected", stdec);
            if (stdec)
            {
                error.fail(stdec.message());
                success = false;
            }
        }
        mConnections.clear();
        return success;
    }
 
 
    template<typename config>
    void WebSocketServerEndPointSetup<config>::registerListener(IWebSocketServer& server)
    {
        std::unique_lock<std::mutex> lock(mListenerMutex);
        mListeners.push_back(&server);
    }
 
 
    template<typename config>
    void WebSocketServerEndPointSetup<config>::unregisterListener(IWebSocketServer& server)
    {
        std::unique_lock<std::mutex> lock(mListenerMutex);
        mListeners.erase(std::remove(mListeners.begin(), mListeners.end(), &server), mListeners.end());
    }
 
 
}