| 1 | // SPDX-License-Identifier: GPL-3.0-or-later |
| 2 | |
| 3 | #include "websocket-internal.h" |
| 4 | #include "websocket-echo.h" |
| 5 | #include "websocket-jsonrpc.h" |
| 6 | #include "../mcp/adapters/mcp-websocket.h" |
| 7 | |
| 8 | ENUM_STR_MAP_DEFINE(WEBSOCKET_PROTOCOL) = { |
| 9 | { .id = WS_PROTOCOL_JSONRPC, .name = "jsonrpc" }, |
| 10 | { .id = WS_PROTOCOL_ECHO, .name = "echo" }, |
| 11 | { .id = WS_PROTOCOL_MCP, .name = "mcp" }, |
| 12 | { .id = WS_PROTOCOL_UNKNOWN, .name = "unknown" }, |
| 13 | |
| 14 | // terminator |
| 15 | { .name = NULL, .id = 0 } |
| 16 | }; |
| 17 | ENUM_STR_DEFINE_FUNCTIONS(WEBSOCKET_PROTOCOL, WS_PROTOCOL_UNKNOWN, "unknown"); |
| 18 | |
| 19 | ENUM_STR_MAP_DEFINE(WEBSOCKET_STATE) = { |
| 20 | { .id = WS_STATE_HANDSHAKE, .name = "handshake" }, |
| 21 | { .id = WS_STATE_OPEN, .name = "open" }, |
| 22 | { .id = WS_STATE_CLOSING_SERVER, .name = "closing_server" }, |
| 23 | { .id = WS_STATE_CLOSING_CLIENT, .name = "closing_client" }, |
| 24 | { .id = WS_STATE_CLOSED, .name = "closed" }, |
| 25 | |
| 26 | // terminator |
| 27 | { .name = NULL, .id = 0 } |
| 28 | }; |
| 29 | ENUM_STR_DEFINE_FUNCTIONS(WEBSOCKET_STATE, WS_STATE_CLOSED, "closed"); |
| 30 | |
| 31 | ENUM_STR_MAP_DEFINE(WEBSOCKET_OPCODE) = { |
| 32 | { .id = WS_OPCODE_CONTINUATION, .name = "continuation" }, |
| 33 | { .id = WS_OPCODE_TEXT, .name = "text" }, |
| 34 | { .id = WS_OPCODE_BINARY, .name = "binary" }, |
| 35 | { .id = WS_OPCODE_CLOSE, .name = "close" }, |
| 36 | { .id = WS_OPCODE_PING, .name = "ping" }, |
| 37 | { .id = WS_OPCODE_PONG, .name = "pong" }, |
| 38 | |
| 39 | // terminator |
| 40 | { .name = NULL, .id = 0 } |
| 41 | }; |
| 42 | ENUM_STR_DEFINE_FUNCTIONS(WEBSOCKET_OPCODE, WS_OPCODE_TEXT, "text"); |
| 43 | |
| 44 | ENUM_STR_MAP_DEFINE(WEBSOCKET_CLOSE_CODE) = { |
| 45 | // Standard WebSocket close codes |
| 46 | { .id = WS_CLOSE_NORMAL, .name = "normal" }, |
| 47 | { .id = WS_CLOSE_GOING_AWAY, .name = "going_away" }, |
| 48 | { .id = WS_CLOSE_PROTOCOL_ERROR, .name = "protocol_error" }, |
| 49 | { .id = WS_CLOSE_UNSUPPORTED_DATA, .name = "unsupported_data" }, |
| 50 | { .id = WS_CLOSE_RESERVED, .name = "reserved" }, |
| 51 | { .id = WS_CLOSE_NO_STATUS, .name = "no_status" }, |
| 52 | { .id = WS_CLOSE_ABNORMAL, .name = "abnormal" }, |
| 53 | { .id = WS_CLOSE_INVALID_PAYLOAD, .name = "invalid_payload" }, |
| 54 | { .id = WS_CLOSE_POLICY_VIOLATION, .name = "policy_violation" }, |
| 55 | { .id = WS_CLOSE_MESSAGE_TOO_BIG, .name = "message_too_big" }, |
| 56 | { .id = WS_CLOSE_EXTENSION_MISSING, .name = "extension_missing" }, |
| 57 | { .id = WS_CLOSE_INTERNAL_ERROR, .name = "internal_error" }, |
| 58 | { .id = WS_CLOSE_TLS_HANDSHAKE, .name = "tls_handshake_error" }, |
| 59 | |
| 60 | // Netdata-specific close codes |
| 61 | { .id = WS_CLOSE_NETDATA_TIMEOUT, .name = "timeout" }, |
| 62 | { .id = WS_CLOSE_NETDATA_SHUTDOWN, .name = "shutdown" }, |
| 63 | { .id = WS_CLOSE_NETDATA_REJECTED, .name = "rejected" }, |
| 64 | { .id = WS_CLOSE_NETDATA_RATE_LIMIT,.name = "rate_limit" }, |
| 65 | |
| 66 | // terminator |
| 67 | { .name = NULL, .id = 0 } |
| 68 | }; |
| 69 | ENUM_STR_DEFINE_FUNCTIONS(WEBSOCKET_CLOSE_CODE, WS_CLOSE_NORMAL, "normal"); |
| 70 | |
| 71 | // Private structure for WebSocket server state |
| 72 | struct websocket_server { |
| 73 | WS_CLIENTS_JudyLSet clients; // JudyL array of WebSocket clients |
| 74 | size_t client_id_counter; // Counter for generating unique client IDs |
| 75 | size_t active_clients; // Number of active clients |
| 76 | SPINLOCK spinlock; // Spinlock to protect the registry |
| 77 | }; |
| 78 | |
| 79 | // The global (but private) instance of the WebSocket server state |
| 80 | static struct websocket_server ws_server = (struct websocket_server){ |
| 81 | .clients = { 0 }, |
| 82 | .client_id_counter = 0, |
| 83 | .active_clients = 0, |
| 84 | .spinlock = SPINLOCK_INITIALIZER |
| 85 | }; |
| 86 | |
| 87 | // Initialize WebSocket subsystem |
| 88 | void websocket_initialize(void) { |
| 89 | // debug_flags |= D_WEBSOCKET; |
| 90 | |
| 91 | // Initialize thread system |
| 92 | websocket_threads_init(); |
| 93 | |
| 94 | // Initialize protocol handlers |
| 95 | websocket_jsonrpc_initialize(); |
| 96 | websocket_echo_initialize(); |
| 97 | mcp_websocket_adapter_initialize(); |
| 98 | |
| 99 | netdata_log_info("WebSocket server subsystem initialized"); |
| 100 | } |
| 101 | |
| 102 | // Create a new WebSocket client with a unique ID |
| 103 | NEVERNULL |
| 104 | WS_CLIENT *websocket_client_create(void) { |
| 105 | WS_CLIENT *wsc = callocz(1, sizeof(WS_CLIENT)); |
| 106 | |
| 107 | spinlock_lock(&ws_server.spinlock); |
| 108 | wsc->id = ++ws_server.client_id_counter; // Generate unique ID |
| 109 | spinlock_unlock(&ws_server.spinlock); |
| 110 | |
| 111 | wsc->connected_t = now_realtime_sec(); |
| 112 | wsc->last_activity_t = wsc->connected_t; |
| 113 | wsc->max_outbound_frame_size = WS_MAX_OUTGOING_FRAME_SIZE; // Default value |
| 114 | |
| 115 | // initialize callbacks to NULL |
| 116 | wsc->on_connect = NULL; |
| 117 | wsc->on_message = NULL; |
| 118 | wsc->on_close = NULL; |
| 119 | wsc->on_disconnect = NULL; |
| 120 | |
| 121 | // initialize the ND_SOCK with the web server's SSL context |
| 122 | nd_sock_init(&wsc->sock, netdata_ssl_web_server_ctx, false); |
| 123 | |
| 124 | // Initialize circular buffers for I/O with WebSocket-specific sizes and max limits |
| 125 | cbuffer_init(&wsc->in_buffer, WEBSOCKET_IN_BUFFER_INITIAL_SIZE, WEBSOCKET_IN_BUFFER_MAX_SIZE, NULL); |
| 126 | cbuffer_init(&wsc->out_buffer, WEBSOCKET_OUT_BUFFER_INITIAL_SIZE, WEBSOCKET_OUT_BUFFER_MAX_SIZE, NULL); |
| 127 | |
| 128 | // Initialize pre-allocated message buffer |
| 129 | wsb_init(&wsc->payload, WEBSOCKET_PAYLOAD_INITIAL_SIZE); |
| 130 | |
| 131 | // Initialize uncompressed buffer with a larger size since decompressed data can expand |
| 132 | // For compressed content, the expanded data can be much larger than the input |
| 133 | wsb_init(&wsc->u_payload, WEBSOCKET_UNPACKED_INITIAL_SIZE); |
| 134 | |
| 135 | // Initialize compressed output buffer for outbound messages |
| 136 | wsb_init(&wsc->c_payload, WEBSOCKET_PAYLOAD_INITIAL_SIZE); |
| 137 | |
| 138 | // Set the initial message state |
| 139 | wsc->opcode = WS_OPCODE_TEXT; // Default opcode |
| 140 | wsc->is_compressed = false; |
| 141 | wsc->message_complete = true; // Not in a fragmented sequence initially |
| 142 | wsc->frame_id = 0; |
| 143 | wsc->message_id = 0; |
| 144 | wsc->compression = WEBSOCKET_COMPRESSION_DEFAULTS; |
| 145 | |
| 146 | return wsc; |
| 147 | } |
| 148 | |
| 149 | // Free a WebSocket client |
| 150 | void websocket_client_free(WS_CLIENT *wsc) { |
| 151 | if (!wsc) |
| 152 | return; |
| 153 | |
| 154 | // First unregister from the client registry |
| 155 | websocket_client_unregister(wsc); |
| 156 | |
| 157 | // We MUST make sure the socket is not in the poll before closing it |
| 158 | // otherwise kernel structures may be corrupted due to socket reuse |
| 159 | if(wsc->wth && wsc->wth->ndpl && wsc->sock.fd >= 0) |
| 160 | (void) nd_poll_del(wsc->wth->ndpl, wsc->sock.fd); |
| 161 | |
| 162 | // Close socket using ND_SOCK abstraction |
| 163 | nd_sock_close(&wsc->sock); |
| 164 | |
| 165 | // Free circular buffers |
| 166 | cbuffer_cleanup(&wsc->in_buffer); |
| 167 | cbuffer_cleanup(&wsc->out_buffer); |
| 168 | |
| 169 | // Cleanup pre-allocated message, uncompressed, and compressed buffers |
| 170 | wsb_cleanup(&wsc->payload); |
| 171 | wsb_cleanup(&wsc->u_payload); |
| 172 | wsb_cleanup(&wsc->c_payload); |
| 173 | |
| 174 | // Clean up compression resources if needed |
| 175 | websocket_compression_cleanup(wsc); |
| 176 | |
| 177 | freez(wsc); |
| 178 | } |
| 179 | |
| 180 | // Register a WebSocket client in the registry |
| 181 | bool websocket_client_register(WS_CLIENT *wsc) { |
| 182 | if (!wsc || wsc->id == 0) |
| 183 | return false; |
| 184 | |
| 185 | spinlock_lock(&ws_server.spinlock); |
| 186 | |
| 187 | int added = WS_CLIENTS_SET(&ws_server.clients, wsc->id, wsc); |
| 188 | if (!added) { |
| 189 | ws_server.active_clients++; |
| 190 | websocket_debug(wsc, "WebSocket client registered, total clients: %u", ws_server.active_clients); |
| 191 | } |
| 192 | |
| 193 | spinlock_unlock(&ws_server.spinlock); |
| 194 | |
| 195 | return added; |
| 196 | } |
| 197 | |
| 198 | // Unregister a WebSocket client from the registry |
| 199 | void websocket_client_unregister(WS_CLIENT *wsc) { |
| 200 | if (!wsc || wsc->id == 0) |
| 201 | return; |
| 202 | |
| 203 | spinlock_lock(&ws_server.spinlock); |
| 204 | |
| 205 | WS_CLIENT *existing = WS_CLIENTS_GET(&ws_server.clients, wsc->id); |
| 206 | if (existing && existing == wsc) { |
| 207 | WS_CLIENTS_DEL(&ws_server.clients, wsc->id); |
| 208 | if (ws_server.active_clients > 0) |
| 209 | ws_server.active_clients--; |
| 210 | |
| 211 | websocket_debug(wsc,"WebSocket client unregistered, total clients: %zu", ws_server.active_clients); |
| 212 | } |
| 213 | |
| 214 | spinlock_unlock(&ws_server.spinlock); |
| 215 | } |
| 216 | |
| 217 | // Find a WebSocket client by ID |
| 218 | ALWAYS_INLINE |
| 219 | WS_CLIENT *websocket_client_find_by_id(size_t id) { |
| 220 | if (id == 0) |
| 221 | return NULL; |
| 222 | |
| 223 | WS_CLIENT *wsc = NULL; |
| 224 | |
| 225 | spinlock_lock(&ws_server.spinlock); |
| 226 | wsc = WS_CLIENTS_GET(&ws_server.clients, id); |
| 227 | spinlock_unlock(&ws_server.spinlock); |
| 228 | |
| 229 | return wsc; |
| 230 | } |
| 231 | |
| 232 | // Broadcast a message to all connected WebSocket clients |
| 233 | int websocket_broadcast_message(const char *message, WEBSOCKET_OPCODE opcode) { |
| 234 | if (!message || (opcode != WS_OPCODE_TEXT && opcode != WS_OPCODE_BINARY)) |
| 235 | return -1; |
| 236 | |
| 237 | int success_count = 0; |
| 238 | |
| 239 | // Send broadcast command to all active threads |
| 240 | for(size_t i = 0; i < WEBSOCKET_MAX_THREADS; i++) { |
| 241 | if(websocket_threads[i].thread && websocket_threads[i].running) { |
| 242 | if(websocket_thread_send_broadcast(&websocket_threads[i], opcode, message)) { |
| 243 | success_count++; |
| 244 | } |
| 245 | } |
| 246 | } |
| 247 | |
| 248 | return success_count; |
| 249 | } |
| 250 | |
| 251 | // Send a WebSocket message to the client |
| 252 | int websocket_send_message(WS_CLIENT *wsc, const char *message, size_t length, WEBSOCKET_OPCODE opcode) { |
| 253 | if (!wsc || !message || wsc->state != WS_STATE_OPEN) |
| 254 | return -1; |
| 255 | |
| 256 | // For other opcodes, use the generic frame sender |
| 257 | bool use_compression = wsc->compression.enabled && |
| 258 | !websocket_frame_is_control_opcode(opcode) && |
| 259 | length >= WS_COMPRESS_MIN_SIZE; |
| 260 | |
| 261 | return websocket_protocol_send_payload(wsc, message, length, opcode, use_compression); |
| 262 | } |