Chapter 97: WebSocket Server Implementation

idf_component_register(SRCS "main.c"
                    INCLUDE_DIRS "."
                    REQUIRES esp_wifi esp_event esp_log nvs_flash netif_stack esp_netif esp_http_server)

#include <stdio.h>
#include <string.h>
#include "esp_wifi.h"
#include "esp_event.h"
#include "esp_log.h"
#include "nvs_flash.h"
#include "esp_netif.h"
#include "esp_http_server.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "freertos/semphr.h"

#define WIFI_SSID      "ESP32-WS-Server"
#define WIFI_PASS      "password" // At least 8 characters
#define MAX_STA_CONN       4

static const char *TAG = "WS_SERVER_EXAMPLE";

// Structure to manage client sessions (socket file descriptors)
#define MAX_WS_CLIENTS 5
static int client_sockets[MAX_WS_CLIENTS];
static uint8_t client_count = 0;
static SemaphoreHandle_t client_list_mutex;


// Function to add a client socket to the list
static void add_client_socket(int sockfd) {
    xSemaphoreTake(client_list_mutex, portMAX_DELAY);
    if (client_count < MAX_WS_CLIENTS) {
        client_sockets[client_count++] = sockfd;
        ESP_LOGI(TAG, "Client socket %d added. Total clients: %d", sockfd, client_count);
    } else {
        ESP_LOGW(TAG, "Max clients reached. Cannot add socket %d", sockfd);
        // Optionally close the new socket if list is full
        close(sockfd);
    }
    xSemaphoreGive(client_list_mutex);
}

// Function to remove a client socket from the list
static void remove_client_socket(int sockfd) {
    xSemaphoreTake(client_list_mutex, portMAX_DELAY);
    int i, j;
    for (i = 0; i < client_count; i++) {
        if (client_sockets[i] == sockfd) {
            // Shift elements to fill the gap
            for (j = i; j < client_count - 1; j++) {
                client_sockets[j] = client_sockets[j + 1];
            }
            client_count--;
            ESP_LOGI(TAG, "Client socket %d removed. Total clients: %d", sockfd, client_count);
            break; // Exit once found and removed
        }
    }
    xSemaphoreGive(client_list_mutex);
}


// WebSocket open handler (called after handshake)
// This is where you'd typically add the client to a list
static esp_err_t on_ws_open(httpd_req_t *req) {
    int sockfd = httpd_req_to_sockfd(req);
    if (sockfd < 0) {
        ESP_LOGE(TAG, "Failed to get socket descriptor for new WebSocket connection");
        return ESP_FAIL;
    }
    ESP_LOGI(TAG, "New WebSocket connection opened, sockfd: %d", sockfd);
    add_client_socket(sockfd);
    // You can set a user context for this session if needed
    // httpd_sess_set_ctx(req->handle, sockfd, your_custom_context, free_context_function);
    return ESP_OK;
}


// WebSocket close handler
// This is called when the HTTP server detects the socket is closed by client or error
// Or when httpd_sess_close is called by the server
static void on_ws_close(httpd_req_t *req) {
    int sockfd = httpd_req_to_sockfd(req);
    ESP_LOGI(TAG, "WebSocket connection closed, sockfd: %d", sockfd);
    remove_client_socket(sockfd);
    // No need to close(sockfd) here, http_server component handles it
    // after this handler returns, or if this handler is not registered.
}


/*
 * Structure holding server handle
 * and task handle of a task sending server side events
 */
struct async_resp_arg {
    httpd_handle_t hd;
    int fd;
};


