1// Copyright 2013 The Chromium Authors. All rights reserved.
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_basic_stream.h"
6
7#include <algorithm>
8#include <limits>
9#include <string>
10#include <vector>
11
12#include "base/basictypes.h"
13#include "base/bind.h"
14#include "base/logging.h"
15#include "base/numerics/safe_conversions.h"
16#include "net/base/io_buffer.h"
17#include "net/base/net_errors.h"
18#include "net/socket/client_socket_handle.h"
19#include "net/websockets/websocket_errors.h"
20#include "net/websockets/websocket_frame.h"
21#include "net/websockets/websocket_frame_parser.h"
22
23namespace net {
24
25namespace {
26
27// This uses type uint64 to match the definition of
28// WebSocketFrameHeader::payload_length in websocket_frame.h.
29const uint64 kMaxControlFramePayload = 125;
30
31// The number of bytes to attempt to read at a time.
32// TODO(ricea): See if there is a better number or algorithm to fulfill our
33// requirements:
34//  1. We would like to use minimal memory on low-bandwidth or idle connections
35//  2. We would like to read as close to line speed as possible on
36//     high-bandwidth connections
37//  3. We can't afford to cause jank on the IO thread by copying large buffers
38//     around
39//  4. We would like to hit any sweet-spots that might exist in terms of network
40//     packet sizes / encryption block sizes / IPC alignment issues, etc.
41const int kReadBufferSize = 32 * 1024;
42
43typedef ScopedVector<WebSocketFrame>::const_iterator WebSocketFrameIterator;
44
45// Returns the total serialized size of |frames|. This function assumes that
46// |frames| will be serialized with mask field. This function forces the
47// masked bit of the frames on.
48int CalculateSerializedSizeAndTurnOnMaskBit(
49    ScopedVector<WebSocketFrame>* frames) {
50  const int kMaximumTotalSize = std::numeric_limits<int>::max();
51
52  int total_size = 0;
53  for (WebSocketFrameIterator it = frames->begin(); it != frames->end(); ++it) {
54    WebSocketFrame* frame = *it;
55    // Force the masked bit on.
56    frame->header.masked = true;
57    // We enforce flow control so the renderer should never be able to force us
58    // to cache anywhere near 2GB of frames.
59    int frame_size = frame->header.payload_length +
60                     GetWebSocketFrameHeaderSize(frame->header);
61    CHECK_GE(kMaximumTotalSize - total_size, frame_size)
62        << "Aborting to prevent overflow";
63    total_size += frame_size;
64  }
65  return total_size;
66}
67
68}  // namespace
69
70WebSocketBasicStream::WebSocketBasicStream(
71    scoped_ptr<ClientSocketHandle> connection,
72    const scoped_refptr<GrowableIOBuffer>& http_read_buffer,
73    const std::string& sub_protocol,
74    const std::string& extensions)
75    : read_buffer_(new IOBufferWithSize(kReadBufferSize)),
76      connection_(connection.Pass()),
77      http_read_buffer_(http_read_buffer),
78      sub_protocol_(sub_protocol),
79      extensions_(extensions),
80      generate_websocket_masking_key_(&GenerateWebSocketMaskingKey) {
81  // http_read_buffer_ should not be set if it contains no data.
82  if (http_read_buffer_.get() && http_read_buffer_->offset() == 0)
83    http_read_buffer_ = NULL;
84  DCHECK(connection_->is_initialized());
85}
86
87WebSocketBasicStream::~WebSocketBasicStream() { Close(); }
88
89int WebSocketBasicStream::ReadFrames(ScopedVector<WebSocketFrame>* frames,
90                                     const CompletionCallback& callback) {
91  DCHECK(frames->empty());
92  // If there is data left over after parsing the HTTP headers, attempt to parse
93  // it as WebSocket frames.
94  if (http_read_buffer_.get()) {
95    DCHECK_GE(http_read_buffer_->offset(), 0);
96    // We cannot simply copy the data into read_buffer_, as it might be too
97    // large.
98    scoped_refptr<GrowableIOBuffer> buffered_data;
99    buffered_data.swap(http_read_buffer_);
100    DCHECK(http_read_buffer_.get() == NULL);
101    ScopedVector<WebSocketFrameChunk> frame_chunks;
102    if (!parser_.Decode(buffered_data->StartOfBuffer(),
103                        buffered_data->offset(),
104                        &frame_chunks))
105      return WebSocketErrorToNetError(parser_.websocket_error());
106    if (!frame_chunks.empty()) {
107      int result = ConvertChunksToFrames(&frame_chunks, frames);
108      if (result != ERR_IO_PENDING)
109        return result;
110    }
111  }
112
113  // Run until socket stops giving us data or we get some frames.
114  while (true) {
115    // base::Unretained(this) here is safe because net::Socket guarantees not to
116    // call any callbacks after Disconnect(), which we call from the
117    // destructor. The caller of ReadFrames() is required to keep |frames|
118    // valid.
119    int result = connection_->socket()->Read(
120        read_buffer_.get(),
121        read_buffer_->size(),
122        base::Bind(&WebSocketBasicStream::OnReadComplete,
123                   base::Unretained(this),
124                   base::Unretained(frames),
125                   callback));
126    if (result == ERR_IO_PENDING)
127      return result;
128    result = HandleReadResult(result, frames);
129    if (result != ERR_IO_PENDING)
130      return result;
131    DCHECK(frames->empty());
132  }
133}
134
135int WebSocketBasicStream::WriteFrames(ScopedVector<WebSocketFrame>* frames,
136                                      const CompletionCallback& callback) {
137  // This function always concatenates all frames into a single buffer.
138  // TODO(ricea): Investigate whether it would be better in some cases to
139  // perform multiple writes with smaller buffers.
140  //
141  // First calculate the size of the buffer we need to allocate.
142  int total_size = CalculateSerializedSizeAndTurnOnMaskBit(frames);
143  scoped_refptr<IOBufferWithSize> combined_buffer(
144      new IOBufferWithSize(total_size));
145
146  char* dest = combined_buffer->data();
147  int remaining_size = total_size;
148  for (WebSocketFrameIterator it = frames->begin(); it != frames->end(); ++it) {
149    WebSocketFrame* frame = *it;
150    WebSocketMaskingKey mask = generate_websocket_masking_key_();
151    int result =
152        WriteWebSocketFrameHeader(frame->header, &mask, dest, remaining_size);
153    DCHECK_NE(ERR_INVALID_ARGUMENT, result)
154        << "WriteWebSocketFrameHeader() says that " << remaining_size
155        << " is not enough to write the header in. This should not happen.";
156    CHECK_GE(result, 0) << "Potentially security-critical check failed";
157    dest += result;
158    remaining_size -= result;
159
160    const int frame_size = frame->header.payload_length;
161    if (frame_size > 0) {
162      CHECK_GE(remaining_size, frame_size);
163      const char* const frame_data = frame->data->data();
164      std::copy(frame_data, frame_data + frame_size, dest);
165      MaskWebSocketFramePayload(mask, 0, dest, frame_size);
166      dest += frame_size;
167      remaining_size -= frame_size;
168    }
169  }
170  DCHECK_EQ(0, remaining_size) << "Buffer size calculation was wrong; "
171                               << remaining_size << " bytes left over.";
172  scoped_refptr<DrainableIOBuffer> drainable_buffer(
173      new DrainableIOBuffer(combined_buffer.get(), total_size));
174  return WriteEverything(drainable_buffer, callback);
175}
176
177void WebSocketBasicStream::Close() { connection_->socket()->Disconnect(); }
178
179std::string WebSocketBasicStream::GetSubProtocol() const {
180  return sub_protocol_;
181}
182
183std::string WebSocketBasicStream::GetExtensions() const { return extensions_; }
184
185/*static*/
186scoped_ptr<WebSocketBasicStream>
187WebSocketBasicStream::CreateWebSocketBasicStreamForTesting(
188    scoped_ptr<ClientSocketHandle> connection,
189    const scoped_refptr<GrowableIOBuffer>& http_read_buffer,
190    const std::string& sub_protocol,
191    const std::string& extensions,
192    WebSocketMaskingKeyGeneratorFunction key_generator_function) {
193  scoped_ptr<WebSocketBasicStream> stream(new WebSocketBasicStream(
194      connection.Pass(), http_read_buffer, sub_protocol, extensions));
195  stream->generate_websocket_masking_key_ = key_generator_function;
196  return stream.Pass();
197}
198
199int WebSocketBasicStream::WriteEverything(
200    const scoped_refptr<DrainableIOBuffer>& buffer,
201    const CompletionCallback& callback) {
202  while (buffer->BytesRemaining() > 0) {
203    // The use of base::Unretained() here is safe because on destruction we
204    // disconnect the socket, preventing any further callbacks.
205    int result = connection_->socket()->Write(
206        buffer.get(),
207        buffer->BytesRemaining(),
208        base::Bind(&WebSocketBasicStream::OnWriteComplete,
209                   base::Unretained(this),
210                   buffer,
211                   callback));
212    if (result > 0) {
213      buffer->DidConsume(result);
214    } else {
215      return result;
216    }
217  }
218  return OK;
219}
220
221void WebSocketBasicStream::OnWriteComplete(
222    const scoped_refptr<DrainableIOBuffer>& buffer,
223    const CompletionCallback& callback,
224    int result) {
225  if (result < 0) {
226    DCHECK_NE(ERR_IO_PENDING, result);
227    callback.Run(result);
228    return;
229  }
230
231  DCHECK_NE(0, result);
232  buffer->DidConsume(result);
233  result = WriteEverything(buffer, callback);
234  if (result != ERR_IO_PENDING)
235    callback.Run(result);
236}
237
238int WebSocketBasicStream::HandleReadResult(
239    int result,
240    ScopedVector<WebSocketFrame>* frames) {
241  DCHECK_NE(ERR_IO_PENDING, result);
242  DCHECK(frames->empty());
243  if (result < 0)
244    return result;
245  if (result == 0)
246    return ERR_CONNECTION_CLOSED;
247  ScopedVector<WebSocketFrameChunk> frame_chunks;
248  if (!parser_.Decode(read_buffer_->data(), result, &frame_chunks))
249    return WebSocketErrorToNetError(parser_.websocket_error());
250  if (frame_chunks.empty())
251    return ERR_IO_PENDING;
252  return ConvertChunksToFrames(&frame_chunks, frames);
253}
254
255int WebSocketBasicStream::ConvertChunksToFrames(
256    ScopedVector<WebSocketFrameChunk>* frame_chunks,
257    ScopedVector<WebSocketFrame>* frames) {
258  for (size_t i = 0; i < frame_chunks->size(); ++i) {
259    scoped_ptr<WebSocketFrame> frame;
260    int result = ConvertChunkToFrame(
261        scoped_ptr<WebSocketFrameChunk>((*frame_chunks)[i]), &frame);
262    (*frame_chunks)[i] = NULL;
263    if (result != OK)
264      return result;
265    if (frame)
266      frames->push_back(frame.release());
267  }
268  // All the elements of |frame_chunks| are now NULL, so there is no point in
269  // calling delete on them all.
270  frame_chunks->weak_clear();
271  if (frames->empty())
272    return ERR_IO_PENDING;
273  return OK;
274}
275
276int WebSocketBasicStream::ConvertChunkToFrame(
277    scoped_ptr<WebSocketFrameChunk> chunk,
278    scoped_ptr<WebSocketFrame>* frame) {
279  DCHECK(frame->get() == NULL);
280  bool is_first_chunk = false;
281  if (chunk->header) {
282    DCHECK(current_frame_header_ == NULL)
283        << "Received the header for a new frame without notification that "
284        << "the previous frame was complete (bug in WebSocketFrameParser?)";
285    is_first_chunk = true;
286    current_frame_header_.swap(chunk->header);
287  }
288  const int chunk_size = chunk->data.get() ? chunk->data->size() : 0;
289  DCHECK(current_frame_header_) << "Unexpected header-less chunk received "
290                                << "(final_chunk = " << chunk->final_chunk
291                                << ", data size = " << chunk_size
292                                << ") (bug in WebSocketFrameParser?)";
293  scoped_refptr<IOBufferWithSize> data_buffer;
294  data_buffer.swap(chunk->data);
295  const bool is_final_chunk = chunk->final_chunk;
296  const WebSocketFrameHeader::OpCode opcode = current_frame_header_->opcode;
297  if (WebSocketFrameHeader::IsKnownControlOpCode(opcode)) {
298    bool protocol_error = false;
299    if (!current_frame_header_->final) {
300      DVLOG(1) << "WebSocket protocol error. Control frame, opcode=" << opcode
301               << " received with FIN bit unset.";
302      protocol_error = true;
303    }
304    if (current_frame_header_->payload_length > kMaxControlFramePayload) {
305      DVLOG(1) << "WebSocket protocol error. Control frame, opcode=" << opcode
306               << ", payload_length=" << current_frame_header_->payload_length
307               << " exceeds maximum payload length for a control message.";
308      protocol_error = true;
309    }
310    if (protocol_error) {
311      current_frame_header_.reset();
312      return ERR_WS_PROTOCOL_ERROR;
313    }
314    if (!is_final_chunk) {
315      DVLOG(2) << "Encountered a split control frame, opcode " << opcode;
316      if (incomplete_control_frame_body_.get()) {
317        DVLOG(3) << "Appending to an existing split control frame.";
318        AddToIncompleteControlFrameBody(data_buffer);
319      } else {
320        DVLOG(3) << "Creating new storage for an incomplete control frame.";
321        incomplete_control_frame_body_ = new GrowableIOBuffer();
322        // This method checks for oversize control frames above, so as long as
323        // the frame parser is working correctly, this won't overflow. If a bug
324        // does cause it to overflow, it will CHECK() in
325        // AddToIncompleteControlFrameBody() without writing outside the buffer.
326        incomplete_control_frame_body_->SetCapacity(kMaxControlFramePayload);
327        AddToIncompleteControlFrameBody(data_buffer);
328      }
329      return OK;
330    }
331    if (incomplete_control_frame_body_.get()) {
332      DVLOG(2) << "Rejoining a split control frame, opcode " << opcode;
333      AddToIncompleteControlFrameBody(data_buffer);
334      const int body_size = incomplete_control_frame_body_->offset();
335      DCHECK_EQ(body_size,
336                static_cast<int>(current_frame_header_->payload_length));
337      scoped_refptr<IOBufferWithSize> body = new IOBufferWithSize(body_size);
338      memcpy(body->data(),
339             incomplete_control_frame_body_->StartOfBuffer(),
340             body_size);
341      incomplete_control_frame_body_ = NULL;  // Frame now complete.
342      DCHECK(is_final_chunk);
343      *frame = CreateFrame(is_final_chunk, body);
344      return OK;
345    }
346  }
347
348  // Apply basic sanity checks to the |payload_length| field from the frame
349  // header. A check for exact equality can only be used when the whole frame
350  // arrives in one chunk.
351  DCHECK_GE(current_frame_header_->payload_length,
352            base::checked_cast<uint64>(chunk_size));
353  DCHECK(!is_first_chunk || !is_final_chunk ||
354         current_frame_header_->payload_length ==
355             base::checked_cast<uint64>(chunk_size));
356
357  // Convert the chunk to a complete frame.
358  *frame = CreateFrame(is_final_chunk, data_buffer);
359  return OK;
360}
361
362scoped_ptr<WebSocketFrame> WebSocketBasicStream::CreateFrame(
363    bool is_final_chunk,
364    const scoped_refptr<IOBufferWithSize>& data) {
365  scoped_ptr<WebSocketFrame> result_frame;
366  const bool is_final_chunk_in_message =
367      is_final_chunk && current_frame_header_->final;
368  const int data_size = data.get() ? data->size() : 0;
369  const WebSocketFrameHeader::OpCode opcode = current_frame_header_->opcode;
370  // Empty frames convey no useful information unless they are the first frame
371  // (containing the type and flags) or have the "final" bit set.
372  if (is_final_chunk_in_message || data_size > 0 ||
373      current_frame_header_->opcode !=
374          WebSocketFrameHeader::kOpCodeContinuation) {
375    result_frame.reset(new WebSocketFrame(opcode));
376    result_frame->header.CopyFrom(*current_frame_header_);
377    result_frame->header.final = is_final_chunk_in_message;
378    result_frame->header.payload_length = data_size;
379    result_frame->data = data;
380    // Ensure that opcodes Text and Binary are only used for the first frame in
381    // the message. Also clear the reserved bits.
382    // TODO(ricea): If a future extension requires the reserved bits to be
383    // retained on continuation frames, make this behaviour conditional on a
384    // flag set at construction time.
385    if (!is_final_chunk && WebSocketFrameHeader::IsKnownDataOpCode(opcode)) {
386      current_frame_header_->opcode = WebSocketFrameHeader::kOpCodeContinuation;
387      current_frame_header_->reserved1 = false;
388      current_frame_header_->reserved2 = false;
389      current_frame_header_->reserved3 = false;
390    }
391  }
392  // Make sure that a frame header is not applied to any chunks that do not
393  // belong to it.
394  if (is_final_chunk)
395    current_frame_header_.reset();
396  return result_frame.Pass();
397}
398
399void WebSocketBasicStream::AddToIncompleteControlFrameBody(
400    const scoped_refptr<IOBufferWithSize>& data_buffer) {
401  if (!data_buffer.get())
402    return;
403  const int new_offset =
404      incomplete_control_frame_body_->offset() + data_buffer->size();
405  CHECK_GE(incomplete_control_frame_body_->capacity(), new_offset)
406      << "Control frame body larger than frame header indicates; frame parser "
407         "bug?";
408  memcpy(incomplete_control_frame_body_->data(),
409         data_buffer->data(),
410         data_buffer->size());
411  incomplete_control_frame_body_->set_offset(new_offset);
412}
413
414void WebSocketBasicStream::OnReadComplete(ScopedVector<WebSocketFrame>* frames,
415                                          const CompletionCallback& callback,
416                                          int result) {
417  result = HandleReadResult(result, frames);
418  if (result == ERR_IO_PENDING)
419    result = ReadFrames(frames, callback);
420  if (result != ERR_IO_PENDING)
421    callback.Run(result);
422}
423
424}  // namespace net
425