From 5e1e5f8be90f23201d6884f53b5031a8dfd80165 Mon Sep 17 00:00:00 2001
From: David Cermak <cermak@espressif.com>
Date: Mon, 7 Sep 2020 16:24:26 +0200
Subject: [PATCH] http_server: Add a flag to enable using control frames in
 user handlers

---
 components/esp_http_server/include/esp_http_server.h        | 6 ++++++
 components/esp_http_server/src/esp_httpd_priv.h             | 1 +
 components/esp_http_server/src/httpd_parse.c                | 5 +++--
 components/esp_http_server/src/httpd_uri.c                  | 2 ++
 .../https_server/wss_server/main/wss_server_example.c       | 3 ++-
 5 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/components/esp_http_server/include/esp_http_server.h b/components/esp_http_server/include/esp_http_server.h
index 6765e7d19d..cfbe0efee2 100644
--- a/components/esp_http_server/include/esp_http_server.h
+++ b/components/esp_http_server/include/esp_http_server.h
@@ -406,6 +406,12 @@ typedef struct httpd_uri {
      * If this flag is true, then method must be HTTP_GET. Otherwise the handshake will not be handled.
      */
     bool is_websocket;
+
+    /**
+     * Flag indicating that control frames (PING, PONG, CLOSE) are also passed to the handler
+     * This is used if a custom processing of the control frames is needed
+     */
+    bool handle_ws_control_frames;
 #endif
 } httpd_uri_t;
 
diff --git a/components/esp_http_server/src/esp_httpd_priv.h b/components/esp_http_server/src/esp_httpd_priv.h
index 27a3051ac2..554f0e69bf 100644
--- a/components/esp_http_server/src/esp_httpd_priv.h
+++ b/components/esp_http_server/src/esp_httpd_priv.h
@@ -75,6 +75,7 @@ struct sock_db {
     bool ws_handshake_done;                 /*!< True if it has done WebSocket handshake (if this socket is a valid WS) */
     bool ws_close;                          /*!< Set to true to close the socket later (when WS Close frame received) */
     esp_err_t (*ws_handler)(httpd_req_t *r);   /*!< WebSocket handler, leave to null if it's not WebSocket */
+    bool ws_control_frames;                         /*!< WebSocket flag indicating that control frames should be passed to user handlers */
 #endif
 };
 
diff --git a/components/esp_http_server/src/httpd_parse.c b/components/esp_http_server/src/httpd_parse.c
index e6aaa3c4c6..cc328a9e66 100644
--- a/components/esp_http_server/src/httpd_parse.c
+++ b/components/esp_http_server/src/httpd_parse.c
@@ -772,8 +772,9 @@ esp_err_t httpd_req_new(struct httpd_data *hd, struct sock_db *sd)
             ESP_LOGD(TAG, LOG_FMT("Received PONG frame"));
         }
 
-        /* Call handler if it's a non-control frame */
-        if (ret == ESP_OK && ra->ws_type <= HTTPD_WS_TYPE_PONG) {
+        /* Call handler if it's a non-control frame (or if handler requests control frames, as well) */
+        if (ret == ESP_OK &&
+            (ra->ws_type < HTTPD_WS_TYPE_CLOSE || sd->ws_control_frames)) {
             ret = sd->ws_handler(r);
         }
 
diff --git a/components/esp_http_server/src/httpd_uri.c b/components/esp_http_server/src/httpd_uri.c
index c3d559ad2a..cdc93df506 100644
--- a/components/esp_http_server/src/httpd_uri.c
+++ b/components/esp_http_server/src/httpd_uri.c
@@ -174,6 +174,7 @@ esp_err_t httpd_register_uri_handler(httpd_handle_t handle,
             hd->hd_calls[i]->user_ctx = uri_handler->user_ctx;
 #ifdef CONFIG_HTTPD_WS_SUPPORT
             hd->hd_calls[i]->is_websocket = uri_handler->is_websocket;
+            hd->hd_calls[i]->handle_ws_control_frames = uri_handler->handle_ws_control_frames;
 #endif
             ESP_LOGD(TAG, LOG_FMT("[%d] installed %s"), i, uri_handler->uri);
             return ESP_OK;
@@ -322,6 +323,7 @@ esp_err_t httpd_uri(struct httpd_data *hd)
 
         aux->sd->ws_handshake_done = true;
         aux->sd->ws_handler = uri->handler;
+        aux->sd->ws_control_frames = uri->handle_ws_control_frames;
 
         /* Return immediately after handshake, no need to call handler here */
         return ESP_OK;
diff --git a/examples/protocols/https_server/wss_server/main/wss_server_example.c b/examples/protocols/https_server/wss_server/main/wss_server_example.c
index 5f1d2978c2..fe2954640f 100644
--- a/examples/protocols/https_server/wss_server/main/wss_server_example.c
+++ b/examples/protocols/https_server/wss_server/main/wss_server_example.c
@@ -84,7 +84,8 @@ static const httpd_uri_t ws = {
         .method     = HTTP_GET,
         .handler    = ws_handler,
         .user_ctx   = NULL,
-        .is_websocket = true
+        .is_websocket = true,
+        .handle_ws_control_frames = true
 };