// Handler for WebSocket messages
static esp_err_t ws_handler(httpd_req_t *req) {
    // First, check if this is an open or close event if using ws_handler directly
    // For this example, we use dedicated open/close handlers via httpd_uri_t config

    if (req->method == HTTP_GET) {
        // This part is for the HTTP GET request that initiates the WebSocket handshake
        ESP_LOGI(TAG, "Handshake done, new connection was opened");
        // The on_ws_open handler (if registered via .open_fn) would have been called
        // If not using .open_fn, you'd add client sockfd here.
        // For this example, .open_fn handles adding the client.
        return ESP_OK;
    }

    // If it's not a GET, it's a data frame (or the server is cleaning up a closed session)
    httpd_ws_frame_t ws_pkt;
    uint8_t *buf = NULL;
    memset(&ws_pkt, 0, sizeof(httpd_ws_frame_t));
    ws_pkt.type = HTTPD_WS_TYPE_TEXT; // Expecting text

    // Set max_len = 0 to get the frame len
    esp_err_t ret = httpd_ws_recv_frame(req, &ws_pkt, 0);
    if (ret != ESP_OK) {
        ESP_LOGE(TAG, "httpd_ws_recv_frame failed to get frame len with %d", ret);
        return ret;
    }
    ESP_LOGI(TAG, "frame len is %d", ws_pkt.len);

    if (ws_pkt.len) {
        buf = calloc(1, ws_pkt.len + 1); // +1 for null terminator
        if (buf == NULL) {
            ESP_LOGE(TAG, "Failed to calloc memory for WebSocket frame");
            return ESP_ERR_NO_MEM;
        }
        ws_pkt.payload = buf;
        // Set max_len = ws_pkt.len to get the frame payload
        ret = httpd_ws_recv_frame(req, &ws_pkt, ws_pkt.len);
        if (ret != ESP_OK) {
            ESP_LOGE(TAG, "httpd_ws_recv_frame failed with %d", ret);
            free(buf);
            return ret;
        }
        // Ensure null termination for text messages
        buf[ws_pkt.len] = 0;

        if (ws_pkt.type == HTTPD_WS_TYPE_TEXT) {
            ESP_LOGI(TAG, "Received Pkt from sockfd %d: type=%d, len=%d, data=%s",
                     httpd_req_to_sockfd(req), ws_pkt.type, ws_pkt.len, (char*)ws_pkt.payload);
            
            // Echo back the received message
            ret = httpd_ws_send_frame(req, &ws_pkt);
            if (ret != ESP_OK) {
                ESP_LOGE(TAG, "httpd_ws_send_frame failed with %d", ret);
            }
        } else if (ws_pkt.type == HTTPD_WS_TYPE_BINARY) {
             ESP_LOGI(TAG, "Received Binary Pkt from sockfd %d: len=%d", httpd_req_to_sockfd(req), ws_pkt.len);
             // Echo back binary
            ret = httpd_ws_send_frame(req, &ws_pkt);
            if (ret != ESP_OK) {
                ESP_LOGE(TAG, "httpd_ws_send_frame failed with %d", ret);
            }
        } else if (ws_pkt.type == HTTPD_WS_TYPE_CLOSE) {
            // Client sent a close frame
            ESP_LOGI(TAG, "Received Close Pkt from sockfd %d", httpd_req_to_sockfd(req));
            // The server will automatically send a close response and then
            // the on_ws_close handler will be called.
        } else {
            ESP_LOGI(TAG, "Received Pkt from sockfd %d: type=%d", httpd_req_to_sockfd(req), ws_pkt.type);
        }
        free(buf);
    }
    return ret;
}


// URI handler structure for WebSocket
static const httpd_uri_t uri_ws = {
    .uri        = "/ws",
    .method     = HTTP_GET, // WebSocket handshake is a GET request
    .handler    = ws_handler, // Handles data frames after handshake
    .user_ctx   = NULL,
    .is_websocket = true,
    .handle_ws_control_frames = true, // HTTP server handles PING/PONG
    .supported_subprotocol = "chat", // Optional: specify a subprotocol
    .open_fn    = on_ws_open,   // Called when new WS connection is opened
    .close_fn   = on_ws_close   // Called when WS connection is closed
};

// Function to start the WebSocket server
static httpd_handle_t start_webserver(void) {
    httpd_handle_t server = NULL;
    httpd_config_t config = HTTPD_DEFAULT_CONFIG();
    config.max_open_sockets = MAX_WS_CLIENTS + 3; // Max WS clients + some buffer for HTTP
    config.lru_purge_enable = true; // Enable LRU purge of sockets

    ESP_LOGI(TAG, "Starting server on port: '%d'", config.server_port);
    if (httpd_start(&server, &config) == ESP_OK) {
        ESP_LOGI(TAG, "Registering URI handlers");
        httpd_register_uri_handler(server, &uri_ws);
        return server;
    }

    ESP_LOGI(TAG, "Error starting server!");
    return NULL;
}

// Function to stop the server
static void stop_webserver(httpd_handle_t server) {
    if (server) {
        httpd_stop(server);
    }
}

