WebSocketBasicStream framing logic

net/websockets/websocket_basic_stream.cc

Index: net/websockets/websocket_basic_stream.cc
diff --git a/net/websockets/websocket_basic_stream.cc b/net/websockets/websocket_basic_stream.cc
new file mode 100644
index 0000000000000000000000000000000000000000..4b0e42c2c731cbb82037e2e712c8143132ee5ef9
--- /dev/null
+++ b/net/websockets/websocket_basic_stream.cc
@@ -0,0 +1,250 @@
+// Copyright 2013 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "net/websockets/websocket_basic_stream.h"
+
+#include <algorithm>
+#include <limits>
+#include <string>
+#include <vector>
+
+#include "base/basictypes.h"
+#include "base/bind.h"

[tyoshino (SeeGerritForStatus), 2013/08/22 21:08:43]

add
base/logging.h

[tyoshino (SeeGerritForStatus), 2013/08/22 21:08:43]

add
net/base/net_errors.h

[Adam Rice, 2013/08/23 06:05:18]

Done.

[Adam Rice, 2013/08/23 06:05:18]

Done.

+#include "net/base/io_buffer.h"
+#include "net/socket/client_socket_handle.h"
+#include "net/websockets/websocket_errors.h"
+#include "net/websockets/websocket_frame.h"
+#include "net/websockets/websocket_frame_parser.h"
+
+namespace net {
+
+namespace {
+
+// The number of bytes to attempt to read at a time.
+// TODO(ricea): See if there is a better number. Should it start small, and get
+// bigger if needed?
+const int kReadAtATime = 32 * 1024;
+
+}  // namespace
+
+WebSocketBasicStream::WebSocketBasicStream(
+    scoped_ptr<ClientSocketHandle> connection)
+    : read_buffer_(new IOBufferWithSize(kReadAtATime)),
+      connection_(connection.Pass()),
+      generate_websocket_masking_key_(&GenerateWebSocketMaskingKey) {}
+
+WebSocketBasicStream::~WebSocketBasicStream() {
+  connection_->socket()->Disconnect();
+}
+
+int WebSocketBasicStream::ReadFrames(
+    ScopedVector<WebSocketFrameChunk>* frame_chunks,
+    const CompletionCallback& callback) {
+  DCHECK(frame_chunks->empty());
+  // If there is data left over after parsing the HTTP headers, attempt to parse
+  // it as WebSocket frames.
+  if (http_read_buffer_) {
+    DCHECK_GE(http_read_buffer_->offset(), 0);
+    // offset() gives the end of the data, but it also controls where data()
+    // points to, so we need to store its value and then reset it to 0.
+    int end_of_data_offset = http_read_buffer_->offset();
+    http_read_buffer_->set_offset(0);
+    if (!parser_.Decode(
+            http_read_buffer_->data(), end_of_data_offset, frame_chunks)) {
+      http_read_buffer_ = NULL;
+      return WebSocketErrorToNetError(parser_.websocket_error());
+    }
+    http_read_buffer_ = NULL;
+  }
+  // Loop until we either have at least one chunk to return, or we get
+  // ERR_IO_PENDING, or something goes wrong.
+  while (frame_chunks->empty()) {
+    // This use of base::Unretained() is safe because WebSocketChannel will
+    // delete us before deleting frame_chunks.
+    int result =
+        connection_->socket()->Read(read_buffer_.get(),
+                                    read_buffer_->size(),
+                                    base::Bind(&WebSocketBasicStream::ReadDone,
+                                               base::Unretained(this),
+                                               base::Unretained(frame_chunks),
+                                               callback));
+    if (result > 0) {
+      if (!parser_.Decode(read_buffer_->data(), result, frame_chunks)) {
+        return WebSocketErrorToNetError(parser_.websocket_error());
+      }
+    } else if (result == 0 && frame_chunks->empty()) {
+      return ERR_CONNECTION_CLOSED;
+    } else {
+      return result;
+    }
+  }
+  return OK;
+}
+
+void WebSocketBasicStream::ReadDone(
+    ScopedVector<WebSocketFrameChunk>* frame_chunks,
+    const CompletionCallback& callback,
+    int result) {
+  if (result > 0) {
+    if (parser_.Decode(read_buffer_->data(), result, frame_chunks)) {
+      if (!frame_chunks->empty()) {
+        result = OK;
+      } else {
+        result = ReadFrames(frame_chunks, callback);
+        if (result == ERR_IO_PENDING) {
+          // This method will be called back again.
+          return;
+        }
+      }
+    } else {
+      result = WebSocketErrorToNetError(parser_.websocket_error());
+    }
+  }
+  if (result == 0 && frame_chunks->empty()) {
+    result = ERR_CONNECTION_CLOSED;
+  }
+  DCHECK_NE(ERR_IO_PENDING, result);
+  callback.Run(result);
+}
+
+int WebSocketBasicStream::WriteFrames(
+    ScopedVector<WebSocketFrameChunk>* frame_chunks,
+    const CompletionCallback& callback) {
+  // This function always concatenates all frames into a single buffer.
+  // TODO(ricea): Investigate whether it would be better in some cases to
+  // perform multiple writes with smaller buffers.
+  //
+  // First calculate the size of the buffer we need to allocate.
+  typedef ScopedVector<WebSocketFrameChunk>::const_iterator Iterator;
+  int total_size = 0;
+  for (Iterator it = frame_chunks->begin(); it != frame_chunks->end(); ++it) {
+    WebSocketFrameChunk* chunk = *it;
+    DCHECK(chunk->header && chunk->final_chunk)
+        << "Only complete frames are supported by WebSocketBasicStream";
+    // Force the masked bit on.
+    chunk->header->masked = true;
+    // We enforce flow control so the renderer should never be able to force us
+    // to cache anywhere near 2GB of frames.
+    int chunk_size =
+        chunk->data->size() + GetWebSocketFrameHeaderSize(*(chunk->header));
+    CHECK_GE(std::numeric_limits<int>::max() - total_size, chunk_size)
+        << "Aborting to prevent overflow";
+    total_size += chunk_size;
+  }
+  scoped_refptr<IOBufferWithSize> dest(new IOBufferWithSize(total_size));
+  char* data = dest->data();
+  int remaining_size = total_size;
+  for (Iterator it = frame_chunks->begin(); it != frame_chunks->end(); ++it) {
+    WebSocketFrameChunk* chunk = *it;
+    WebSocketMaskingKey mask = generate_websocket_masking_key_();
+    int result = WriteWebSocketFrameHeader(
+        *(chunk->header), &mask, data, remaining_size);
+    DCHECK(result != ERR_INVALID_ARGUMENT)
+        << "WriteWebSocketFrameHeader() says that " << remaining_size
+        << " is not enough to write the header in. This should not happen.";
+    CHECK_GE(result, 0) << "Potentially security-critical check failed";
+    data += result;
+    remaining_size -= result;
+
+    const char* const frame_data = chunk->data->data();
+    const int frame_size = chunk->data->size();
+    CHECK_GE(remaining_size, frame_size);
+    std::copy(frame_data, frame_data + frame_size, data);
+    MaskWebSocketFramePayload(mask, 0, data, frame_size);
+    data += frame_size;
+    remaining_size -= frame_size;
+  }
+  DCHECK_EQ(0, remaining_size) << "Buffer size calculation was wrong; "
+                               << remaining_size << " bytes left over.";
+  scoped_refptr<DrainableIOBuffer> drainable_buffer(
+      new DrainableIOBuffer(dest, total_size));
+  return WriteEverything(drainable_buffer, callback);
+}
+
+int WebSocketBasicStream::WriteEverything(
+    const scoped_refptr<DrainableIOBuffer>& buffer,
+    const CompletionCallback& callback) {
+  while (buffer->BytesRemaining() > 0) {
+    // The use of base::Unretained() here is safe because on destruction we
+    // disconnect the socket, preventing any further callbacks.
+    int result = connection_->socket()->Write(
+        buffer.get(),
+        buffer->BytesRemaining(),
+        base::Bind(&WebSocketBasicStream::WriteDone,
+                   base::Unretained(this),
+                   buffer,
+                   callback));
+    if (result > 0) {
+      buffer->DidConsume(result);
+    } else {
+      return result;
+    }
+  }
+  return OK;
+}
+
+void WebSocketBasicStream::WriteDone(
+    const scoped_refptr<DrainableIOBuffer>& buffer,
+    const CompletionCallback& callback,
+    int result) {
+  if (result > 0) {
+    buffer->DidConsume(result);
+    if (buffer->BytesRemaining() > 0) {
+      int result = WriteEverything(buffer, callback);
+      if (result != ERR_IO_PENDING) {
+        callback.Run(result);
+      }
+    } else {
+      callback.Run(OK);
+    }
+  } else {
+    DCHECK(result != ERR_IO_PENDING);
+    callback.Run(result);
+  }
+}
+
+void WebSocketBasicStream::Close() { connection_->socket()->Disconnect(); }
+
+std::string WebSocketBasicStream::GetSubProtocol() const {
+  return sub_protocol_;
+}
+
+std::string WebSocketBasicStream::GetExtensions() const { return extensions_; }
+
+int WebSocketBasicStream::SendHandshakeRequest(
+    const GURL& url,
+    const HttpRequestHeaders& headers,
+    HttpResponseInfo* response_info,
+    const CompletionCallback& callback) {
+  // TODO(ricea): Implement handshake-related functionality.
+  NOTREACHED();
+  return ERR_NOT_IMPLEMENTED;
+}
+
+int WebSocketBasicStream::ReadHandshakeResponse(
+    const CompletionCallback& callback) {
+  NOTREACHED();
+  return ERR_NOT_IMPLEMENTED;
+}
+
+/*static*/
+scoped_ptr<WebSocketBasicStream>
+WebSocketBasicStream::CreateWebSocketBasicStreamForTesting(
+    scoped_ptr<ClientSocketHandle> connection,
+    const scoped_refptr<GrowableIOBuffer>& http_read_buffer,
+    const std::string& sub_protocol,
+    const std::string& extensions,
+    WebSocketMaskingKeyGeneratorFunction key_generator_function) {
+  scoped_ptr<WebSocketBasicStream> stream(
+      new WebSocketBasicStream(connection.Pass()));
+  if (http_read_buffer) {
+    stream->http_read_buffer_ = http_read_buffer;
+  }
+  stream->sub_protocol_ = sub_protocol;
+  stream->extensions_ = extensions;
+  stream->generate_websocket_masking_key_ = key_generator_function;
+  return stream.Pass();
+}
+
+}  // namespace net