rublon-ssh/PAM/ssh/include/rublon/websockets.hpp

#pragma once

#include <cstddef>
#include <memory>

#include <rublon/error.hpp>
#include <rublon/pam_action.hpp>
#include <rublon/utils.hpp>

#include <condition_variable>
#include <mutex>
#include <sio_client.h>
#include <string>

namespace rublon {

class WebSocketConnectionListener : public std::enable_shared_from_this< WebSocketConnectionListener > {
  public:
    enum class Status { Disconnected, Conected };

  private:
    sio::client & _ws;

    std::mutex & _lock;
    std::condition_variable_any & _cond;

    bool handshake{false};
    Status _status{Status::Disconnected};

  public:
    WebSocketConnectionListener(sio::client & ws, std::mutex & lock, std::condition_variable_any & cond)
     : _ws(ws), _lock{lock}, _cond{cond} {}

    void makeCallbacks() {
        auto _this = shared_from_this();
        _ws.set_open_listener([_this]() { _this->on_connected(); });
        _ws.set_close_listener([_this](sio::client::close_reason const & reason) { _this->on_close(reason); });
        _ws.set_fail_listener([_this]() { _this->on_fail(); });
    }

    tl::expected< Status, bool > waitForConnection() {
        std::lock_guard lock{_lock};
        log(LogLevel::Info, "Waiting for connection");
        _cond.wait(_lock);
        return _status;
    }

    void on_connected() {
        std::unique_lock{_lock};
        log(LogLevel::Info, "WebSocket connected");
        handshake = true;
        _status   = Status::Conected;
        _cond.notify_all();
    }

    void on_close(sio::client::close_reason const & reason) {
        std::unique_lock{_lock};
        log(LogLevel::Info,
          "WebSocket connection %s",
          reason == sio::client::close_reason::close_reason_drop ? "dropped" : "closed normally");
        handshake = true;
        _status   = Status::Disconnected;
        _cond.notify_all();
    }

    void on_fail() {
        std::unique_lock{_lock};
        log(LogLevel::Info, "WebSocket connection fail");
        handshake = true;
        _status   = Status::Disconnected;
        _cond.notify_all();
    }
};

class WebSocketSingleShotEventListener : public std::enable_shared_from_this< WebSocketSingleShotEventListener > {
  public:
    enum class Status { Unknown, Approved, Denied, Expired };

  private:
    sio::client & _ws;

    std::mutex & _lock;
    std::condition_variable_any & _cond;

    bool event_received = false;
    Status _status{Status::Unknown};

  public:
    std::string _data{};

  private:
    void printobj(std::size_t i, const std::map< std::string, sio::message::ptr > & objects) {
        std::size_t o{};
        for(const auto & [name, msg] : objects) {
            log(LogLevel::Debug, "Object %ld: '%s'", i, name.c_str());
            printMessage(i, o++, msg);
        }
    };

    void printMessage(std::size_t i, std::size_t o, const sio::message::ptr & message) {
        switch(message->get_flag()) {
            case sio::message::flag_integer:
                log(LogLevel::Debug, "message %ld/%ld integer   : %d", i, o, message->get_int());
                break;
            case sio::message::flag_double:
                log(LogLevel::Debug, "message %ld/%ld double    : %f", i, o, message->get_double());
                break;
            case sio::message::flag_string:
                log(LogLevel::Debug, "message %ld/%ld string    : %s", i, o, message->get_string().c_str());
                break;
            case sio::message::flag_binary:
                log(LogLevel::Debug, "message %ld/%ld binary    : %s", i, o, message->get_binary()->c_str());
                break;
            case sio::message::flag_array:
                log(LogLevel::Debug, "message %ld/%ld array     : blah", i, o);
                break;
            case sio::message::flag_object:
                log(LogLevel::Debug, "message %ld/%ld object    : obj ->", i, o);
                printobj(i, message->get_map());
                break;
            case sio::message::flag_boolean:
                log(LogLevel::Debug, "message %ld/%ld bool     : %s", i, o, message->get_bool() ? "yes" : "no");
                break;
            case sio::message::flag_null:
                log(LogLevel::Debug, "message %ld/%ld NULL", i, o);
                break;
            default:
                log(LogLevel::Info, "Message with unknown type");
        }
    };

    void notify(Status status) {
        event_received = true;
        _status        = status;
        _cond.notify_all();
    }

    void onConfirmed(sio::event & e) {
        std::lock_guard lock{_lock};
        log(LogLevel::Info, "Autentication confirmed");
        _data = e.get_messages().at(1)->get_map()["data"]->get_string();
        notify(Status::Approved);
    }

    void onDenied(sio::event & e) {
        std::lock_guard lock{_lock};
        log(LogLevel::Info, "Autentication denied by user");
        // _data = e.get_messages().at(1)->get_map()["redirectUrl"]->get_string();
        for(std::size_t i{}; i < e.get_messages().size(); i++)
            printMessage(i, 0, e.get_messages().at(i));
        notify(Status::Denied);
    }