// Wi-Fi event handler
static void wifi_event_handler(void* arg, esp_event_base_t event_base,
                               int32_t event_id, void* event_data) {
    if (event_id == WIFI_EVENT_AP_STACONNECTED) {
        wifi_event_ap_staconnected_t* event = (wifi_event_ap_staconnected_t*) event_data;
        ESP_LOGI(TAG, "Station "MACSTR" joined, AID=%d", MAC2STR(event->mac), event->aid);
    } else if (event_id == WIFI_EVENT_AP_STADISCONNECTED) {
        wifi_event_ap_stadisconnected_t* event = (wifi_event_ap_stadisconnected_t*) event_data;
        ESP_LOGI(TAG, "Station "MACSTR" left, AID=%d", MAC2STR(event->mac), event->aid);
    }
}

// Wi-Fi AP Initialization
void wifi_init_softap(void) {
    ESP_ERROR_CHECK(esp_netif_init());
    ESP_ERROR_CHECK(esp_event_loop_create_default());
    esp_netif_t *p_netif = esp_netif_create_default_wifi_ap();
    assert(p_netif);

    wifi_init_config_t cfg = WIFI_INIT_CONFIG_DEFAULT();
    ESP_ERROR_CHECK(esp_wifi_init(&cfg));

    ESP_ERROR_CHECK(esp_event_handler_instance_register(WIFI_EVENT,
                                                        ESP_EVENT_ANY_ID,
                                                        &wifi_event_handler,
                                                        NULL,
                                                        NULL));

    wifi_config_t wifi_config = {
        .ap = {
            .ssid = WIFI_SSID,
            .ssid_len = strlen(WIFI_SSID),
            .password = WIFI_PASS,
            .max_connection = MAX_STA_CONN,
            .authmode = WIFI_AUTH_WPA_WPA2_PSK,
            .pmf_cfg = {
                .required = false,
            },
        },
    };
    if (strlen(WIFI_PASS) == 0) {
        wifi_config.ap.authmode = WIFI_AUTH_OPEN;
    }

    ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_AP));
    ESP_ERROR_CHECK(esp_wifi_set_config(WIFI_IF_AP, &wifi_config));
    ESP_ERROR_CHECK(esp_wifi_start());

    ESP_LOGI(TAG, "wifi_init_softap finished. SSID:%s password:%s", WIFI_SSID, WIFI_PASS);

    esp_netif_ip_info_t ip_info;
    esp_netif_get_ip_info(p_netif, &ip_info);
    ESP_LOGI(TAG, "AP IP Address: " IPSTR, IP2STR(&ip_info.ip));
}


void app_main(void) {
    // Initialize NVS
    esp_err_t ret = nvs_flash_init();
    if (ret == ESP_ERR_NVS_NO_FREE_PAGES || ret == ESP_ERR_NVS_NEW_VERSION_FOUND) {
      ESP_ERROR_CHECK(nvs_flash_erase());
      ret = nvs_flash_init();
    }
    ESP_ERROR_CHECK(ret);

    client_list_mutex = xSemaphoreCreateMutex();
    if (client_list_mutex == NULL) {
        ESP_LOGE(TAG, "Failed to create client list mutex");
        return;
    }

    ESP_LOGI(TAG, "ESP_WEBSOCKET_SERVER_EXAMPLE");

    wifi_init_softap(); // Initialize Wi-Fi in SoftAP mode
    start_webserver();  // Start the HTTP/WebSocket server
}

// ... (Keep previous includes, defines, TAG, client_sockets, client_count, client_list_mutex, 
//      add_client_socket, remove_client_socket, on_ws_open, on_ws_close, ws_handler, uri_ws, 
//      stop_webserver, wifi_event_handler, wifi_init_softap)

// Global server handle for broadcasting task
static httpd_handle_t server_handle = NULL;

// Task to broadcast messages
static void broadcast_task(void *pvParameters) {
    char broadcast_msg[50];
    int counter = 0;

    while (1) {
        vTaskDelay(pdMS_TO_TICKS(5000)); // Broadcast every 5 seconds

        if (server_handle == NULL || client_count == 0) {
            //ESP_LOGI(TAG, "No server or no clients, skipping broadcast");
            continue;
        }

        sprintf(broadcast_msg, "Server broadcast: Hello #%d", counter++);
        ESP_LOGI(TAG, "Broadcasting to %d clients: %s", client_count, broadcast_msg);

        httpd_ws_frame_t ws_pkt;
        memset(&ws_pkt, 0, sizeof(httpd_ws_frame_t));
        ws_pkt.payload = (uint8_t*)broadcast_msg;
        ws_pkt.len = strlen(broadcast_msg);
        ws_pkt.type = HTTPD_WS_TYPE_TEXT;
        // ws_pkt.final = true; // This is implicitly true for non-fragmented frames

        xSemaphoreTake(client_list_mutex, portMAX_DELAY);
        for (int i = 0; i < client_count; i++) {
            int client_sockfd = client_sockets[i];
            
            // httpd_ws_send_frame_async requires the server handle and client sockfd
            esp_err_t ret = httpd_ws_send_frame_async(server_handle, client_sockfd, &ws_pkt);
            if (ret != ESP_OK) {
                ESP_LOGE(TAG, "httpd_ws_send_frame_async failed for sockfd %d with error: %s", client_sockfd, esp_err_to_name(ret));
                // Consider removing client if send fails persistently
            } else {
                ESP_LOGI(TAG, "Sent to sockfd %d", client_sockfd);
            }
        }
        xSemaphoreGive(client_list_mutex);
    }
}

