Add unit tests for websocket hardening

.gitignore

@@ -21,6 +21,7 @@ build/config/*
 
 # Test folder
 test/build/*
+test/managed_components/*
 test-ci/
 
 # HTML documentation

main/http_server/websocket.c

@@ -13,10 +13,9 @@
 #include "websocket_api.h"
 #include "http_server.h"
 #include "log_buffer.h"
+#include "websocket_internal.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";
 
@@ -31,12 +30,12 @@ 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)
+WEBSOCKET_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)
+WEBSOCKET_STATIC bool websocket_origin_matches_host(const char *origin, const char *host)
 {
     if (origin == NULL || host == NULL || host[0] == 0) {
         return false;
@@ -64,7 +63,7 @@ static bool websocket_origin_matches_host(const char *origin, const char *host)
 }
 
 
-static bool websocket_has_free_slot(void)
+WEBSOCKET_STATIC bool websocket_has_free_slot(void)
 {
     if (clients_mutex == NULL ||
         xSemaphoreTake(clients_mutex, pdMS_TO_TICKS(100)) != pdTRUE) {
@@ -84,7 +83,7 @@ static bool websocket_has_free_slot(void)
     return has_free_slot;
 }
 
-static esp_err_t websocket_origin_is_allowed(httpd_req_t *req)
+WEBSOCKET_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) {

main/http_server/websocket_internal.h

@@ -0,0 +1,24 @@
+#ifndef WEBSOCKET_INTERNAL_H_
+#define WEBSOCKET_INTERNAL_H_
+
+#include <stdbool.h>
+#include <stddef.h>
+#include "esp_err.h"
+#include "esp_http_server.h"
+#include "websocket.h"
+
+#define WS_MAX_WEBSOCKET_PAYLOAD_SIZE 1024U
+#define WS_HANDSHAKE_HEADER_SIZE 256
+
+#ifdef UNIT_TESTING
+#define WEBSOCKET_STATIC
+#else
+#define WEBSOCKET_STATIC static
+#endif
+
+WEBSOCKET_STATIC bool websocket_payload_fits(size_t payload_len);
+WEBSOCKET_STATIC bool websocket_origin_matches_host(const char *origin, const char *host);
+WEBSOCKET_STATIC bool websocket_has_free_slot(void);
+WEBSOCKET_STATIC esp_err_t websocket_origin_is_allowed(httpd_req_t *req);
+
+#endif /* WEBSOCKET_INTERNAL_H_ */

test-ci/main/CMakeLists.txt

@@ -1,2 +1,2 @@
 idf_component_register(SRCS "unit_test_all.c"
-                    INCLUDE_DIRS ".")
+                    INCLUDE_DIRS ".")

test/CMakeLists.txt

@@ -17,5 +17,6 @@ include($ENV{IDF_PATH}/tools/cmake/project.cmake)
 set(IDF_TARGET "esp32s3")
 
 idf_build_set_property(COMPILE_DEFINITIONS "-DCONFIG_ASIC_FREQUENCY=100" APPEND)
+idf_build_set_property(COMPILE_DEFINITIONS "-DUNIT_TESTING=1" APPEND)
 
 project(unit_test_stratum)

test/main/CMakeLists.txt

@@ -1,2 +1,3 @@
-idf_component_register(SRCS "unit_test_all.c"
-                    INCLUDE_DIRS ".")
+idf_component_register(SRCS "unit_test_all.c" "test_websocket.c"
+                    INCLUDE_DIRS "." "../../main" "../../main/http_server"
+                    REQUIRES main unity esp_http_server)

test/main/test_websocket.c

