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/websockets/websocket_chunk_assembler.h"
6
7 #include "base/compiler_specific.h"
8 #include "base/containers/extend.h"
9 #include "base/containers/span.h"
10 #include "base/types/expected.h"
11 #include "net/base/net_errors.h"
12 #include "net/websockets/websocket_errors.h"
13 #include "net/websockets/websocket_frame.h"
14
15 namespace net {
16
17 namespace {
18
19 // This uses type uint64_t to match the definition of
20 // WebSocketFrameHeader::payload_length in websocket_frame.h.
21 constexpr uint64_t kMaxControlFramePayload = 125;
22
23 // Utility function to create a WebSocketFrame
MakeWebSocketFrame(const WebSocketFrameHeader & header,base::span<uint8_t> payload)24 std::unique_ptr<WebSocketFrame> MakeWebSocketFrame(
25 const WebSocketFrameHeader& header,
26 base::span<uint8_t> payload) {
27 auto frame = std::make_unique<WebSocketFrame>(header.opcode);
28 frame->header.CopyFrom(header);
29
30 if (header.masked) {
31 MaskWebSocketFramePayload(header.masking_key, 0, payload);
32 }
33 frame->payload = payload;
34
35 return frame;
36 }
37
38 } // namespace
39
40 WebSocketChunkAssembler::WebSocketChunkAssembler() = default;
41
42 WebSocketChunkAssembler::~WebSocketChunkAssembler() = default;
43
Reset()44 void WebSocketChunkAssembler::Reset() {
45 current_frame_header_.reset();
46 chunk_buffer_.clear();
47 state_ = AssemblyState::kInitialFrame;
48 }
49
50 base::expected<std::unique_ptr<WebSocketFrame>, net::Error>
HandleChunk(std::unique_ptr<WebSocketFrameChunk> chunk)51 WebSocketChunkAssembler::HandleChunk(
52 std::unique_ptr<WebSocketFrameChunk> chunk) {
53 if (state_ == AssemblyState::kMessageFinished) {
54 Reset();
55 }
56
57 if (chunk->header) {
58 CHECK_EQ(state_, AssemblyState::kInitialFrame);
59 CHECK(!current_frame_header_);
60 current_frame_header_ = std::move(chunk->header);
61 }
62
63 CHECK(current_frame_header_);
64
65 const WebSocketFrameHeader::OpCode opcode = current_frame_header_->opcode;
66 const bool is_control_frame =
67 WebSocketFrameHeader::IsKnownControlOpCode(opcode) ||
68 WebSocketFrameHeader::IsReservedControlOpCode(opcode);
69 const bool is_data_frame = WebSocketFrameHeader::IsKnownDataOpCode(opcode) ||
70 WebSocketFrameHeader::IsReservedDataOpCode(opcode);
71
72 CHECK(is_control_frame || is_data_frame);
73
74 if (is_control_frame && !current_frame_header_->final) {
75 return base::unexpected(ERR_WS_PROTOCOL_ERROR);
76 }
77
78 if (is_control_frame &&
79 current_frame_header_->payload_length > kMaxControlFramePayload) {
80 return base::unexpected(ERR_WS_PROTOCOL_ERROR);
81 }
82
83 const bool is_first_chunk = state_ == AssemblyState::kInitialFrame;
84 const bool is_final_chunk = chunk->final_chunk;
85
86 const bool is_empty_middle_chunk =
87 !is_first_chunk && !is_final_chunk && chunk->payload.empty();
88 if (is_empty_middle_chunk) {
89 return base::unexpected(ERR_IO_PENDING);
90 }
91
92 // Handle single-chunk frame without buffering
93 const bool is_single_chunk_frame = is_first_chunk && is_final_chunk;
94 if (is_single_chunk_frame) {
95 CHECK_EQ(current_frame_header_->payload_length, chunk->payload.size());
96
97 auto frame = MakeWebSocketFrame(*current_frame_header_,
98 base::as_writable_bytes(chunk->payload));
99 state_ = AssemblyState::kMessageFinished;
100 return frame;
101 }
102
103 // For data frames, process each chunk separately without accumulating all
104 // in memory (streaming to render process)
105 if (is_data_frame) {
106 auto frame = MakeWebSocketFrame(*current_frame_header_,
107 base::as_writable_bytes(chunk->payload));
108
109 // Since we are synthesizing a frame that the origin server didn't send,
110 // we need to comply with the requirement ourselves.
111 if (state_ == AssemblyState::kContinuationFrame) {
112 // This is needed to satisfy the constraint of RFC7692:
113 //
114 // An endpoint MUST NOT set the "Per-Message Compressed" bit of control
115 // frames and non-first fragments of a data message.
116 frame->header.opcode = WebSocketFrameHeader::kOpCodeContinuation;
117 frame->header.reserved1 = false;
118 frame->header.reserved2 = false;
119 frame->header.reserved3 = false;
120 }
121 frame->header.payload_length = chunk->payload.size();
122 frame->header.final = current_frame_header_->final && chunk->final_chunk;
123
124 if (is_final_chunk) {
125 state_ = AssemblyState::kMessageFinished;
126 } else {
127 state_ = AssemblyState::kContinuationFrame;
128 }
129
130 return frame;
131 }
132
133 CHECK(is_control_frame && current_frame_header_->final);
134
135 // Control frames should be processed as a unit as they are small in size.
136 base::Extend(chunk_buffer_, chunk->payload);
137
138 if (!chunk->final_chunk) {
139 state_ = AssemblyState::kControlFrame;
140 return base::unexpected(ERR_IO_PENDING);
141 }
142 state_ = AssemblyState::kMessageFinished;
143
144 CHECK_EQ(current_frame_header_->payload_length, chunk_buffer_.size());
145
146 auto frame = MakeWebSocketFrame(*current_frame_header_,
147 base::as_writable_byte_span(chunk_buffer_));
148
149 state_ = AssemblyState::kMessageFinished;
150 return frame;
151 }
152
153 } // namespace net
154