LNXRNT/Sources/websocket.h

#pragma once

#ifdef WITH_NETWORKING

#include "websocket_config.h"
#include "lockfree_queue.h"
#include "global_thread_pool.h"
#include <kinc/network/socket.h>
#include <functional>
#include <string>
#include <vector>
#include <atomic>
#include <thread>
#include <v8.h>
#include <memory>

#ifdef _WIN32
typedef void* SOCKET_HANDLE;
#else
typedef int SOCKET_HANDLE;
#endif

#ifdef WITH_SSL
#ifndef _WIN32
#include <openssl/ssl.h>
#include <openssl/err.h>
#include <openssl/crypto.h>
#endif
#endif

namespace WebSocketWrapper {
    enum ReadyState {
        CONNECTING = 0,
        OPEN = 1,
        CLOSING = 2,
        CLOSED = 3
    };

    enum EventType {
        EVENT_OPEN,
        EVENT_MESSAGE,
        EVENT_ERROR,
        EVENT_CLOSE
    };

    struct WebSocketEvent {
        EventType type;
        std::string data;
        int code;
        std::string reason;
        
        WebSocketEvent() : type(EVENT_OPEN), code(0) {}
        WebSocketEvent(EventType t) : type(t), code(0) {}
        WebSocketEvent(EventType t, const std::string& d) : type(t), data(d), code(0) {}
        WebSocketEvent(EventType t, int c, const std::string& r) : type(t), code(c), reason(r) {}
    };

    class WebSocketClient {
    public:
        WebSocketClient(v8::Isolate* isolate, v8::Global<v8::Context>* global_context, const std::string& url);
        ~WebSocketClient();

        // this is for browser compatible methods
        void send(const std::string& data);
        void sendBinary(const std::string& data);
        void close(int code = 1000, const std::string& reason = "");
        // process pending events on main thread
        void processEvents();
        
        ReadyState getReadyState() const { return ready_state_; }
        const std::string& getUrl() const { return url_; }
        const std::string& getProtocol() const { return protocol_; }
        int getBufferedAmount() const { return buffered_amount_; }
        
        // V8 callback setters
        void setOnOpen(v8::Local<v8::Function> callback);
        void setOnMessage(v8::Local<v8::Function> callback);
        void setOnError(v8::Local<v8::Function> callback);
        void setOnClose(v8::Local<v8::Function> callback);

    private:
        v8::Isolate* isolate_;
        v8::Global<v8::Context>* global_context_;
        std::string url_;
        std::string protocol_;
        ReadyState ready_state_;
        bool is_ssl_;
        int buffered_amount_;
        
        // V8 callback storage
        v8::Global<v8::Function> on_open_;
        v8::Global<v8::Function> on_message_;
        v8::Global<v8::Function> on_error_;
        v8::Global<v8::Function> on_close_;
        
        LockFreeQueue<WebSocketEvent> event_queue_;
        
        void* ws_;
        void* ssl_ws_;
        
        #ifdef WITH_SSL
        #ifdef _WIN32
        // windows SChannel opaque pointers 
        void* ssl_cred_handle_;
        void* ssl_context_handle_;
        bool ssl_context_initialized_;
        int ssl_socket_;  
        std::vector<char> ssl_buffer_;  
        #else
        SSL_CTX* ssl_ctx_;
        SSL* ssl_;
        bool ssl_initialized_;
        #endif
        
        bool initWSL();
        void cleanupWSL();
        bool performWSLHandshake(int socket, const std::string& host);
        int wslRead(char* buffer, int length);
        int webSocketSSLSend(const void* data, int len);
        int webSocketSSLSend(const char* buffer, int length);
        int webSocketSSLReceive(char* buffer, int bufferSize);
        #endif
        
        std::string base64Encode(const std::string& data);
        
        // called from worker thread
        void handleOpen();
        void handleMessage(const std::string& message);
        void handleError(const std::string& error);
        void handleClose(int code, const std::string& reason);
        
        // called from main thread
        void processOpenEvent();
        void processMessageEvent(const std::string& message);
        void processMessageBatch(const std::vector<std::string>& messages);
        void processErrorEvent(const std::string& error);
        void processCloseEvent(int code, const std::string& reason);
        // main thread only
        void callCallback(v8::Global<v8::Function>& callback, int argc, v8::Local<v8::Value> argv[]);
        
        bool connectToServer(const std::string& host, int port, const std::string& path);
        void messageLoop();
        void processFrame(uint8_t opcode, bool fin, const std::vector<uint8_t>& payload);
        
        // RFC 6455 protocol 
        bool performWebSocketHandshake(int sock, const std::string& host, int port, const std::string& path);
        void fireOpenEvent();
        void fireMessageEvent(const std::string& message, bool binary = false);
        void fireErrorEvent(const std::string& error);
        void fireCloseEvent(int code, const std::string& reason);
        
        bool parseWebSocketFrame(const std::vector<uint8_t>& buffer, size_t& offset);
        std::vector<uint8_t> createWebSocketFrame(const std::string& message, uint8_t opcode = 0x1);
        std::vector<uint8_t> createWebSocketBinaryFrame(const std::string& data);
        std::vector<uint8_t> createCloseFrame(uint16_t code, const std::string& reason);
        std::vector<uint8_t> createPongFrame(const std::vector<uint8_t>& payload);
        
        std::string generateWebSocketKey();
        std::string base64Encode(const std::vector<uint8_t>& data);
        std::vector<uint8_t> sha1Hash(const std::string& data);
        void maskData(std::vector<uint8_t>& data, uint32_t maskKey);
        
        std::vector<uint8_t> partial_frame_buffer_;
        bool expecting_continuation_ = false;
        uint8_t continuation_opcode_ = 0;
        
        bool initializeWebSocketSSL(int sock, const std::string& host);
        void cleanupWebSocketSSL();
        
        bool parseUrl(const std::string& url, std::string& host, int& port, std::string& path);
    };

    void initialize();
    void cleanup();
    // called from main event loop
    void processEvents(); 
    
    int createWebSocketConnection(v8::Isolate* isolate, const std::string& url);
    
    // get global context from main.cpp 
    v8::Global<v8::Context>* getGlobalContext();
}

// V8 binding functions called from main.cpp
void runt_websocket_create(const v8::FunctionCallbackInfo<v8::Value>& args);
void runt_websocket_send(const v8::FunctionCallbackInfo<v8::Value>& args);
void runt_websocket_send_binary(const v8::FunctionCallbackInfo<v8::Value>& args);
void runt_websocket_close(const v8::FunctionCallbackInfo<v8::Value>& args);
void runt_websocket_get_ready_state(const v8::FunctionCallbackInfo<v8::Value>& args);
void runt_websocket_get_url(const v8::FunctionCallbackInfo<v8::Value>& args);
void runt_websocket_get_buffered_amount(const v8::FunctionCallbackInfo<v8::Value>& args);
void runt_websocket_set_onopen(const v8::FunctionCallbackInfo<v8::Value>& args);
void runt_websocket_set_onmessage(const v8::FunctionCallbackInfo<v8::Value>& args);
void runt_websocket_set_onerror(const v8::FunctionCallbackInfo<v8::Value>& args);
void runt_websocket_set_onclose(const v8::FunctionCallbackInfo<v8::Value>& args);

void createWebSocketClass(v8::Isolate* isolate, v8::Local<v8::ObjectTemplate>& global);
void createWebSocketEventClasses(v8::Isolate* isolate, v8::Local<v8::ObjectTemplate>& global);
#endif