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 #include "net/test/embedded_test_server/websocket_message_assembler.h" 6 7 #include "base/containers/extend.h" 8 #include "base/containers/span.h" 9 #include "base/logging.h" 10 #include "base/numerics/safe_conversions.h" 11 #include "net/base/net_errors.h" 12 13 namespace net::test_server { 14 15 WebSocketMessageAssembler::WebSocketMessageAssembler() = default; 16 WebSocketMessageAssembler::~WebSocketMessageAssembler() = default; 17 HandleFrame(bool is_final,WebSocketFrameHeader::OpCode opcode,base::span<const char> payload)18MessageOrError WebSocketMessageAssembler::HandleFrame( 19 bool is_final, 20 WebSocketFrameHeader::OpCode opcode, 21 base::span<const char> payload) { 22 if (state_ == MessageState::kFinished) { 23 Reset(); 24 } 25 26 switch (opcode) { 27 case WebSocketFrameHeader::kOpCodeText: 28 if (state_ != MessageState::kIdle) { 29 DVLOG(1) << "Unexpected text frame while expecting continuation"; 30 return base::unexpected(ERR_WS_PROTOCOL_ERROR); 31 } 32 is_text_message_ = true; 33 break; 34 35 case WebSocketFrameHeader::kOpCodeBinary: 36 if (state_ != MessageState::kIdle) { 37 DVLOG(1) << "Unexpected binary frame while expecting continuation"; 38 return base::unexpected(ERR_WS_PROTOCOL_ERROR); 39 } 40 // Explicitly set to indicate binary handling. 41 is_text_message_ = false; 42 break; 43 44 case WebSocketFrameHeader::kOpCodeContinuation: 45 if (state_ == MessageState::kIdle) { 46 DVLOG(1) << "Unexpected continuation frame in idle state"; 47 return base::unexpected(ERR_WS_PROTOCOL_ERROR); 48 } 49 break; 50 51 default: 52 DVLOG(1) << "Invalid frame opcode: " << opcode; 53 return base::unexpected(ERR_WS_PROTOCOL_ERROR); 54 } 55 56 // If it's the final frame and we haven't received previous fragments, return 57 // the current payload directly as the message. This avoids using an internal 58 // buffer, optimizing memory usage by eliminating unnecessary copies. 59 if (is_final && multi_frame_buffer_.empty()) { 60 return Message(is_text_message_, base::as_bytes(payload)); 61 } 62 63 base::Extend(multi_frame_buffer_, base::as_byte_span(payload)); 64 65 if (is_final) { 66 Message complete_message(is_text_message_, base::span(multi_frame_buffer_)); 67 state_ = MessageState::kFinished; 68 return complete_message; 69 } 70 71 // Update the state to expect a continuation frame. 72 state_ = is_text_message_ ? MessageState::kExpectTextContinuation 73 : MessageState::kExpectBinaryContinuation; 74 return base::unexpected(ERR_IO_PENDING); 75 } 76 Reset()77void WebSocketMessageAssembler::Reset() { 78 multi_frame_buffer_.clear(); 79 state_ = MessageState::kIdle; 80 is_text_message_ = false; 81 } 82 83 } // namespace net::test_server 84