diff --git a/include/cinatra/coro_http_client.hpp b/include/cinatra/coro_http_client.hpp index 858f3d64..72f4f8a8 100644 --- a/include/cinatra/coro_http_client.hpp +++ b/include/cinatra/coro_http_client.hpp @@ -1835,7 +1835,6 @@ class coro_http_client : public std::enable_shared_from_this { resp_data data{}; head_buf_.consume(head_buf_.size()); - size_t header_size = 2; std::shared_ptr sock = socket_; asio::streambuf &read_buf = sock->head_buf_; bool has_init_ssl = false; @@ -1844,8 +1843,8 @@ class coro_http_client : public std::enable_shared_from_this { #endif websocket ws{}; while (true) { - if (auto [ec, _] = - co_await async_read_ws(sock, read_buf, header_size, has_init_ssl); + if (auto [ec, _] = co_await async_read_ws( + sock, read_buf, ws.left_header_len(), has_init_ssl); ec) { data.net_err = ec; data.status = 404; @@ -1859,27 +1858,31 @@ class coro_http_client : public std::enable_shared_from_this { } const char *data_ptr = asio::buffer_cast(read_buf.data()); - auto ret = ws.parse_header(data_ptr, header_size, false); - if (ret == -2) { - header_size += ws.left_header_len(); + auto ret = ws.parse_header(data_ptr, read_buf.size(), false); + if (ret == ws_header_status::incomplete) { continue; } + else if (ret == ws_header_status::error) { + data.net_err = std::make_error_code(std::errc::protocol_error); + data.status = 404; + close_socket(*sock); + co_return data; + } + frame_header *header = (frame_header *)data_ptr; bool is_close_frame = header->opcode == opcode::close; - read_buf.consume(header_size); + read_buf.consume(read_buf.size()); size_t payload_len = ws.payload_length(); - if (payload_len > read_buf.size()) { - size_t size_to_read = payload_len - read_buf.size(); - if (auto [ec, size] = co_await async_read_ws( - sock, read_buf, size_to_read, has_init_ssl); - ec) { - data.net_err = ec; - data.status = 404; - close_socket(*sock); - co_return data; - } + + if (auto [ec, size] = + co_await async_read_ws(sock, read_buf, payload_len, has_init_ssl); + ec) { + data.net_err = ec; + data.status = 404; + close_socket(*sock); + co_return data; } data_ptr = asio::buffer_cast(read_buf.data()); @@ -1913,7 +1916,6 @@ class coro_http_client : public std::enable_shared_from_this { } #endif read_buf.consume(read_buf.size()); - header_size = 2; if (is_close_frame) { std::string reason = "close"; diff --git a/include/cinatra/websocket.hpp b/include/cinatra/websocket.hpp index e684112e..5f4e76c2 100644 --- a/include/cinatra/websocket.hpp +++ b/include/cinatra/websocket.hpp @@ -253,7 +253,7 @@ class websocket { size_t payload_length_ = 0; - size_t left_header_len_ = 0; + size_t left_header_len_ = 2; uint8_t mask_key_[4] = {}; unsigned char msg_opcode_ = 0; unsigned char msg_fin_ = 0;