Harden HTTP request body and upload handling

main/http_server/websocket.c

@@ -1,3 +1,6 @@
+#include <string.h>
+#include <strings.h>
+
 #include <stdint.h>
 #include <unistd.h>
 #include "freertos/FreeRTOS.h"
@@ -12,6 +15,8 @@
 #include "log_buffer.h"
 
 #define WS_LOG_SCRATCH_SIZE 2048
+#define WS_MAX_WEBSOCKET_PAYLOAD_SIZE 1024U
+#define WS_HANDSHAKE_HEADER_SIZE 256
 
 static const char * TAG = "websocket";
 
@@ -26,6 +31,90 @@ static SemaphoreHandle_t clients_mutex = NULL;
 static httpd_handle_t server_handle = NULL;
 static TaskHandle_t s_websocket_log_task_handle = NULL;
 
+static bool websocket_payload_fits(size_t payload_len)
+{
+    return payload_len <= WS_MAX_WEBSOCKET_PAYLOAD_SIZE;
+}
+
+static bool websocket_origin_matches_host(const char *origin, const char *host)
+{
+    if (origin == NULL || host == NULL || host[0] == 0) {
+        return false;
+    }
+
+    const char *authority = NULL;
+    static const char http_prefix[] = "http://";
+    static const char https_prefix[] = "https://";
+
+    if (strncasecmp(origin, http_prefix, sizeof(http_prefix) - 1) == 0) {
+        authority = origin + sizeof(http_prefix) - 1;
+    } else if (strncasecmp(origin, https_prefix, sizeof(https_prefix) - 1) == 0) {
+        authority = origin + sizeof(https_prefix) - 1;
+    } else {
+        return false;
+    }
+
+    size_t authority_len = strcspn(authority, "/?#");
+    if (authority_len == 0 || authority[authority_len] != 0) {
+        return false;
+    }
+
+    size_t host_len = strlen(host);
+    return authority_len == host_len && strncasecmp(authority, host, host_len) == 0;
+}
+
+
+static bool websocket_has_free_slot(void)
+{
+    if (clients_mutex == NULL ||
+        xSemaphoreTake(clients_mutex, pdMS_TO_TICKS(100)) != pdTRUE) {
+        ESP_LOGE(TAG, "Failed to acquire mutex while checking client capacity");
+        return false;
+    }
+
+    bool has_free_slot = false;
+    for (int i = 0; i < MAX_WEBSOCKET_CLIENTS; i++) {
+        if (clients[i].fd == -1) {
+            has_free_slot = true;
+            break;
+        }
+    }
+
+    xSemaphoreGive(clients_mutex);
+    return has_free_slot;
+}
+
+static esp_err_t websocket_origin_is_allowed(httpd_req_t *req)
+{
+    size_t origin_len = httpd_req_get_hdr_value_len(req, "Origin");
+    if (origin_len == 0) {
+        // Non-browser clients such as websocat do not necessarily send Origin.
+        return ESP_OK;
+    }
+
+    size_t host_len = httpd_req_get_hdr_value_len(req, "Host");
+    if (origin_len >= WS_HANDSHAKE_HEADER_SIZE || host_len == 0 ||
+        host_len >= WS_HANDSHAKE_HEADER_SIZE) {
+        ESP_LOGW(TAG, "Rejecting WebSocket handshake with invalid Origin/Host length");
+        return ESP_FAIL;
+    }
+
+    char origin[WS_HANDSHAKE_HEADER_SIZE];
+    char host[WS_HANDSHAKE_HEADER_SIZE];
+    if (httpd_req_get_hdr_value_str(req, "Origin", origin, sizeof(origin)) != ESP_OK ||
+        httpd_req_get_hdr_value_str(req, "Host", host, sizeof(host)) != ESP_OK) {
+        ESP_LOGW(TAG, "Rejecting WebSocket handshake with unreadable Origin/Host");
+        return ESP_FAIL;
+    }
+
+    if (!websocket_origin_matches_host(origin, host)) {
+        ESP_LOGW(TAG, "Rejecting cross-origin WebSocket handshake");
+        return ESP_FAIL;
+    }
+
+    return ESP_OK;
+}
+
 void websocket_set_log_task_handle(TaskHandle_t task_handle)
 {
     s_websocket_log_task_handle = task_handle;
@@ -175,17 +264,28 @@ void websocket_init(httpd_handle_t server)
 
 esp_err_t websocket_pre_handshake(httpd_req_t *req)
 {
+    if (websocket_origin_is_allowed(req) != ESP_OK) {
+        httpd_resp_send_err(req, HTTPD_403_FORBIDDEN,
+                            "Forbidden WebSocket origin");
+        return ESP_FAIL;
+    }
+
     if (is_network_allowed(req) != ESP_OK) {
         httpd_resp_send_err(req, HTTPD_401_UNAUTHORIZED, "Unauthorized");
         return ESP_FAIL;
     }
 
-    int active_clients = 0;
-    for (int i = 0; i < WS_TYPE_MAX; i++) {
-        active_clients += type_counts[i];
+    WebSocketClientType type = (WebSocketClientType)(uintptr_t)req->user_ctx;
+    if (type < 0 || type >= WS_TYPE_MAX) {
+        ESP_LOGE(TAG, "Rejecting WebSocket connection with invalid client type: %d",
+                 type);
+        httpd_resp_send_err(req, HTTPD_500_INTERNAL_SERVER_ERROR,
+                            "Invalid WebSocket endpoint");
+        return ESP_FAIL;
     }
-    if (active_clients >= MAX_WEBSOCKET_CLIENTS) {
-        ESP_LOGE(TAG, "Max WebSocket clients reached, rejecting new connection");
+
+    if (!websocket_has_free_slot()) {
+        ESP_LOGW(TAG, "Max WebSocket clients reached, rejecting handshake");
         httpd_resp_send_custom_err(req, "429 Too Many Requests", "Max WebSocket clients reached");
         return ESP_FAIL;
     }
@@ -221,15 +321,18 @@ esp_err_t websocket_handler(httpd_req_t *req)
         return ret;
     }
 
-    // If there's a payload, drain it
+    // Inbound application data is ignored, but it must be drained to keep the
+    // WebSocket stream synchronized. Never allocate based on a peer-provided
+    // frame length.
     if (ws_pkt.len > 0) {
-        uint8_t *buf = (uint8_t *)calloc(1, ws_pkt.len + 1);
-        if (buf) {
-            ws_pkt.payload = buf;
-            ret = httpd_ws_recv_frame(req, &ws_pkt, ws_pkt.len);
-            free(buf);
-            return ret;
+        if (!websocket_payload_fits(ws_pkt.len)) {
+            ESP_LOGW(TAG, "Rejecting oversized WebSocket frame: %zu bytes", ws_pkt.len);
+            return ESP_ERR_INVALID_SIZE;
         }
+
+        uint8_t buf[WS_MAX_WEBSOCKET_PAYLOAD_SIZE];
+        ws_pkt.payload = buf;
+        return httpd_ws_recv_frame(req, &ws_pkt, sizeof(buf));
     }
 
     return ESP_OK;