@@ -0,0 +1,73 @@
+#include <stdint.h>
+#include <stdbool.h>
+#include <string.h>
+
+#include "unity.h"
+#ifndef UNIT_TESTING
+#define UNIT_TESTING 1
+#endif
+#include "websocket_internal.h"
+
+TEST_CASE("WebSocket payload limit has strict boundary", "[websocket]")
+{
+    TEST_ASSERT_TRUE(websocket_payload_fits(0));
+    TEST_ASSERT_TRUE(websocket_payload_fits(WS_MAX_WEBSOCKET_PAYLOAD_SIZE));
+    TEST_ASSERT_FALSE(websocket_payload_fits(WS_MAX_WEBSOCKET_PAYLOAD_SIZE + 1U));
+    TEST_ASSERT_FALSE(websocket_payload_fits(SIZE_MAX));
+}
+
+TEST_CASE("WebSocket origin must match request host", "[websocket]")
+{
+    // Valid matching cases
+    TEST_ASSERT_TRUE(websocket_origin_matches_host("http://192.168.1.42", "192.168.1.42"));
+    TEST_ASSERT_TRUE(websocket_origin_matches_host("http://bitaxe.local", "bitaxe.local"));
+    TEST_ASSERT_TRUE(websocket_origin_matches_host("https://BITAXE.local", "bitaxe.LOCAL"));
+    TEST_ASSERT_TRUE(websocket_origin_matches_host("http://bitaxe.local:8080", "bitaxe.local:8080"));
+
+    // Host mismatches and cross-origin attacks
+    TEST_ASSERT_FALSE(websocket_origin_matches_host("http://evil.local", "bitaxe.local"));
+    TEST_ASSERT_FALSE(websocket_origin_matches_host("http://bitaxe.local:8080", "bitaxe.local"));
+    TEST_ASSERT_FALSE(websocket_origin_matches_host("http://bitaxe.local", "bitaxe.local:8080"));
+
+    // Untrusted/malicious origins & URL tricks
+    TEST_ASSERT_FALSE(websocket_origin_matches_host("null", "bitaxe.local"));
+    TEST_ASSERT_FALSE(websocket_origin_matches_host("file://bitaxe.local", "bitaxe.local"));
+    TEST_ASSERT_FALSE(websocket_origin_matches_host("ftp://bitaxe.local", "bitaxe.local"));
+    TEST_ASSERT_FALSE(websocket_origin_matches_host("ws://bitaxe.local", "bitaxe.local"));
+    TEST_ASSERT_FALSE(websocket_origin_matches_host("http://bitaxe.local/", "bitaxe.local"));
+    TEST_ASSERT_FALSE(websocket_origin_matches_host("http://bitaxe.local@example.com", "example.com"));
+    TEST_ASSERT_FALSE(websocket_origin_matches_host("http://bitaxe.local.attacker.com", "bitaxe.local"));
+
+    // Edge cases and NULL pointers
+    TEST_ASSERT_FALSE(websocket_origin_matches_host(NULL, "bitaxe.local"));
+    TEST_ASSERT_FALSE(websocket_origin_matches_host("http://bitaxe.local", NULL));
+    TEST_ASSERT_FALSE(websocket_origin_matches_host("http://bitaxe.local", ""));
+    TEST_ASSERT_FALSE(websocket_origin_matches_host("", "bitaxe.local"));
+}
+
+TEST_CASE("WebSocket client capacity tracking", "[websocket]")
+{
+    // After initialization, slots are empty
+    websocket_init(NULL);
+    TEST_ASSERT_TRUE(websocket_has_free_slot());
+
+    // Add maximum clients
+    for (int i = 0; i < MAX_WEBSOCKET_CLIENTS; i++) {
+        TEST_ASSERT_EQUAL(ESP_OK, websocket_add_client(100 + i, WS_TYPE_API));
+    }
+
+    // Capacity is now full
+    TEST_ASSERT_FALSE(websocket_has_free_slot());
+
+    // Adding 11th client fails
+    TEST_ASSERT_EQUAL(ESP_FAIL, websocket_add_client(999, WS_TYPE_API));
+
+    // Remove 1 client and check capacity opens up
+    websocket_remove_client(100);
+    TEST_ASSERT_TRUE(websocket_has_free_slot());
+
+    // Cleanup remaining clients
+    for (int i = 1; i < MAX_WEBSOCKET_CLIENTS; i++) {
+        websocket_remove_client(100 + i);
+    }
+}

test/sdkconfig.defaults

@@ -1,2 +1,3 @@
 CONFIG_ESP_INT_WDT=n
 CONFIG_ESP_TASK_WDT=n
+CONFIG_HTTPD_WS_SUPPORT=y