websockets.h

/*NB_REVISION*/
 
/*NB_COPYRIGHT*/
 
#ifndef __WEBSOCKETS_H
#define __WEBSOCKETS_H
 
// NB Definitions
#include <predef.h>
 
// NB Libs
#include <buffers.h>
#include <http.h>
#include <iosys.h>
#include <nbrtos.h>
#include <nettypes.h>
 
 
#define WS_BUFFER_SEGMENTS 4                        
#define WS_MAX_SOCKS 4                              
#define WS_BUFFER_MIN 10                            
#define WS_ACCUM_DLY ((TICKS_PER_SECOND / 4) + 1)   
#define WS_FLUSH_TIMEOUT 0                          
 
#define WS_FIN_BIT  0x80      
#define WS_OP_CONT  0x0       
#define WS_OP_TEXT  0x1       
#define WS_OP_BIN   0x2       
#define WS_OP_CLOSE 0x8       
#define WS_OP_PING  0x9       
#define WS_OP_PONG  0xA       
 
#define WS_SO_TEXT 0x8        
 
namespace NB
{
class WSCritObj;  // WebSocket critical section object for thread synchronization
 
class WebSocket
{
    /* Client-side short frame header (payload <=125 bytes, with mask) */
    struct WS_Client_Short
    {
        uint8_t opAndFin;      // Opcode (4 bits) and FIN bit (1 bit) combined
        uint8_t maskAndLen;    // Mask bit (1 bit) and payload length (7 bits)
        uint32_t mask;         // 32-bit masking key for payload
        uint8_t pData[];       // Flexible array member for payload data
    } __attribute__((packed));
    
    /* Client-side medium frame header (payload 126-65535 bytes, with mask) */
    struct WS_Client_Med
    {
        uint8_t opAndFin;      // Opcode (4 bits) and FIN bit (1 bit) combined
        uint8_t maskAndLen;    // Mask bit (1 bit) and payload length indicator (7 bits = 126)
        beuint16_t length;     // 16-bit big-endian actual payload length
        uint32_t mask;         // 32-bit masking key for payload
        uint8_t pData[];       // Flexible array member for payload data
    } __attribute__((packed));
    
    /* Client-side long frame header (payload >65535 bytes, with mask) */
    struct WS_Client_Long
    {
        uint8_t opAndFin;      // Opcode (4 bits) and FIN bit (1 bit) combined
        uint8_t maskAndLen;    // Mask bit (1 bit) and payload length indicator (7 bits = 127)
        beuint32_t lengthHi;   // Upper 32 bits of 64-bit big-endian payload length
        beuint32_t lengthLo;   // Lower 32 bits of 64-bit big-endian payload length
        uint32_t mask;         // 32-bit masking key for payload
        uint8_t pData[];       // Flexible array member for payload data
    } __attribute__((packed));
    
    /* Server-side short frame header (payload <=125 bytes, no mask) */
    struct WS_Server_Short
    {
        uint8_t opAndFin;      // Opcode (4 bits) and FIN bit (1 bit) combined
        uint8_t maskAndLen;    // Mask bit (1 bit = 0) and payload length (7 bits)
        uint8_t pData[];       // Flexible array member for payload data
    } __attribute__((packed));
    
    /* Server-side medium frame header (payload 126-65535 bytes, no mask) */
    struct WS_Server_Med
    {
        uint8_t opAndFin;      // Opcode (4 bits) and FIN bit (1 bit) combined
        uint8_t maskAndLen;    // Mask bit (1 bit = 0) and payload length indicator (7 bits = 126)
        beuint16_t length;     // 16-bit big-endian actual payload length
        uint8_t pData[];       // Flexible array member for payload data
    } __attribute__((packed));
    
    /* Server-side long frame header (payload >65535 bytes, no mask) */
    struct WS_Server_Long
    {
        uint8_t opAndFin;      // Opcode (4 bits) and FIN bit (1 bit) combined
        uint8_t maskAndLen;    // Mask bit (1 bit = 0) and payload length indicator (7 bits = 127)
        beuint32_t lengthHi;   // Upper 32 bits of 64-bit big-endian payload length
        beuint32_t lengthLo;   // Lower 32 bits of 64-bit big-endian payload length
        uint8_t pData[];       // Flexible array member for payload data
    } __attribute__((packed));
    
