From 67339c5f4a8076c4e578e45e710bb7fd20d4a7bf Mon Sep 17 00:00:00 2001 From: mrsobakin <68982655+mrsobakin@users.noreply.github.com> Date: Fri, 21 Mar 2025 22:16:19 +0300 Subject: [PATCH] feat(websocket): Support partial frame payload reads --- .../esp_http_server/include/esp_http_server.h | 21 ++++++++++ components/esp_http_server/src/httpd_ws.c | 40 +++++++++++++------ 2 files changed, 49 insertions(+), 12 deletions(-) diff --git a/components/esp_http_server/include/esp_http_server.h b/components/esp_http_server/include/esp_http_server.h index a264c18390..51a0939eb4 100644 --- a/components/esp_http_server/include/esp_http_server.h +++ b/components/esp_http_server/include/esp_http_server.h @@ -1748,6 +1748,8 @@ typedef struct httpd_ws_frame { httpd_ws_type_t type; /*!< WebSocket frame type */ uint8_t *payload; /*!< Pre-allocated data buffer */ size_t len; /*!< Length of the WebSocket data */ + size_t left_len; /*!< Length of the WebSocket data that is yet to be received. + This field should not be modified by user. */ } httpd_ws_frame_t; /** @@ -1768,11 +1770,30 @@ typedef void (*transfer_complete_cb)(esp_err_t err, int socket, void *arg); * @return * - ESP_OK : On successful * - ESP_FAIL : Socket errors occurs + * - ESP_ERR_INVALID_SIZE : max_len is too small to fit the entire payload * - ESP_ERR_INVALID_STATE : Handshake was already done beforehand * - ESP_ERR_INVALID_ARG : Argument is invalid (null or non-WebSocket) */ esp_err_t httpd_ws_recv_frame(httpd_req_t *req, httpd_ws_frame_t *pkt, size_t max_len); +/** + * @brief Receive and parse a WebSocket frame part + * + * @note Calling httpd_ws_recv_frame_part() with max_len as 0 will give actual frame size in pkt->len. + * In contrast to httpd_ws_recv_frame, this method is able to read frame payload partially. The amount of data that is yet to be received is stored in pkt->left_len + * Keep in mind however, that you have to read the entire packet before completing the request successfully. + * + * @param[in] req Current request + * @param[out] pkt WebSocket packet + * @param[in] max_len Maximum length for receive + * @return + * - ESP_OK : On successful + * - ESP_FAIL : Socket errors occurs + * - ESP_ERR_INVALID_STATE : Handshake was already done beforehand + * - ESP_ERR_INVALID_ARG : Argument is invalid (null or non-WebSocket) + */ +esp_err_t httpd_ws_recv_frame_part(httpd_req_t *req, httpd_ws_frame_t *pkt, size_t max_len); + /** * @brief Construct and send a WebSocket frame * @param[in] req Current request diff --git a/components/esp_http_server/src/httpd_ws.c b/components/esp_http_server/src/httpd_ws.c index 949addc674..49eb409fb2 100644 --- a/components/esp_http_server/src/httpd_ws.c +++ b/components/esp_http_server/src/httpd_ws.c @@ -249,7 +249,7 @@ static esp_err_t httpd_ws_check_req(httpd_req_t *req) return ESP_OK; } -static esp_err_t httpd_ws_unmask_payload(uint8_t *payload, size_t len, const uint8_t *mask_key) +static esp_err_t httpd_ws_unmask_payload(uint8_t *payload, size_t len, const uint8_t *mask_key, size_t mask_offset) { if (len < 1 || !payload) { ESP_LOGW(TAG, LOG_FMT("Invalid payload provided")); @@ -257,13 +257,13 @@ static esp_err_t httpd_ws_unmask_payload(uint8_t *payload, size_t len, const uin } for (size_t idx = 0; idx < len; idx++) { - payload[idx] = (payload[idx] ^ mask_key[idx % 4]); + payload[idx] = (payload[idx] ^ mask_key[(idx + mask_offset) % 4]); } return ESP_OK; } -esp_err_t httpd_ws_recv_frame(httpd_req_t *req, httpd_ws_frame_t *frame, size_t max_len) +static esp_err_t httpd_ws_recv_frame_internal(httpd_req_t *req, httpd_ws_frame_t *frame, size_t max_len, bool partial) { esp_err_t ret = httpd_ws_check_req(req); if (ret != ESP_OK) { @@ -328,6 +328,8 @@ esp_err_t httpd_ws_recv_frame(httpd_req_t *req, httpd_ws_frame_t *frame, size_t ((uint64_t)length_bytes[6] << 8U) | ((uint64_t)length_bytes[7])); } + frame->left_len = frame->len; + /* If this frame is masked, dump the mask as well */ if (masked) { if (httpd_recv_with_opt(req, (char *)aux->mask_key, sizeof(aux->mask_key), HTTPD_RECV_OPT_BLOCKING) < sizeof(aux->mask_key)) { @@ -341,20 +343,23 @@ esp_err_t httpd_ws_recv_frame(httpd_req_t *req, httpd_ws_frame_t *frame, size_t return ESP_ERR_INVALID_STATE; } } - /* We only accept the incoming packet length that is smaller than the max_len (or it will overflow the buffer!) */ /* If max_len is 0, regard it OK for userspace to get frame len */ + if (max_len == 0) { + ESP_LOGD(TAG, "regard max_len == 0 is OK for user to get frame len"); + return ESP_OK; + } if (frame->len > max_len) { - if (max_len == 0) { - ESP_LOGD(TAG, "regard max_len == 0 is OK for user to get frame len"); - return ESP_OK; + /* When reading entire packet at once, we only accept the incoming packet length that is smaller than the max_len (or it will overflow the buffer!) */ + if (!partial) { + ESP_LOGW(TAG, LOG_FMT("WS Message too long")); + return ESP_ERR_INVALID_SIZE; } - ESP_LOGW(TAG, LOG_FMT("WS Message too long")); - return ESP_ERR_INVALID_SIZE; + ESP_LOGD(TAG, LOG_FMT("WS Message too long. User will have to call read again")); } /* Receive buffer */ /* If there's nothing to receive, return and stop here. */ - if (frame->len == 0) { + if (frame->left_len == 0) { return ESP_OK; } @@ -363,7 +368,7 @@ esp_err_t httpd_ws_recv_frame(httpd_req_t *req, httpd_ws_frame_t *frame, size_t return ESP_FAIL; } - size_t left_len = frame->len; + size_t left_len = (max_len < frame->left_len) ? max_len : frame->left_len; size_t offset = 0; while (left_len > 0) { @@ -378,12 +383,23 @@ esp_err_t httpd_ws_recv_frame(httpd_req_t *req, httpd_ws_frame_t *frame, size_t ESP_LOGD(TAG, "Frame length: %"NEWLIB_NANO_COMPAT_FORMAT", Bytes Read: %"NEWLIB_NANO_COMPAT_FORMAT, NEWLIB_NANO_COMPAT_CAST(frame->len), NEWLIB_NANO_COMPAT_CAST(offset)); } + size_t mask_offset = frame->len - frame->left_len; + frame->left_len -= offset; + /* Unmask payload */ - httpd_ws_unmask_payload(frame->payload, frame->len, aux->mask_key); + httpd_ws_unmask_payload(frame->payload, offset, aux->mask_key, mask_offset); return ESP_OK; } +esp_err_t httpd_ws_recv_frame(httpd_req_t *req, httpd_ws_frame_t *frame, size_t max_len) { + return httpd_ws_recv_frame_internal(req, frame, max_len, false); +} + +esp_err_t httpd_ws_recv_frame_part(httpd_req_t *req, httpd_ws_frame_t *frame, size_t max_len) { + return httpd_ws_recv_frame_internal(req, frame, max_len, true); +} + esp_err_t httpd_ws_send_frame(httpd_req_t *req, httpd_ws_frame_t *frame) { esp_err_t ret = httpd_ws_check_req(req);