// Modify start_webserver to store the handle
static httpd_handle_t start_webserver(void) {
    //httpd_handle_t server = NULL; // server_handle is now global
    httpd_config_t config = HTTPD_DEFAULT_CONFIG();
    config.max_open_sockets = MAX_WS_CLIENTS + 3;
    config.lru_purge_enable = true;

    ESP_LOGI(TAG, "Starting server on port: '%d'", config.server_port);
    if (httpd_start(&server_handle, &config) == ESP_OK) { // Use global server_handle
        ESP_LOGI(TAG, "Registering URI handlers");
        httpd_register_uri_handler(server_handle, &uri_ws);
        return server_handle;
    }

    ESP_LOGI(TAG, "Error starting server!");
    server_handle = NULL;
    return NULL;
}


void app_main(void) {
    esp_err_t ret = nvs_flash_init();
    if (ret == ESP_ERR_NVS_NO_FREE_PAGES || ret == ESP_ERR_NVS_NEW_VERSION_FOUND) {
      ESP_ERROR_CHECK(nvs_flash_erase());
      ret = nvs_flash_init();
    }
    ESP_ERROR_CHECK(ret);

    client_list_mutex = xSemaphoreCreateMutex();
    if (client_list_mutex == NULL) {
        ESP_LOGE(TAG, "Failed to create client list mutex");
        return;
    }
    memset(client_sockets, 0, sizeof(client_sockets)); // Initialize client sockets array

    ESP_LOGI(TAG, "ESP_WEBSOCKET_SERVER_BROADCAST_EXAMPLE");

    wifi_init_softap();
    if (start_webserver() != NULL) {
        xTaskCreate(broadcast_task, "broadcast_task", 4096, NULL, 5, NULL);
    }
}

openssl req -newkey rsa:2048 -nodes -keyout prvtkey.pem -x509 -days 365 -out servercert.pem -subj "/CN=esp32-wss-server"

#ifndef MAIN_SERVER_CERT_H_
#define MAIN_SERVER_CERT_H_

// -----BEGIN CERTIFICATE-----
// ... (Your server certificate PEM data) ...
// -----END CERTIFICATE-----
const unsigned char server_cert_pem_start[] = "-----BEGIN CERTIFICATE-----\n"
    // ... content of your servercert.pem ...
    "-----END CERTIFICATE-----\n";
const unsigned int server_cert_pem_len = sizeof(server_cert_pem_start);

// -----BEGIN PRIVATE KEY-----
// ... (Your private key PEM data) ...
// -----END PRIVATE KEY-----
const unsigned char server_key_pem_start[] = "-----BEGIN PRIVATE KEY-----\n"
    // ... content of your prvtkey.pem ...
    "-----END PRIVATE KEY-----\n";
const unsigned int server_key_pem_len = sizeof(server_key_pem_start);

#endif /* MAIN_SERVER_CERT_H_ */

// ... (include server_cert.h)
#include "server_cert.h" // Assuming you created this file

