1 // Copyright 2024 The Chromium Authors 2 // Use of this source code is governed by a BSD-style license that can be 3 // found in the LICENSE file. 4 5 #ifndef NET_TEST_EMBEDDED_TEST_SERVER_WEBSOCKET_MESSAGE_ASSEMBLER_H_ 6 #define NET_TEST_EMBEDDED_TEST_SERVER_WEBSOCKET_MESSAGE_ASSEMBLER_H_ 7 8 #include <string_view> 9 #include <vector> 10 11 #include "base/containers/span.h" 12 #include "base/memory/raw_span.h" 13 #include "base/types/expected.h" 14 #include "net/base/net_errors.h" 15 #include "net/websockets/websocket_frame.h" 16 17 namespace net::test_server { 18 19 struct Message { 20 bool is_text_message; 21 // `body` either points into the `payload` passed to `HandleFrame()` or into 22 // storage owned by the WebSocketMessageAssembler object. It will be 23 // invalidated by the next call to `HandleFrame()`. 24 base::raw_span<const uint8_t> body; 25 }; 26 27 using MessageOrError = base::expected<Message, net::Error>; 28 29 class WebSocketMessageAssembler final { 30 public: 31 WebSocketMessageAssembler(); 32 ~WebSocketMessageAssembler(); 33 34 // Handles incoming WebSocket frames and assembles messages. 35 // If `final` is true and the message is complete, it returns a `Message`. 36 // If more frames are expected, it returns `ERR_IO_PENDING`. 37 // Possible errors are `ERR_IO_PENDING` and `ERR_WS_PROTOCOL_ERROR`. 38 // Note: Validation of text messages as UTF-8 is the responsibility of the 39 // client. 40 MessageOrError HandleFrame(bool final, 41 WebSocketFrameHeader::OpCode opcode, 42 base::span<const char> payload); 43 44 // Resets internal state when a message is fully processed or in case of 45 // errors. 46 void Reset(); 47 48 private: 49 // Buffer to hold partial frames for multi-frame messages. 50 std::vector<uint8_t> multi_frame_buffer_; 51 52 // State to track if we are expecting a continuation frame. 53 enum class MessageState { 54 // kIdle: No message is being processed. 55 kIdle, 56 kExpectTextContinuation, 57 kExpectBinaryContinuation, 58 kFinished 59 } state_ = MessageState::kIdle; 60 61 bool is_text_message_ = false; 62 }; 63 } // namespace net::test_server 64 65 #endif // NET_TEST_EMBEDDED_TEST_SERVER_WEBSOCKET_MESSAGE_ASSEMBLER_H_ 66