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..d37205d1508d002356257c5bec412c7a5596dfcd
--- /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"
+#include "base/logging.h"
+#include "net/base/io_buffer.h"
+#include "net/base/net_errors.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;

[szym, 2013/08/27 22:54:04]

suggest: kReadBufferSize

This is merely the size of the buffer passed to Socket::Read. Why would you
change this value dynamically?

[Adam Rice, 2013/08/28 12:01:11]

Some applications of WebSockets may require a large number of low-bandwidth
sockets, so it would be good to limit the memory used by each socket.

+
+}  // namespace
+
+WebSocketBasicStream::WebSocketBasicStream(
+    scoped_ptr<ClientSocketHandle> connection)
+    : read_buffer_(new IOBufferWithSize(kReadAtATime)),
+      connection_(connection.Pass()),
+      generate_websocket_masking_key_(&GenerateWebSocketMaskingKey) {}

[szym, 2013/08/27 22:54:04]

suggest: DCHECK(connection_->is_initialized())

[Adam Rice, 2013/08/28 12:01:11]

Done.

+
+WebSocketBasicStream::~WebSocketBasicStream() {
+  connection_->socket()->Disconnect();

[szym, 2013/08/27 22:54:04]

Suggest calling Close() instead.

+}
+
+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);
+    if (!parser_.Decode(http_read_buffer_->StartOfBuffer(),
+                        http_read_buffer_->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.

[szym, 2013/08/27 22:54:04]

Do not refer to WebSocketChannel. This use is safe, because of the semantics of
ReadFrames. "|frame_chunks| remains owned by the caller and must be valid until
the operation completes..." So this comment is unnecessary, but if you want to
be explicit, just say "The semantics of ReadFrames() guarantee |frame_chunks|
will live until callback or Close()."

[Adam Rice, 2013/08/28 12:01:11]

I plan to rename base::Unretained to base::CriticalSecurityVulnerability in the
future, so I am putting these comments in place as insurance.

+    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(

[szym, 2013/08/27 22:54:04]

Chromium code style: Function declaration order should match function definition
order.

[Adam Rice, 2013/08/28 12:01:11]

Done.

+    ScopedVector<WebSocketFrameChunk>* frame_chunks,
+    const CompletionCallback& callback,
+    int result) {
+  if (result > 0) {

[szym, 2013/08/27 22:54:04]

The logic here and in ReadFrames is very similar but rather difficult to follow.
I suggest extracting it to one location in a design pattern similar to "int
DoLoop(int result)". Here, I suggest:

// Returns ERR_IO_PENDING if needs more data.
int HandleReadResult(int result, * frame_chunks) {
  DCHECK_NE(ERR_IO_PENDING, result);
  DCHECK(frame_chunks->empty());
  if (result < 0)
    return result;
  if (result == 0)
    return ERR_CONNECTION_CLOSED;
  if (!parser_.Decode(read_buffer_, result, frame_chunks))
    return WebSocketErrorToNetError(parser_.websocket_error())
  if (!frame_chunks->empty())
    return OK;
  return ERR_IO_PENDING;
}

int ReadFrames(* frame_chunks, callback) {
  DCHECK(frame_chunks->empty());
  // Run until socket stops giving us data or we get some chunks.
  while(true) {
    int result = socket->Read(..., Bind(OnReadComplete, ...));
    if (result == ERR_IO_PENDING)
      return result;
    result = HandleReadResult(result, frame_chunks);
    if (result != ERR_IO_PENDING)
      return result;
  }
}

void OnReadComplete(int result, * frame_chunks, callback) {
  int result = HandleReadResult(result, frame_chunks);
  if (result == ERR_IO_PENDING)
    result = ReadFrames(frame_chunks, callback);
  if (result != ERR_IO_PENDING)
    callback.Run(result);
}

[Adam Rice, 2013/08/28 12:01:11]

Done.

+    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)

[szym, 2013/08/27 22:54:04]

Suggest breaking DCHECK(..&&..) into two DCHECKs

[Adam Rice, 2013/08/28 12:01:11]

Done.

+        << "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)

[szym, 2013/08/27 22:54:04]

Define a constant 
const int kMaximumTotalSize = ...
and use it instead of std::numeric_limits here

[Adam Rice, 2013/08/28 12:01:11]

Done.

+        << "Aborting to prevent overflow";
+    total_size += chunk_size;
+  }
+  scoped_refptr<IOBufferWithSize> combined_buffer(
+      new IOBufferWithSize(total_size));
+  char* dest = combined_buffer->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, dest, 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";
+    dest += 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, dest);
+    MaskWebSocketFramePayload(mask, 0, dest, frame_size);
+    dest += 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(combined_buffer, 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) {

[szym, 2013/08/27 22:54:04]

Suggest unrolling the if{} folds in this function:

if (result < 0) {
  callback.Run(result);
  return;
}

buffer->DidConsume(result);
...

[Adam Rice, 2013/08/28 12:01:11]

Done.

+    buffer->DidConsume(result);
+    if (buffer->BytesRemaining() > 0) {

[szym, 2013/08/27 22:54:04]

You don't need this check. Just call WriteEverything().

[Adam Rice, 2013/08/28 12:01:11]

Done.

+      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);

[szym, 2013/08/27 22:54:04]

BUG: It's possible to end up in here with result == 0 and return OK. This does
not match the spec in websocket_stream.h

Although I'm not sure Socket::Write could return 0.

[Adam Rice, 2013/08/28 12:01:11]

I've put in a DCHECK for result != 0 for the time being. socket.h has the
unhelpful comment "The return value when the connection is closed is undefined,
and may be OS dependent." It is possible that a Socket implementer could take
that as license to use 0 for "Connection Closed", but I hope no-one would do
that.

+  }
+}
+
+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