1 // Copyright 2013 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_deflate_stream.h"
6
7 #include <stdint.h>
8
9 #include <algorithm>
10 #include <string>
11 #include <utility>
12 #include <vector>
13
14 #include "base/functional/bind.h"
15 #include "base/logging.h"
16 #include "base/memory/scoped_refptr.h"
17 #include "net/base/io_buffer.h"
18 #include "net/base/net_errors.h"
19 #include "net/websockets/websocket_deflate_parameters.h"
20 #include "net/websockets/websocket_deflate_predictor.h"
21 #include "net/websockets/websocket_deflater.h"
22 #include "net/websockets/websocket_errors.h"
23 #include "net/websockets/websocket_frame.h"
24 #include "net/websockets/websocket_inflater.h"
25 #include "net/websockets/websocket_stream.h"
26
27 class GURL;
28
29 namespace net {
30
31 namespace {
32
33 const int kWindowBits = 15;
34 const size_t kChunkSize = 4 * 1024;
35
36 } // namespace
37
WebSocketDeflateStream(std::unique_ptr<WebSocketStream> stream,const WebSocketDeflateParameters & params,std::unique_ptr<WebSocketDeflatePredictor> predictor)38 WebSocketDeflateStream::WebSocketDeflateStream(
39 std::unique_ptr<WebSocketStream> stream,
40 const WebSocketDeflateParameters& params,
41 std::unique_ptr<WebSocketDeflatePredictor> predictor)
42 : stream_(std::move(stream)),
43 deflater_(params.client_context_take_over_mode()),
44 inflater_(kChunkSize, kChunkSize),
45 predictor_(std::move(predictor)) {
46 DCHECK(stream_);
47 DCHECK(params.IsValidAsResponse());
48 int client_max_window_bits = 15;
49 if (params.is_client_max_window_bits_specified()) {
50 DCHECK(params.has_client_max_window_bits_value());
51 client_max_window_bits = params.client_max_window_bits();
52 }
53 deflater_.Initialize(client_max_window_bits);
54 inflater_.Initialize(kWindowBits);
55 }
56
57 WebSocketDeflateStream::~WebSocketDeflateStream() = default;
58
ReadFrames(std::vector<std::unique_ptr<WebSocketFrame>> * frames,CompletionOnceCallback callback)59 int WebSocketDeflateStream::ReadFrames(
60 std::vector<std::unique_ptr<WebSocketFrame>>* frames,
61 CompletionOnceCallback callback) {
62 read_callback_ = std::move(callback);
63 inflater_outputs_.clear();
64 int result = stream_->ReadFrames(
65 frames, base::BindOnce(&WebSocketDeflateStream::OnReadComplete,
66 base::Unretained(this), base::Unretained(frames)));
67 if (result < 0)
68 return result;
69 DCHECK_EQ(OK, result);
70 DCHECK(!frames->empty());
71
72 return InflateAndReadIfNecessary(frames);
73 }
74
WriteFrames(std::vector<std::unique_ptr<WebSocketFrame>> * frames,CompletionOnceCallback callback)75 int WebSocketDeflateStream::WriteFrames(
76 std::vector<std::unique_ptr<WebSocketFrame>>* frames,
77 CompletionOnceCallback callback) {
78 deflater_outputs_.clear();
79 int result = Deflate(frames);
80 if (result != OK)
81 return result;
82 if (frames->empty())
83 return OK;
84 return stream_->WriteFrames(frames, std::move(callback));
85 }
86
Close()87 void WebSocketDeflateStream::Close() { stream_->Close(); }
88
GetSubProtocol() const89 std::string WebSocketDeflateStream::GetSubProtocol() const {
90 return stream_->GetSubProtocol();
91 }
92
GetExtensions() const93 std::string WebSocketDeflateStream::GetExtensions() const {
94 return stream_->GetExtensions();
95 }
96
GetNetLogWithSource() const97 const NetLogWithSource& WebSocketDeflateStream::GetNetLogWithSource() const {
98 return stream_->GetNetLogWithSource();
99 }
100
OnReadComplete(std::vector<std::unique_ptr<WebSocketFrame>> * frames,int result)101 void WebSocketDeflateStream::OnReadComplete(
102 std::vector<std::unique_ptr<WebSocketFrame>>* frames,
103 int result) {
104 if (result != OK) {
105 frames->clear();
106 std::move(read_callback_).Run(result);
107 return;
108 }
109
110 int r = InflateAndReadIfNecessary(frames);
111 if (r != ERR_IO_PENDING)
112 std::move(read_callback_).Run(r);
113 }
114
Deflate(std::vector<std::unique_ptr<WebSocketFrame>> * frames)115 int WebSocketDeflateStream::Deflate(
116 std::vector<std::unique_ptr<WebSocketFrame>>* frames) {
117 std::vector<std::unique_ptr<WebSocketFrame>> frames_to_write;
118 // Store frames of the currently processed message if writing_state_ equals to
119 // WRITING_POSSIBLY_COMPRESSED_MESSAGE.
120 std::vector<std::unique_ptr<WebSocketFrame>> frames_of_message;
121 for (size_t i = 0; i < frames->size(); ++i) {
122 DCHECK(!(*frames)[i]->header.reserved1);
123 if (!WebSocketFrameHeader::IsKnownDataOpCode((*frames)[i]->header.opcode)) {
124 frames_to_write.push_back(std::move((*frames)[i]));
125 continue;
126 }
127 if (writing_state_ == NOT_WRITING)
128 OnMessageStart(*frames, i);
129
130 std::unique_ptr<WebSocketFrame> frame(std::move((*frames)[i]));
131 predictor_->RecordInputDataFrame(frame.get());
132
133 if (writing_state_ == WRITING_UNCOMPRESSED_MESSAGE) {
134 if (frame->header.final)
135 writing_state_ = NOT_WRITING;
136 predictor_->RecordWrittenDataFrame(frame.get());
137 frames_to_write.push_back(std::move(frame));
138 current_writing_opcode_ = WebSocketFrameHeader::kOpCodeContinuation;
139 } else {
140 if (frame->payload &&
141 !deflater_.AddBytes(
142 frame->payload,
143 static_cast<size_t>(frame->header.payload_length))) {
144 DVLOG(1) << "WebSocket protocol error. "
145 << "deflater_.AddBytes() returns an error.";
146 return ERR_WS_PROTOCOL_ERROR;
147 }
148 if (frame->header.final && !deflater_.Finish()) {
149 DVLOG(1) << "WebSocket protocol error. "
150 << "deflater_.Finish() returns an error.";
151 return ERR_WS_PROTOCOL_ERROR;
152 }
153
154 if (writing_state_ == WRITING_COMPRESSED_MESSAGE) {
155 if (deflater_.CurrentOutputSize() >= kChunkSize ||
156 frame->header.final) {
157 int result = AppendCompressedFrame(frame->header, &frames_to_write);
158 if (result != OK)
159 return result;
160 }
161 if (frame->header.final)
162 writing_state_ = NOT_WRITING;
163 } else {
164 DCHECK_EQ(WRITING_POSSIBLY_COMPRESSED_MESSAGE, writing_state_);
165 bool final = frame->header.final;
166 frames_of_message.push_back(std::move(frame));
167 if (final) {
168 int result = AppendPossiblyCompressedMessage(&frames_of_message,
169 &frames_to_write);
170 if (result != OK)
171 return result;
172 frames_of_message.clear();
173 writing_state_ = NOT_WRITING;
174 }
175 }
176 }
177 }
178 DCHECK_NE(WRITING_POSSIBLY_COMPRESSED_MESSAGE, writing_state_);
179 frames->swap(frames_to_write);
180 return OK;
181 }
182
OnMessageStart(const std::vector<std::unique_ptr<WebSocketFrame>> & frames,size_t index)183 void WebSocketDeflateStream::OnMessageStart(
184 const std::vector<std::unique_ptr<WebSocketFrame>>& frames,
185 size_t index) {
186 WebSocketFrame* frame = frames[index].get();
187 current_writing_opcode_ = frame->header.opcode;
188 DCHECK(current_writing_opcode_ == WebSocketFrameHeader::kOpCodeText ||
189 current_writing_opcode_ == WebSocketFrameHeader::kOpCodeBinary);
190 WebSocketDeflatePredictor::Result prediction =
191 predictor_->Predict(frames, index);
192
193 switch (prediction) {
194 case WebSocketDeflatePredictor::DEFLATE:
195 writing_state_ = WRITING_COMPRESSED_MESSAGE;
196 return;
197 case WebSocketDeflatePredictor::DO_NOT_DEFLATE:
198 writing_state_ = WRITING_UNCOMPRESSED_MESSAGE;
199 return;
200 case WebSocketDeflatePredictor::TRY_DEFLATE:
201 writing_state_ = WRITING_POSSIBLY_COMPRESSED_MESSAGE;
202 return;
203 }
204 NOTREACHED();
205 }
206
AppendCompressedFrame(const WebSocketFrameHeader & header,std::vector<std::unique_ptr<WebSocketFrame>> * frames_to_write)207 int WebSocketDeflateStream::AppendCompressedFrame(
208 const WebSocketFrameHeader& header,
209 std::vector<std::unique_ptr<WebSocketFrame>>* frames_to_write) {
210 const WebSocketFrameHeader::OpCode opcode = current_writing_opcode_;
211 scoped_refptr<IOBufferWithSize> compressed_payload =
212 deflater_.GetOutput(deflater_.CurrentOutputSize());
213 if (!compressed_payload.get()) {
214 DVLOG(1) << "WebSocket protocol error. "
215 << "deflater_.GetOutput() returns an error.";
216 return ERR_WS_PROTOCOL_ERROR;
217 }
218 deflater_outputs_.push_back(compressed_payload);
219 auto compressed = std::make_unique<WebSocketFrame>(opcode);
220 compressed->header.CopyFrom(header);
221 compressed->header.opcode = opcode;
222 compressed->header.final = header.final;
223 compressed->header.reserved1 =
224 (opcode != WebSocketFrameHeader::kOpCodeContinuation);
225 compressed->payload = compressed_payload->data();
226 compressed->header.payload_length = compressed_payload->size();
227
228 current_writing_opcode_ = WebSocketFrameHeader::kOpCodeContinuation;
229 predictor_->RecordWrittenDataFrame(compressed.get());
230 frames_to_write->push_back(std::move(compressed));
231 return OK;
232 }
233
AppendPossiblyCompressedMessage(std::vector<std::unique_ptr<WebSocketFrame>> * frames,std::vector<std::unique_ptr<WebSocketFrame>> * frames_to_write)234 int WebSocketDeflateStream::AppendPossiblyCompressedMessage(
235 std::vector<std::unique_ptr<WebSocketFrame>>* frames,
236 std::vector<std::unique_ptr<WebSocketFrame>>* frames_to_write) {
237 DCHECK(!frames->empty());
238
239 const WebSocketFrameHeader::OpCode opcode = current_writing_opcode_;
240 scoped_refptr<IOBufferWithSize> compressed_payload =
241 deflater_.GetOutput(deflater_.CurrentOutputSize());
242 if (!compressed_payload.get()) {
243 DVLOG(1) << "WebSocket protocol error. "
244 << "deflater_.GetOutput() returns an error.";
245 return ERR_WS_PROTOCOL_ERROR;
246 }
247 deflater_outputs_.push_back(compressed_payload);
248
249 uint64_t original_payload_length = 0;
250 for (size_t i = 0; i < frames->size(); ++i) {
251 WebSocketFrame* frame = (*frames)[i].get();
252 // Asserts checking that frames represent one whole data message.
253 DCHECK(WebSocketFrameHeader::IsKnownDataOpCode(frame->header.opcode));
254 DCHECK_EQ(i == 0,
255 WebSocketFrameHeader::kOpCodeContinuation !=
256 frame->header.opcode);
257 DCHECK_EQ(i == frames->size() - 1, frame->header.final);
258 original_payload_length += frame->header.payload_length;
259 }
260 if (original_payload_length <=
261 static_cast<uint64_t>(compressed_payload->size())) {
262 // Compression is not effective. Use the original frames.
263 for (auto& frame : *frames) {
264 predictor_->RecordWrittenDataFrame(frame.get());
265 frames_to_write->push_back(std::move(frame));
266 }
267 frames->clear();
268 return OK;
269 }
270 auto compressed = std::make_unique<WebSocketFrame>(opcode);
271 compressed->header.CopyFrom((*frames)[0]->header);
272 compressed->header.opcode = opcode;
273 compressed->header.final = true;
274 compressed->header.reserved1 = true;
275 compressed->payload = compressed_payload->data();
276 compressed->header.payload_length = compressed_payload->size();
277
278 predictor_->RecordWrittenDataFrame(compressed.get());
279 frames_to_write->push_back(std::move(compressed));
280 return OK;
281 }
282
Inflate(std::vector<std::unique_ptr<WebSocketFrame>> * frames)283 int WebSocketDeflateStream::Inflate(
284 std::vector<std::unique_ptr<WebSocketFrame>>* frames) {
285 std::vector<std::unique_ptr<WebSocketFrame>> frames_to_output;
286 std::vector<std::unique_ptr<WebSocketFrame>> frames_passed;
287 frames->swap(frames_passed);
288 for (auto& frame_passed : frames_passed) {
289 std::unique_ptr<WebSocketFrame> frame(std::move(frame_passed));
290 frame_passed = nullptr;
291 DVLOG(3) << "Input frame: opcode=" << frame->header.opcode
292 << " final=" << frame->header.final
293 << " reserved1=" << frame->header.reserved1
294 << " payload_length=" << frame->header.payload_length;
295
296 if (!WebSocketFrameHeader::IsKnownDataOpCode(frame->header.opcode)) {
297 frames_to_output.push_back(std::move(frame));
298 continue;
299 }
300
301 if (reading_state_ == NOT_READING) {
302 if (frame->header.reserved1)
303 reading_state_ = READING_COMPRESSED_MESSAGE;
304 else
305 reading_state_ = READING_UNCOMPRESSED_MESSAGE;
306 current_reading_opcode_ = frame->header.opcode;
307 } else {
308 if (frame->header.reserved1) {
309 DVLOG(1) << "WebSocket protocol error. "
310 << "Receiving a non-first frame with RSV1 flag set.";
311 return ERR_WS_PROTOCOL_ERROR;
312 }
313 }
314
315 if (reading_state_ == READING_UNCOMPRESSED_MESSAGE) {
316 if (frame->header.final)
317 reading_state_ = NOT_READING;
318 current_reading_opcode_ = WebSocketFrameHeader::kOpCodeContinuation;
319 frames_to_output.push_back(std::move(frame));
320 } else {
321 DCHECK_EQ(reading_state_, READING_COMPRESSED_MESSAGE);
322 if (frame->payload &&
323 !inflater_.AddBytes(
324 frame->payload,
325 static_cast<size_t>(frame->header.payload_length))) {
326 DVLOG(1) << "WebSocket protocol error. "
327 << "inflater_.AddBytes() returns an error.";
328 return ERR_WS_PROTOCOL_ERROR;
329 }
330 if (frame->header.final) {
331 if (!inflater_.Finish()) {
332 DVLOG(1) << "WebSocket protocol error. "
333 << "inflater_.Finish() returns an error.";
334 return ERR_WS_PROTOCOL_ERROR;
335 }
336 }
337 // TODO(yhirano): Many frames can be generated by the inflater and
338 // memory consumption can grow.
339 // We could avoid it, but avoiding it makes this class much more
340 // complicated.
341 while (inflater_.CurrentOutputSize() >= kChunkSize ||
342 frame->header.final) {
343 size_t size = std::min(kChunkSize, inflater_.CurrentOutputSize());
344 auto inflated =
345 std::make_unique<WebSocketFrame>(WebSocketFrameHeader::kOpCodeText);
346 scoped_refptr<IOBufferWithSize> data = inflater_.GetOutput(size);
347 inflater_outputs_.push_back(data);
348 bool is_final = !inflater_.CurrentOutputSize() && frame->header.final;
349 if (!data.get()) {
350 DVLOG(1) << "WebSocket protocol error. "
351 << "inflater_.GetOutput() returns an error.";
352 return ERR_WS_PROTOCOL_ERROR;
353 }
354 inflated->header.CopyFrom(frame->header);
355 inflated->header.opcode = current_reading_opcode_;
356 inflated->header.final = is_final;
357 inflated->header.reserved1 = false;
358 inflated->payload = data->data();
359 inflated->header.payload_length = data->size();
360 DVLOG(3) << "Inflated frame: opcode=" << inflated->header.opcode
361 << " final=" << inflated->header.final
362 << " reserved1=" << inflated->header.reserved1
363 << " payload_length=" << inflated->header.payload_length;
364 frames_to_output.push_back(std::move(inflated));
365 current_reading_opcode_ = WebSocketFrameHeader::kOpCodeContinuation;
366 if (is_final)
367 break;
368 }
369 if (frame->header.final)
370 reading_state_ = NOT_READING;
371 }
372 }
373 frames->swap(frames_to_output);
374 return frames->empty() ? ERR_IO_PENDING : OK;
375 }
376
InflateAndReadIfNecessary(std::vector<std::unique_ptr<WebSocketFrame>> * frames)377 int WebSocketDeflateStream::InflateAndReadIfNecessary(
378 std::vector<std::unique_ptr<WebSocketFrame>>* frames) {
379 int result = Inflate(frames);
380 while (result == ERR_IO_PENDING) {
381 DCHECK(frames->empty());
382
383 result = stream_->ReadFrames(
384 frames,
385 base::BindOnce(&WebSocketDeflateStream::OnReadComplete,
386 base::Unretained(this), base::Unretained(frames)));
387 if (result < 0)
388 break;
389 DCHECK_EQ(OK, result);
390 DCHECK(!frames->empty());
391
392 result = Inflate(frames);
393 }
394 if (result < 0)
395 frames->clear();
396 return result;
397 }
398
399 } // namespace net
400