// 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 or algorithm to fulfill our
// requirements:
//  1. We would like to use minimal memory on low-bandwidth or idle connections
//  2. We would like to read as close to line speed as possible on
//     high-bandwidth connections
//  3. We can't afford to cause jank on the IO thread by copying large buffers
//     around
//  4. We would like to hit any sweet-spots that might exist in terms of network
//     packet sizes / encryption block sizes / IPC alignment issues, etc.
const int kReadBufferSize = 32 * 1024;

}  // namespace

WebSocketBasicStream::WebSocketBasicStream(
    scoped_ptr<ClientSocketHandle> connection)
    : read_buffer_(new IOBufferWithSize(kReadBufferSize)),
      connection_(connection.Pass()),
      generate_websocket_masking_key_(&GenerateWebSocketMaskingKey) {
  DCHECK(connection_->is_initialized());
}

WebSocketBasicStream::~WebSocketBasicStream() { Close(); }

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);
    // We cannot simply copy the data into read_buffer_, as it might be too
    // large.
    scoped_refptr<GrowableIOBuffer> buffered_data;
    buffered_data.swap(http_read_buffer_);
    DCHECK(http_read_buffer_.get() == NULL);
    if (!parser_.Decode(buffered_data->StartOfBuffer(),
                        buffered_data->offset(),
                        frame_chunks))
      return WebSocketErrorToNetError(parser_.websocket_error());
    if (!frame_chunks->empty())
      return OK;
  }

  // Run until socket stops giving us data or we get some chunks.
  while (true) {
    // base::Unretained(this) here is safe because net::Socket guarantees not to
    // call any callbacks after Disconnect(), which we call from the
    // destructor. The caller of ReadFrames() is required to keep |frame_chunks|
    // valid.
    int result = connection_->socket()
                     ->Read(read_buffer_.get(),
                            read_buffer_->size(),
                            base::Bind(&WebSocketBasicStream::OnReadComplete,
                                       base::Unretained(this),
                                       base::Unretained(frame_chunks),
                                       callback));
    if (result == ERR_IO_PENDING)
      return result;
    result = HandleReadResult(result, frame_chunks);
    if (result != ERR_IO_PENDING)
      return 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;
  const int kMaximumTotalSize = std::numeric_limits<int>::max();
  int total_size = 0;
  for (Iterator it = frame_chunks->begin(); it != frame_chunks->end(); ++it) {
    WebSocketFrameChunk* chunk = *it;
    DCHECK(chunk->header)
        << "Only complete frames are supported by WebSocketBasicStream";
    DCHECK(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(kMaximumTotalSize - total_size, chunk_size)
        << "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);
}

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();
}

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::OnWriteComplete,
                                        base::Unretained(this),
                                        buffer,
                                        callback));
    if (result > 0) {
      buffer->DidConsume(result);
    } else {
      return result;
    }
  }
  return OK;
}

void WebSocketBasicStream::OnWriteComplete(
    const scoped_refptr<DrainableIOBuffer>& buffer,
    const CompletionCallback& callback,
    int result) {
  if (result < 0) {
    DCHECK(result != ERR_IO_PENDING);
    callback.Run(result);
    return;
  }

  DCHECK(result != 0);
  buffer->DidConsume(result);
  result = WriteEverything(buffer, callback);
  if (result != ERR_IO_PENDING)
    callback.Run(result);
}

int WebSocketBasicStream::HandleReadResult(
    int result,
    ScopedVector<WebSocketFrameChunk>* 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_->data(), result, frame_chunks))
    return WebSocketErrorToNetError(parser_.websocket_error());
  if (!frame_chunks->empty())
    return OK;
  return ERR_IO_PENDING;
}

void WebSocketBasicStream::OnReadComplete(
    ScopedVector<WebSocketFrameChunk>* frame_chunks,
    const CompletionCallback& callback,
    int result) {
  result = HandleReadResult(result, frame_chunks);
  if (result == ERR_IO_PENDING)
    result = ReadFrames(frame_chunks, callback);
  if (result != ERR_IO_PENDING)
    callback.Run(result);
}

}  // namespace net