    /*
     * WebSocket connection state machine enumeration
     * 
     * Tracks the lifecycle of a WebSocket connection from initialization through
     * establishment, closure handshake, and final cleanup.
     */
    enum WS_State
    {
        WS_INIT,           // Initial state, connection not yet established
        WS_CONNECTING,     // Handshake in progress, transitioning to WebSocket protocol
        WS_ESTABLISHED,    // Connection fully established, ready for data exchange
        WS_CLOSE_PENDING,  // Close requested locally, waiting to send close frame
        WS_CLOSE_SENT,     // Close frame sent, awaiting peer's close response
        WS_CLOSE_RECV,     // Close frame received from peer, sending acknowledgment
        WS_TCP_CLOSED,     // Underlying TCP connection has been closed
        WS_CLOSED,         // WebSocket fully closed, cleanup complete
        WS_ABORT,          // Connection aborted due to error
        WS_FAIL            // Connection failed (handshake failure, protocol error)
    };
    
    static const char *WS_StateStr(WS_State state);  // Convert WS_State enum to string for debugging
    
    /*
     * WebSocket close status codes (RFC 6455)
     * 
     * Standard status codes sent in close frames to indicate the reason for
     * connection closure. Some codes must not be sent on the wire.
     */
    enum WS_StatusCode
    {
        WS_STAT_NORM_CLOSE = 1000,        // Normal closure, connection purpose fulfilled
        WS_STAT_GOING_AWAY = 1001,        // Endpoint going away (server shutdown, page navigation)
        WS_STAT_PROT_ERROR = 1002,        // Protocol error detected
        WS_STAT_UNACCEPT_TYPE = 1003,     // Unacceptable data type received
        WS_STAT_NONE = 1005,              // Reserved: No status code present (must not be sent)
        WS_STAT_NONE_ABNORMAL = 1006,     // Reserved: Abnormal closure, no close frame (must not be sent)
        WS_STAT_NOT_CONSISTENT = 1007,    // Data not consistent with message type (e.g., invalid UTF-8)
        WS_STAT_POLICY_VIOLATION = 1008,  // Policy violation detected
        WS_STAT_MSG_TOO_BIG = 1009,       // Message too large to process
        WS_STAT_EXPECTED_EXTENSION = 1010,// Expected extension not negotiated
        WS_STAT_UNEXPECTED_COND = 1011,   // Unexpected condition prevented request fulfillment
        WS_STAT_TLS_FAILURE = 1015,       // Reserved: TLS handshake failure (must not be sent)
    };
    
    fifo_buffer_storage rxBuffer;   // Receive buffer for incoming data from TCP connection
    fifo_buffer_storage ctlBuffer;  // Control frame buffer for ping/pong/close frames
    fifo_buffer_storage txBuffer;   // Transmit buffer for outgoing data to TCP connection
 
public:
 
    int m_fd_tcp;
 
    int m_fd_ws;
 
    OS_CRIT socket_crit;
 
private:
    WS_State m_state;            // Current connection state (INIT, CONNECTING, ESTABLISHED, etc.)
    int m_options;               // Socket options bitmap (e.g., WS_SO_TEXT for text mode)
    uint32_t m_SendTime;         // Timestamp of last send operation (in ticks)
    uint32_t m_AccumDly;         // Accumulation delay for batching writes (in ticks)
    bool m_serverNotClient;      // True if this is server-side connection, false for client
 
    volatile bool m_CtlBlock;    // Flag indicating control frame processing is blocked
    uint32_t m_currentIndex;     // Current buffer index for frame processing
    uint32_t m_remLen;           // Remaining length to process in current operation
 
    uint8_t m_frameHeadBuf[14];  // Buffer for incoming frame header (max 14 bytes)
    uint8_t m_frameHeadNext;     // Next byte position in frame header buffer
    bool m_gotLen;               // True if complete length field has been received
    uint32_t m_currentMask;      // Current masking key for frame payload
    uint32_t m_remainingLen;     // Remaining payload bytes to receive in current frame
 
    uint8_t m_txFrameHeadBuf[14];// Buffer for outgoing frame header (max 14 bytes)
    uint8_t m_txFrameHeadNext;   // Next byte position in TX frame header buffer
    uint8_t m_txFrameHeadLen;    // Total length of TX frame header
    uint32_t m_txCurrMask;       // Current masking key for outgoing payload
    uint32_t m_txRemLen;         // Remaining payload bytes to send in current frame
    int m_txFlushInProgress;     // Non-zero if flush operation is currently in progress
    bool m_txCtlSendInProgress;  // True if control frame transmission is in progress
 
    /* Linked list node for tasks waiting on control frame operations */
    struct PendingCtlWait {
        OS_SEM          sem;     // Semaphore for blocking the waiting task
        PendingCtlWait *pNext;   // Pointer to next waiting task in linked list
    };
    PendingCtlWait *m_ctlWait;   // Head of linked list for tasks waiting on control frames
 
    uint32_t m_pingTick;         // Tick count when last ping was sent
    int      m_pingLen;          // Length of last ping payload
    OS_SEM  *m_pingWaitSem;      // Semaphore for blocking on ping reply
 
