Support WebSocket per-message deflate extension in http server.

net/server/web_socket_encoder.cc

+// Copyright 2014 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/server/web_socket_encoder.h"
+
+#include "base/logging.h"
+#include "base/strings/string_number_conversions.h"
+#include "base/strings/stringprintf.h"
+#include "net/base/io_buffer.h"
+#include "net/websockets/websocket_extension_parser.h"
+
+namespace net {
+
+const char WebSocketEncoder::kClientExtensions[] =
+    "Sec-WebSocket-Extensions: permessage-deflate; client_max_window_bits";
+
+namespace {
+
+const int kInflaterChunkSize = 16 * 1024;
+
+// Constants for hybi-10 frame format.
+
+typedef int OpCode;
+
+const OpCode kOpCodeContinuation = 0x0;
+const OpCode kOpCodeText = 0x1;
+const OpCode kOpCodeBinary = 0x2;
+const OpCode kOpCodeClose = 0x8;
+const OpCode kOpCodePing = 0x9;
+const OpCode kOpCodePong = 0xA;
+
+const unsigned char kFinalBit = 0x80;
+const unsigned char kReserved1Bit = 0x40;
+const unsigned char kReserved2Bit = 0x20;
+const unsigned char kReserved3Bit = 0x10;
+const unsigned char kOpCodeMask = 0xF;
+const unsigned char kMaskBit = 0x80;
+const unsigned char kPayloadLengthMask = 0x7F;
+
+const size_t kMaxSingleBytePayloadLength = 125;
+const size_t kTwoBytePayloadLengthField = 126;
+const size_t kEightBytePayloadLengthField = 127;
+const size_t kMaskingKeyWidthInBytes = 4;
+
+WebSocket::ParseResult DecodeFrameHybi17(const base::StringPiece& frame,

[yhirano, 2014/12/08 05:06:29]

You can use WebSocketFrameParser in net/websockets/websocket_frame_parser.h and
WriteWebSocketFrameHeader in net/websockets/websocket_frame.h.

[dgozman, 2014/12/08 17:53:00]

We can try to completely switch over to net/websockets infrastructure, but
definitely not in this patch.
I'm worried it's too heavyweight for our usecase, but I may be wrong here.

[yhirano, 2014/12/09 11:12:01]

I see, no problem.

+                                         bool client_frame,
+                                         int* bytes_consumed,
+                                         std::string* output,
+                                         bool* compressed) {
+  size_t data_length = frame.length();
+  if (data_length < 2)
+    return WebSocket::FRAME_INCOMPLETE;
+
+  const char* buffer_begin = const_cast<char*>(frame.data());
+  const char* p = buffer_begin;
+  const char* buffer_end = p + data_length;
+
+  unsigned char first_byte = *p++;
+  unsigned char second_byte = *p++;
+
+  bool final = (first_byte & kFinalBit) != 0;
+  bool reserved1 = (first_byte & kReserved1Bit) != 0;
+  bool reserved2 = (first_byte & kReserved2Bit) != 0;
+  bool reserved3 = (first_byte & kReserved3Bit) != 0;
+  int op_code = first_byte & kOpCodeMask;
+  bool masked = (second_byte & kMaskBit) != 0;
+  *compressed = reserved1;
+  if (!final || reserved2 || reserved3)
+    return WebSocket::FRAME_ERROR;  // Only compression extension is supported.
+
+  bool closed = false;
+  switch (op_code) {
+    case kOpCodeClose:
+      closed = true;
+      break;
+    case kOpCodeText:
+      break;
+    case kOpCodeBinary:        // We don't support binary frames yet.
+    case kOpCodeContinuation:  // We don't support binary frames yet.
+    case kOpCodePing:          // We don't support binary frames yet.
+    case kOpCodePong:          // We don't support binary frames yet.
+    default:
+      return WebSocket::FRAME_ERROR;
+  }
+
+  if (client_frame && !masked)  // In Hybi-17 spec client MUST mask his frame.
+    return WebSocket::FRAME_ERROR;
+
+  uint64 payload_length64 = second_byte & kPayloadLengthMask;
+  if (payload_length64 > kMaxSingleBytePayloadLength) {
+    int extended_payload_length_size;
+    if (payload_length64 == kTwoBytePayloadLengthField)
+      extended_payload_length_size = 2;
+    else {
+      DCHECK(payload_length64 == kEightBytePayloadLengthField);
+      extended_payload_length_size = 8;
+    }
+    if (buffer_end - p < extended_payload_length_size)
+      return WebSocket::FRAME_INCOMPLETE;
+    payload_length64 = 0;
+    for (int i = 0; i < extended_payload_length_size; ++i) {
+      payload_length64 <<= 8;
+      payload_length64 |= static_cast<unsigned char>(*p++);
+    }
+  }
+
+  size_t actual_masking_key_length = masked ? kMaskingKeyWidthInBytes : 0;
+  static const uint64 max_payload_length = 0x7FFFFFFFFFFFFFFFull;
+  static size_t max_length = std::numeric_limits<size_t>::max();
+  if (payload_length64 > max_payload_length ||
+      payload_length64 + actual_masking_key_length > max_length) {
+    // WebSocket frame length too large.
+    return WebSocket::FRAME_ERROR;
+  }
+  size_t payload_length = static_cast<size_t>(payload_length64);
+
+  size_t total_length = actual_masking_key_length + payload_length;
+  if (static_cast<size_t>(buffer_end - p) < total_length)
+    return WebSocket::FRAME_INCOMPLETE;
+
+  if (masked) {
+    output->resize(payload_length);
+    const char* masking_key = p;
+    char* payload = const_cast<char*>(p + kMaskingKeyWidthInBytes);
+    for (size_t i = 0; i < payload_length; ++i)  // Unmask the payload.
+      (*output)[i] = payload[i] ^ masking_key[i % kMaskingKeyWidthInBytes];
+  } else {
+    output->assign(p, p + payload_length);
+  }
+
+  size_t pos = p + actual_masking_key_length + payload_length - buffer_begin;
+  *bytes_consumed = pos;
+  return closed ? WebSocket::FRAME_CLOSE : WebSocket::FRAME_OK;
+}
+
+void EncodeFrameHybi17(const std::string& message,
+                       int masking_key,
+                       bool compressed,
+                       std::string* output) {
+  std::vector<char> frame;
+  OpCode op_code = kOpCodeText;
+  size_t data_length = message.length();
+
+  int reserved1 = compressed ? kReserved1Bit : 0;
+  frame.push_back(kFinalBit | op_code | reserved1);
+  char mask_key_bit = masking_key != 0 ? kMaskBit : 0;
+  if (data_length <= kMaxSingleBytePayloadLength)
+    frame.push_back(data_length | mask_key_bit);
+  else if (data_length <= 0xFFFF) {
+    frame.push_back(kTwoBytePayloadLengthField | mask_key_bit);
+    frame.push_back((data_length & 0xFF00) >> 8);
+    frame.push_back(data_length & 0xFF);
+  } else {
+    frame.push_back(kEightBytePayloadLengthField | mask_key_bit);
+    char extended_payload_length[8];
+    size_t remaining = data_length;
+    // Fill the length into extended_payload_length in the network byte order.
+    for (int i = 0; i < 8; ++i) {
+      extended_payload_length[7 - i] = remaining & 0xFF;
+      remaining >>= 8;
+    }
+    frame.insert(frame.end(), extended_payload_length,
+                 extended_payload_length + 8);
+    DCHECK(!remaining);
+  }
+
+  const char* data = const_cast<char*>(message.data());
+  if (masking_key != 0) {
+    const char* mask_bytes = reinterpret_cast<char*>(&masking_key);
+    frame.insert(frame.end(), mask_bytes, mask_bytes + 4);
+    for (size_t i = 0; i < data_length; ++i)  // Mask the payload.
+      frame.push_back(data[i] ^ mask_bytes[i % kMaskingKeyWidthInBytes]);
+  } else {
+    frame.insert(frame.end(), data, data + data_length);
+  }
+  *output = std::string(&frame[0], frame.size());
+}
+
+}  // anonymous namespace
+
+// static
+WebSocketEncoder* WebSocketEncoder::CreateServer(
+    const std::string& request_extensions,
+    std::string* response_extensions) {
+  bool deflate;
+  int client_window_bits;
+  int server_window_bits;
+  ParseExtensions(request_extensions, &deflate, &client_window_bits,
+                  &server_window_bits);
+
+  if (deflate) {
+    *response_extensions = base::StringPrintf(
+        "permessage-deflate; server_max_window_bits=%d; "
+        "client_max_window_bits=%d",
+        server_window_bits, client_window_bits);
+    return new WebSocketEncoder(true /* is_server */, server_window_bits,
+                                client_window_bits);
+  } else {
+    *response_extensions = std::string();
+    return new WebSocketEncoder(true /* is_server */);
+  }
+}
+
+// static
+WebSocketEncoder* WebSocketEncoder::CreateClient(
+    const std::string& response_extensions) {
+  bool deflate;
+  int client_window_bits;
+  int server_window_bits;
+  ParseExtensions(response_extensions, &deflate, &client_window_bits,
+                  &server_window_bits);
+
+  if (deflate) {
+    return new WebSocketEncoder(false /* is_server */, client_window_bits,
+                                server_window_bits);
+  } else {
+    return new WebSocketEncoder(false /* is_server */);
+  }
+}
+
+// static
+void WebSocketEncoder::ParseExtensions(const std::string& extensions,
+                                       bool* deflate,
+                                       int* client_window_bits,
+                                       int* server_window_bits) {
+  *deflate = false;
+  *client_window_bits = 15;
+  *server_window_bits = 15;
+
+  if (extensions.empty())
+    return;
+
+  // TODO(dgozman): split extensions header if another extension is introduced.
+  WebSocketExtensionParser parser;
+  parser.Parse(extensions);
+  if (parser.has_error())
+    return;
+  if (parser.extension().name() != "permessage-deflate")
+    return;
+
+  const std::vector<WebSocketExtension::Parameter>& parameters =
+      parser.extension().parameters();
+  for (const auto& param : parameters) {
+    const std::string& name = param.name();
+    if (name == "client_max_window_bits" && param.HasValue()) {
+      int bits = 0;
+      if (base::StringToInt(param.value(), &bits) && bits >= 8 && bits <= 15)
+        *client_window_bits = bits;
+    }
+    if (name == "server_max_window_bits" && param.HasValue()) {
+      int bits = 0;
+      if (base::StringToInt(param.value(), &bits) && bits >= 8 && bits <= 15)
+        *server_window_bits = bits;
+    }
+  }
+  *deflate = true;
+}
+
+WebSocketEncoder::WebSocketEncoder(bool is_server) : is_server_(is_server) {
+}
+
+WebSocketEncoder::WebSocketEncoder(bool is_server,
+                                   int deflate_bits,
+                                   int inflate_bits)
+    : is_server_(is_server) {
+  deflater_.reset(new WebSocketDeflater(WebSocketDeflater::TAKE_OVER_CONTEXT));

[yhirano, 2014/12/08 05:06:28]

If you don't verify peer's preference on context takeover control (i.e.
*_no_context_takeover parameters), DO_NOT_TAKE_OVER_CONTEXT should be used.

[dgozman, 2014/12/08 17:53:00]

Added no_context_takeover support.

+  deflater_->Initialize(deflate_bits);

[yhirano, 2014/12/08 05:06:29]

WebSocketDeflater::Initialize may fail.
If you ignore the possibility, it's OK, but adding a comment or DCHECK may be
helpful.

[dgozman, 2014/12/08 17:53:00]

Fixed.

+
+  inflater_.reset(
+      new WebSocketInflater(kInflaterChunkSize, kInflaterChunkSize));
+  inflater_->Initialize(inflate_bits);

[yhirano, 2014/12/08 05:06:29]

ditto

+}
+
+WebSocketEncoder::~WebSocketEncoder() {
+}
+
+WebSocket::ParseResult WebSocketEncoder::DecodeFrame(
+    const base::StringPiece& frame,
+    int* bytes_consumed,
+    std::string* output) {
+  bool compressed;
+  WebSocket::ParseResult result =
+      DecodeFrameHybi17(frame, is_server_, bytes_consumed, output, &compressed);
+  if (result == WebSocket::FRAME_OK && compressed) {
+    if (!Inflate(output))
+      result = WebSocket::FRAME_ERROR;
+  }
+  return result;
+}
+
+void WebSocketEncoder::EncodeFrame(const std::string& frame,
+                                   int masking_key,
+                                   std::string* output) {
+  std::string compressed;
+  if (Deflate(frame, &compressed))
+    EncodeFrameHybi17(compressed, masking_key, true, output);
+  else
+    EncodeFrameHybi17(frame, masking_key, false, output);
+}
+
+bool WebSocketEncoder::Inflate(std::string* message) {
+  if (!inflater_)
+    return false;
+
+  const char* data = message->data();
+  int length = message->length();
+  while (length > 0) {

[yhirano, 2014/12/08 05:06:29]

I don't understand why this loop is needed: Doesn't inflater_->AddBytes(data,
length) work?

[dgozman, 2014/12/08 17:53:00]

This is not needed. I just read the interface wrong.

+    int length_to_add = std::min(length, kInflaterChunkSize);
+    if (!inflater_->AddBytes(data, length_to_add))
+      return false;
+    data += length_to_add;
+    length -= length_to_add;
+  }
+
+  if (!inflater_->Finish())
+    return false;
+
+  std::vector<char> output;
+  while (inflater_->CurrentOutputSize() > 0) {
+    scoped_refptr<IOBufferWithSize> chunk =
+        inflater_->GetOutput(inflater_->CurrentOutputSize());
+    if (!chunk.get())
+      return false;
+    output.insert(output.end(), chunk->data(), chunk->data() + chunk->size());
+  }
+
+  *message = std::string(&output[0], output.size());

[yhirano, 2014/12/08 05:06:29]

This is problematic when output is empty.

[dgozman, 2014/12/08 17:53:00]

Nice catch! Done.

+  return true;
+}
+
+bool WebSocketEncoder::Deflate(const std::string& message,
+                               std::string* output) {
+  if (!deflater_)
+    return false;
+  if (!deflater_->AddBytes(message.data(), message.length())) {
+    deflater_->Finish();
+    return false;
+  }
+  if (!deflater_->Finish())
+    return false;
+  scoped_refptr<IOBufferWithSize> buffer =
+      deflater_->GetOutput(deflater_->CurrentOutputSize());
+  if (!buffer.get())
+    return false;
+  *output = std::string(buffer->data(), buffer->size());
+  return true;
+}
+
+}  // namespace net