    void onExpired(sio::event & /*e*/) {
        std::lock_guard lock{_lock};
        log(LogLevel::Error, "Autentication call expierd");
        notify(Status::Expired);
    }

    void onErrorEvent(sio::message::ptr const & message) {
        std::lock_guard lock{_lock};
        log(LogLevel::Error, "Error: %s", message->get_string().c_str());
        event_received = true;
        _cond.notify_all();
    }

  public:
    WebSocketSingleShotEventListener(sio::client & ws, std::mutex & lock, std::condition_variable_any & cond)
     : _ws{ws}, _lock{lock}, _cond{cond} {}

    void makeCallbacks() {
        auto _this = shared_from_this();
        _ws.socket()->on("transactionConfirmed", [_this](sio::event & e) { _this->onConfirmed(e); });
        _ws.socket()->on("transactionDenied", [_this](sio::event & e) { _this->onDenied(e); });
        _ws.socket()->on("transactionExpired", [_this](sio::event & e) { _this->onExpired(e); });

        _ws.socket()->on_error([_this](sio::message::ptr const & message) { _this->onErrorEvent(message); });
        // _ws.socket()->on_any([_this](sio::event & e) { _this->onAny(e); });
    }

    /// TODO pass timeout
    tl::expected< AuthenticationStatus, Error > waitForEvent() {
        std::lock_guard lock{_lock};
        if(!event_received) {
            log(LogLevel::Info, "Waiting for WS event");
            auto timeout = std::chrono::system_clock::now() + std::chrono::minutes{2};
            while(!event_received) {
                if(_cond.wait_until(_lock, timeout) == std::cv_status::timeout) {
                    log(LogLevel::Info, "Waiting for WS event failed due to WS timeout");
                    break;
                } else {
                    if(event_received) {
                        log(LogLevel::Info, "Event arrived");
                        break;
                    } else
                        log(LogLevel::Info, "Event arrived, but it's not an action");
                }
            }
        } else {
            log(LogLevel::Info, "event already received");
        }

        AuthenticationStatus status{
          _status == Status::Approved ? AuthenticationStatus::Action::Confirmed : AuthenticationStatus::Action::Denied, _data};
        return status;
    }
};

class WebSocket {
    std::mutex _lock;
    std::condition_variable_any _cond;
    sio::client ws;

  public:
    WebSocket()  = default;
    ~WebSocket() = default;

    std::shared_ptr< WebSocketSingleShotEventListener > listen() {
        const auto createEventListener = [&]() -> std::shared_ptr< WebSocketSingleShotEventListener > {
            log(LogLevel::Debug, "Initializiing WebSocketSingleShotEventListener");
            auto wsSingleShotEventListener = std::make_shared< WebSocketSingleShotEventListener >(ws, _lock, _cond);
            wsSingleShotEventListener->makeCallbacks();
            return wsSingleShotEventListener;
        };
        return createEventListener();
    }

    bool AttachToCore(std::string_view uri, std::string_view tid) {
        /// needed here only for rublon core api URL, so 1k fine
        memory::StrictMonotonic_1k_HeapResource memoryResource;
        std::pmr::string str{uri.data(), uri.size(), &memoryResource};
        std::pmr::string httpsPrefix = {"https://", &memoryResource};

        if(uri.rfind(httpsPrefix, 0) == 0) {
            str.replace(str.find(httpsPrefix), httpsPrefix.size(), "wss://");
        }
        str += "/ws/socket.io/";

        // ws.set_logs_verbose();
        ws.connect(str.c_str());

        const auto attachToTransactionConfirmationChannel = [&](const auto & /*status*/) -> tl::expected< bool, bool > {
            log(LogLevel::Debug, "Initializiing attachToTransactionConfirmationChannel");

            memory::StrictMonotonic_1k_HeapResource resource;
            std::pmr::string channel{&resource};
            channel.reserve(100);
            channel += "transactionConfirmation.";
            channel += tid;

            auto message = std::dynamic_pointer_cast< sio::object_message >(sio::object_message::create());
            message->insert("channel", channel.c_str());
            log(LogLevel::Debug, "emiting %s message on subscribe channel", channel.c_str());
            ws.socket()->emit("subscribe", {message});

            return true; /// TODO
        };

        auto wsConnectionListener = std::make_shared< WebSocketConnectionListener >(ws, _lock, _cond);
        wsConnectionListener->makeCallbacks();

        auto v = wsConnectionListener
                   ->waitForConnection() //
                   .and_then(attachToTransactionConfirmationChannel);

        log(LogLevel::Debug, "WS Connected");
        return v.value_or(false);
    }
};
} // namespace rublon