    uint32_t m_txBufferMin;      // Minimum TX buffer space threshold
 
    void ServerRxFromTCP();      // Process incoming data from TCP socket (server-side)
    bool ReadRemainingHeader();  // Continue reading frame header bytes, return true when complete
    void ProcessCtlFrame();      // Process received control frame (ping/pong/close)
    bool ValidatePong(uint32_t seed, uint32_t len);  // Validate pong payload matches ping seed
    void Close(bool closedAfterSend, uint16_t code = 0, const char *reason = NULL, int len = 0);  // Initiate WebSocket close handshake
    void RemoteClose();          // Handle close frame received from peer
    void CleanUp();              // Clean up resources after connection closure
    void SendCtlFrame();         // Send queued control frame (ping/pong/close)
    void Init(int fd_tcp, int fd_ws, bool serverNotClient);  // Initialize WebSocket connection state
    void Flush_Header(uint8_t *buf, uint32_t mask, uint32_t len);  // Write frame header to buffer
 
    void AddCtlWait(PendingCtlWait *pWait);  // Add task to control frame wait queue
    void PostToCtlWait();        // Wake up tasks waiting on control frame completion
 
    static uint32_t s_flushTick;       // Global tick counter for flush timing
    static fd_set s_pendingCtlSocks;   // Set of sockets with pending control frames
    static OS_CRIT s_ws_crit_obj;      // Global critical section for WebSocket operations
    static bool s_bFinishedInit;       // True if static initialization is complete
 
    static void StaticInit();          // Initialize static/global WebSocket subsystem
    static int GetNewSocket(int fd_tcp, bool serverNotClient);  // Allocate new WebSocket structure
 
    static WebSocket *GetRecordFromTCP(int fd);  // Find WebSocket record by TCP file descriptor
 
    // Core connection establishment logic
    static int CoreConnect(IPADDR ip, const char *host, const char *resource, const char *protocol, int portnum, bool useSSL = false);  
 
 public:
 
    int WriteData(const char *buf, int nbytes);
 
    int ReadData(char *buf, int nbytes);
 
    void Flush();
 
    void Pong(uint32_t len);
 
    inline WS_State GetState() { return m_state; }
 
    int Ping(uint32_t len, uint32_t *sentTick);
 
    int WaitForPingReply(uint32_t timeout);
 
    int GetPingReplyTick(uint32_t *replyTick);
 
    int setoption(int option);
 
    int clroption(int option);
 
    int getoption();
 
    inline int GetWriteSpace() { return txBuffer.SpaceAvail(); }
 
    static int s_openSocketCount;
 
    static WebSocket WSSockStructs[WS_MAX_SOCKS];
 
    static int ws_readwto(int fd, char *buf, int nbytes, int timeout);
 
    static int ws_read(int fd, char *buf, int nbytes);
 
    static int ws_write(int fd, const char *buf, int nbytes);
 
    static int ws_externalclose(int fd);
 
    static void ws_flush(int fd);
 
    static int ws_setoption(int fd, int option);
 
    static int ws_clroption(int fd, int option);
 
    static int ws_getoption(int fd);
 
    static WebSocket *GetWebSocketRecord(int fd);
 
    static int promote_tcp_ws(int tcp_fd);
 
    static void ws_read_notify(int fd);
 
    static void ws_write_notify(int fd);
 
    static void GetFlushTime();
 
    static void RunSkippedCallBack();          // Process any skipped callbacks
 
    static int Connect(const char * host, const char * resource, const char * protocol, int portnum = 80, bool useSSL = false);
 
    static inline int Connect(const char * host, const char * resource, int portnum = 80, bool useSSL = false)
    { return Connect(host, resource, NULL, portnum, useSSL); }
 
    static int Connect(IPADDR host, const char * resource, const char * protocol, int portnum = 80, bool useSSL = false);
 
 
    static inline int Connect(IPADDR host, const char * resource, int portnum = 80, bool useSSL = false)
    { return Connect(host, resource, NULL, portnum, useSSL); }
 
    void DumpSock();
 
    static void DumpSockets();
 
    static int GetFreeWebSocketCount();
 
};  // class WebSocket
 
/* } groupWebSocketClass */
 
 
}   // namespace NB
 
 
 
int WSPing(int fd, uint32_t len, uint32_t *sentTick);
 
int WSGetPingReplyTick(int fd, uint32_t *replyTick);
 
int WSWaitForPingReply(int fd, uint32_t timeout);
 
int WSUpgrade(HTTP_Request *req, int sock);
 
 
#endif /* ----- #ifndef __WEBSOCKETS_H  ----- */