// Modify start_webserver for WSS
static httpd_handle_t start_secure_webserver(void) {
    httpd_config_t config = HTTPD_DEFAULT_CONFIG();
    // For WSS, the default port is 443. HTTPD_DEFAULT_CONFIG() sets port 80.
    // If you want to use a different port for HTTPS/WSS, set config.server_port
    // For WSS, the URI scheme will be wss://<ip>:<port>/ws
    // If using standard port 443, then wss://<ip>/ws
    config.server_port = 443; // Standard HTTPS/WSS port
    config.ctrl_port = 32769; // Different from server_port

    config.max_open_sockets = MAX_WS_CLIENTS + 3;
    config.lru_purge_enable = true;

    // SSL configuration
    httpd_ssl_config_t ssl_config = HTTPD_SSL_CONFIG_DEFAULT();
    ssl_config.servercert = server_cert_pem_start;
    ssl_config.servercert_len = server_cert_pem_len;
    ssl_config.prvtkey_pem = server_key_pem_start;
    ssl_config.prvtkey_len = server_key_pem_len;
    // For production, you might want to set httpd_ssl_config.transport_mode and .port_secure

    ESP_LOGI(TAG, "Starting WSS server on port: '%d'", config.server_port);
    // The httpd_ssl_start function is deprecated. Use httpd_start with ssl_config.
    // esp_err_t ret = httpd_ssl_start(&server_handle, &config, &ssl_config);
    // The modern way is to pass ssl_config to httpd_start via config.global_transport_ctx
    // and config.global_transport_ctx_free_fn.
    // However, for basic SSL, httpd_config_t has direct members for this in later IDF versions.
    // Let's check ESP-IDF v5.x httpd_config_t:
    // It has .server_port and .ctrl_port.
    // For SSL, we pass the ssl_config to httpd_start.
    // No, httpd_ssl_config is separate. httpd_start takes httpd_config_t.
    // The function to start an SSL server is httpd_ssl_start in older IDFs.
    // In ESP-IDF v5.x, you use httpd_start and provide SSL config through a different mechanism
    // if httpd_ssl_start is fully removed.
    //
    // Re-checking ESP-IDF v5.x docs for http_server with SSL:
    // You create an httpd_config_t.
    // Then, to start an HTTPS server, you call httpd_ssl_start(&server_handle, &config).
    // This function is indeed available in v5.0, v5.1, v5.2.
    // So, the old way is still the way for esp_http_server with SSL.
    
    // Correction: The `httpd_ssl_config_t` is now part of `httpd_config_t` as `global_user_ctx`
    // and you set a specific session_creation_cb or use `HTTPD_SSL_CONFIG_DEFAULT()` which populates these.
    // The `httpd_ssl_start` and `httpd_ssl_stop` are indeed the functions to use.

    // Let's use the recommended way from ESP-IDF examples for v5.x if they've changed.
    // The `esp_http_server` component's `CMakeLists.txt` might conditionally compile SSL support.
    // Ensure `CONFIG_ESP_TLS_SERVER` is enabled in menuconfig if not by default.

    // The httpd_ssl_config_t structure is used with httpd_ssl_start()
    // This is the correct approach for ESP-IDF v5.x based on current API docs.
    
    esp_err_t ret = httpd_ssl_start(&server_handle, &ssl_config); // Pass the httpd_config_t via ssl_config.base_cfg
    // No, the httpd_ssl_config_t has a base_path member, not base_cfg.
    // The httpd_ssl_config_t IS the config for httpd_ssl_start.
    // Let's assume httpd_config_t is implicitly used for port etc.
    // The `httpd_ssl_config_t` contains the `httpd_config_t` as `httpd_config_t base;`
    // So we need to populate that.

    httpd_ssl_config_t https_conf = HTTPD_SSL_CONFIG_DEFAULT(); // This macro populates base config too.
    https_conf.httpd.max_open_sockets = MAX_WS_CLIENTS + 3;
    https_conf.httpd.lru_purge_enable = true;
    https_conf.httpd.server_port = 443; // Default for HTTPD_SSL_CONFIG_DEFAULT is 443
    https_conf.httpd.ctrl_port = 32769; // Ensure this is different

    https_conf.servercert = server_cert_pem_start;
    https_conf.servercert_len = server_cert_pem_len;
    https_conf.prvtkey_pem = server_key_pem_start;
    https_conf.prvtkey_len = server_key_pem_len;

    ESP_LOGI(TAG, "Starting WSS server on port: '%d'", https_conf.httpd.server_port);
    if (httpd_ssl_start(&server_handle, &https_conf) == ESP_OK) {
        ESP_LOGI(TAG, "Registering URI handlers for WSS");
        httpd_register_uri_handler(server_handle, &uri_ws); // Same URI handler works
        return server_handle;
    }

    ESP_LOGE(TAG, "Error starting secure server!");
    server_handle = NULL;
    return NULL;
}

// In app_main, call start_secure_webserver()
void app_main(void) {
    // ... (NVS init, mutex init, logging) ...
    wifi_init_softap();
    // server_handle = start_webserver(); // For WS
    server_handle = start_secure_webserver(); // For WSS

    if (server_handle != NULL) {
        xTaskCreate(broadcast_task, "broadcast_task", 4096, NULL, 5, NULL);
    }
}