text: // 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 "google_apis/gcm/engine/connection_handler_impl.h"
  
  #include "base/message_loop/message_loop.h"
  #include "google/protobuf/io/coded_stream.h"
  #include "google_apis/gcm/base/mcs_util.h"
  #include "google_apis/gcm/base/socket_stream.h"
  #include "google_apis/gcm/protocol/mcs.pb.h"
  #include "net/base/net_errors.h"
  #include "net/socket/stream_socket.h"
  
  using namespace google::protobuf::io;
  
  namespace gcm {
  
  namespace {
  
  // # of bytes a MCS version packet consumes.
  const int kVersionPacketLen = 1;
  // # of bytes a tag packet consumes.
  const int kTagPacketLen = 1;
  // Max # of bytes a length packet consumes.
  const int kSizePacketLenMin = 1;
  const int kSizePacketLenMax = 2;
  
  // The current MCS protocol version.
  const int kMCSVersion = 41;
  
  }  // namespace
  
  ConnectionHandlerImpl::ConnectionHandlerImpl(
      base::TimeDelta read_timeout,
      const ProtoReceivedCallback& read_callback,
      const ProtoSentCallback& write_callback,
      const ConnectionChangedCallback& connection_callback)
      : read_timeout_(read_timeout),
        socket_(NULL),
        handshake_complete_(false),
        message_tag_(0),
        message_size_(0),
        read_callback_(read_callback),
        write_callback_(write_callback),
        connection_callback_(connection_callback),
        weak_ptr_factory_(this) {
  }
  
  ConnectionHandlerImpl::~ConnectionHandlerImpl() {
  }
  
  void ConnectionHandlerImpl::Init(
      const mcs_proto::LoginRequest& login_request,
      net::StreamSocket* socket) {
    DCHECK(!read_callback_.is_null());
    DCHECK(!write_callback_.is_null());
    DCHECK(!connection_callback_.is_null());
  
    // Invalidate any previously outstanding reads.
    weak_ptr_factory_.InvalidateWeakPtrs();
  
    handshake_complete_ = false;
    message_tag_ = 0;
    message_size_ = 0;
    socket_ = socket;
    input_stream_.reset(new SocketInputStream(socket_));
    output_stream_.reset(new SocketOutputStream(socket_));
  
    Login(login_request);
  }
  
  void ConnectionHandlerImpl::Reset() {
    CloseConnection();
  }
  
  bool ConnectionHandlerImpl::CanSendMessage() const {
    return handshake_complete_ && output_stream_.get() &&
        output_stream_->GetState() == SocketOutputStream::EMPTY;
  }
  
  void ConnectionHandlerImpl::SendMessage(
      const google::protobuf::MessageLite& message) {
    DCHECK_EQ(output_stream_->GetState(), SocketOutputStream::EMPTY);
    DCHECK(handshake_complete_);
  
    {
      CodedOutputStream coded_output_stream(output_stream_.get());
      DVLOG(1) << "Writing proto of size " << message.ByteSize();
      int tag = GetMCSProtoTag(message);
      DCHECK_NE(tag, -1);
      coded_output_stream.WriteRaw(&tag, 1);
      coded_output_stream.WriteVarint32(message.ByteSize());
      message.SerializeToCodedStream(&coded_output_stream);
    }
  
    if (output_stream_->Flush(
            base::Bind(&ConnectionHandlerImpl::OnMessageSent,
                       weak_ptr_factory_.GetWeakPtr())) != net::ERR_IO_PENDING) {
      OnMessageSent();
    }
  }
  
  void ConnectionHandlerImpl::Login(
      const google::protobuf::MessageLite& login_request) {
    DCHECK_EQ(output_stream_->GetState(), SocketOutputStream::EMPTY);
  
    const char version_byte[1] = {kMCSVersion};
    const char login_request_tag[1] = {kLoginRequestTag};
    {
      CodedOutputStream coded_output_stream(output_stream_.get());
      coded_output_stream.WriteRaw(version_byte, 1);
      coded_output_stream.WriteRaw(login_request_tag, 1);
      coded_output_stream.WriteVarint32(login_request.ByteSize());
      login_request.SerializeToCodedStream(&coded_output_stream);
    }
  
    if (output_stream_->Flush(
            base::Bind(&ConnectionHandlerImpl::OnMessageSent,
                       weak_ptr_factory_.GetWeakPtr())) != net::ERR_IO_PENDING) {
      base::MessageLoop::current()->PostTask(
          FROM_HERE,
          base::Bind(&ConnectionHandlerImpl::OnMessageSent,
                     weak_ptr_factory_.GetWeakPtr()));
    }
  
    read_timeout_timer_.Start(FROM_HERE,
                              read_timeout_,
                              base::Bind(&ConnectionHandlerImpl::OnTimeout,
                                         weak_ptr_factory_.GetWeakPtr()));
    WaitForData(MCS_VERSION_TAG_AND_SIZE);
  }
  
  void ConnectionHandlerImpl::OnMessageSent() {
    if (!output_stream_.get()) {
      // The connection has already been closed. Just return.
      DCHECK(!input_stream_.get());
      DCHECK(!read_timeout_timer_.IsRunning());
      return;
    }
  
    if (output_stream_->GetState() != SocketOutputStream::EMPTY) {
      int last_error = output_stream_->last_error();
      CloseConnection();
      // If the socket stream had an error, plumb it up, else plumb up FAILED.
      if (last_error == net::OK)
        last_error = net::ERR_FAILED;
      connection_callback_.Run(last_error);
      return;
    }
  
    write_callback_.Run();
  }
  
  void ConnectionHandlerImpl::GetNextMessage() {
    DCHECK(SocketInputStream::EMPTY == input_stream_->GetState() ||
           SocketInputStream::READY == input_stream_->GetState());
    message_tag_ = 0;
    message_size_ = 0;
  
    WaitForData(MCS_TAG_AND_SIZE);
  }
  
  void ConnectionHandlerImpl::WaitForData(ProcessingState state) {
    DVLOG(1) << "Waiting for MCS data: state == " << state;
  
    if (!input_stream_) {
      // The connection has already been closed. Just return.
      DCHECK(!output_stream_.get());
      DCHECK(!read_timeout_timer_.IsRunning());
      return;
    }
  
    if (input_stream_->GetState() != SocketInputStream::EMPTY &&
        input_stream_->GetState() != SocketInputStream::READY) {
      // An error occurred.
      int last_error = output_stream_->last_error();
      CloseConnection();
      // If the socket stream had an error, plumb it up, else plumb up FAILED.
      if (last_error == net::OK)
        last_error = net::ERR_FAILED;
      connection_callback_.Run(last_error);
      return;
    }
  
    // Used to determine whether a Socket::Read is necessary.
    size_t min_bytes_needed = 0;
    // Used to limit the size of the Socket::Read.
    size_t max_bytes_needed = 0;
  
    switch(state) {
      case MCS_VERSION_TAG_AND_SIZE:
        min_bytes_needed = kVersionPacketLen + kTagPacketLen + kSizePacketLenMin;
        max_bytes_needed = kVersionPacketLen + kTagPacketLen + kSizePacketLenMax;
        break;
      case MCS_TAG_AND_SIZE:
        min_bytes_needed = kTagPacketLen + kSizePacketLenMin;
        max_bytes_needed = kTagPacketLen + kSizePacketLenMax;
        break;
      case MCS_FULL_SIZE:
        // If in this state, the minimum size packet length must already have been
        // insufficient, so set both to the max length.
        min_bytes_needed = kSizePacketLenMax;
        max_bytes_needed = kSizePacketLenMax;
        break;
      case MCS_PROTO_BYTES:
        read_timeout_timer_.Reset();
        // No variability in the message size, set both to the same.
        min_bytes_needed = message_size_;
        max_bytes_needed = message_size_;
        break;
      default:
        NOTREACHED();
    }
    DCHECK_GE(max_bytes_needed, min_bytes_needed);
  
    size_t unread_byte_count = input_stream_->UnreadByteCount();
    if (min_bytes_needed > unread_byte_count &&
        input_stream_->Refresh(
            base::Bind(&ConnectionHandlerImpl::WaitForData,
                       weak_ptr_factory_.GetWeakPtr(),
                       state),
            max_bytes_needed - unread_byte_count) == net::ERR_IO_PENDING) {
      return;
    }
  
    // Check for refresh errors.
    if (input_stream_->GetState() != SocketInputStream::READY) {
      // An error occurred.
      int last_error = input_stream_->last_error();
      CloseConnection();
      // If the socket stream had an error, plumb it up, else plumb up FAILED.
      if (last_error == net::OK)
        last_error = net::ERR_FAILED;
      connection_callback_.Run(last_error);
      return;
    }
  
    // Check whether read is complete, or needs to be continued (
    // SocketInputStream::Refresh can finish without reading all the data).
    if (input_stream_->UnreadByteCount() < min_bytes_needed) {
      DVLOG(1) << "Socket read finished prematurely. Waiting for "
               << min_bytes_needed - input_stream_->UnreadByteCount()
               << " more bytes.";
      base::MessageLoop::current()->PostTask(
          FROM_HERE,
          base::Bind(&ConnectionHandlerImpl::WaitForData,
                     weak_ptr_factory_.GetWeakPtr(),
                     MCS_PROTO_BYTES));
      return;
    }
  
    // Received enough bytes, process them.
    DVLOG(1) << "Processing MCS data: state == " << state;
    switch(state) {
      case MCS_VERSION_TAG_AND_SIZE:
        OnGotVersion();
        break;
      case MCS_TAG_AND_SIZE:
        OnGotMessageTag();
        break;
      case MCS_FULL_SIZE:
        OnGotMessageSize();
        break;
      case MCS_PROTO_BYTES:
        OnGotMessageBytes();
        break;
      default:
        NOTREACHED();
    }
  }
  
  void ConnectionHandlerImpl::OnGotVersion() {
    uint8 version = 0;
    {
      CodedInputStream coded_input_stream(input_stream_.get());
      coded_input_stream.ReadRaw(&version, 1);
    }
    // TODO(zea): remove this when the server is ready.
    if (version < kMCSVersion && version != 38) {
      LOG(ERROR) << "Invalid GCM version response: " << static_cast<int>(version);
      connection_callback_.Run(net::ERR_FAILED);
      return;
    }
  
    input_stream_->RebuildBuffer();
  
    // Process the LoginResponse message tag.
    OnGotMessageTag();
  }
  
  void ConnectionHandlerImpl::OnGotMessageTag() {
    if (input_stream_->GetState() != SocketInputStream::READY) {
      LOG(ERROR) << "Failed to receive protobuf tag.";
      read_callback_.Run(scoped_ptr<google::protobuf::MessageLite>());
      return;
    }
  
    {
      CodedInputStream coded_input_stream(input_stream_.get());
      coded_input_stream.ReadRaw(&message_tag_, 1);
    }
  
    DVLOG(1) << "Received proto of type "
             << static_cast<unsigned int>(message_tag_);
  
    if (!read_timeout_timer_.IsRunning()) {
      read_timeout_timer_.Start(FROM_HERE,
                                read_timeout_,
                                base::Bind(&ConnectionHandlerImpl::OnTimeout,
                                           weak_ptr_factory_.GetWeakPtr()));
    }
    OnGotMessageSize();
  }
  
  void ConnectionHandlerImpl::OnGotMessageSize() {
    if (input_stream_->GetState() != SocketInputStream::READY) {
      LOG(ERROR) << "Failed to receive message size.";
      read_callback_.Run(scoped_ptr<google::protobuf::MessageLite>());
      return;
    }
  
    bool need_another_byte = false;
    int prev_byte_count = input_stream_->ByteCount();
    {
      CodedInputStream coded_input_stream(input_stream_.get());
      if (!coded_input_stream.ReadVarint32(&message_size_))
        need_another_byte = true;
    }
  
    if (need_another_byte) {
      DVLOG(1) << "Expecting another message size byte.";
      if (prev_byte_count >= kSizePacketLenMax) {
        // Already had enough bytes, something else went wrong.
        LOG(ERROR) << "Failed to process message size.";
        read_callback_.Run(scoped_ptr<google::protobuf::MessageLite>());
        return;
      }
      // Back up by the amount read (should always be 1 byte).
      int bytes_read = prev_byte_count - input_stream_->ByteCount();
      DCHECK_EQ(bytes_read, 1);
      input_stream_->BackUp(bytes_read);
      WaitForData(MCS_FULL_SIZE);
      return;
    }
  
    DVLOG(1) << "Proto size: " << message_size_;
  
    if (message_size_ > 0)
      WaitForData(MCS_PROTO_BYTES);
    else
      OnGotMessageBytes();
  }
  
  void ConnectionHandlerImpl::OnGotMessageBytes() {
    read_timeout_timer_.Stop();
    scoped_ptr<google::protobuf::MessageLite> protobuf(
        BuildProtobufFromTag(message_tag_));
    // Messages with no content are valid; just use the default protobuf for
    // that tag.
    if (protobuf.get() && message_size_ == 0) {
      base::MessageLoop::current()->PostTask(
          FROM_HERE,
          base::Bind(&ConnectionHandlerImpl::GetNextMessage,
                     weak_ptr_factory_.GetWeakPtr()));
      read_callback_.Run(protobuf.Pass());
      return;
    }
  
    if (!protobuf.get() ||
        input_stream_->GetState() != SocketInputStream::READY) {
      LOG(ERROR) << "Failed to extract protobuf bytes of type "
                 << static_cast<unsigned int>(message_tag_);
      // Reset the connection.
      connection_callback_.Run(net::ERR_FAILED);
      return;
    }
  
    {
      CodedInputStream coded_input_stream(input_stream_.get());
      if (!protobuf->ParsePartialFromCodedStream(&coded_input_stream)) {
        LOG(ERROR) << "Unable to parse GCM message of type "
                   << static_cast<unsigned int>(message_tag_);
        // Reset the connection.
        connection_callback_.Run(net::ERR_FAILED);
        return;
      }
    }
  
    input_stream_->RebuildBuffer();
    base::MessageLoop::current()->PostTask(
        FROM_HERE,
        base::Bind(&ConnectionHandlerImpl::GetNextMessage,
                   weak_ptr_factory_.GetWeakPtr()));
    if (message_tag_ == kLoginResponseTag) {
      if (handshake_complete_) {
        LOG(ERROR) << "Unexpected login response.";
      } else {
        handshake_complete_ = true;
        DVLOG(1) << "GCM Handshake complete.";
        connection_callback_.Run(net::OK);
      }
    }
    read_callback_.Run(protobuf.Pass());
  }
  
  void ConnectionHandlerImpl::OnTimeout() {
    LOG(ERROR) << "Timed out waiting for GCM Protocol buffer.";
    CloseConnection();
    connection_callback_.Run(net::ERR_TIMED_OUT);
  }
  
  void ConnectionHandlerImpl::CloseConnection() {
    DVLOG(1) << "Closing connection.";
    read_timeout_timer_.Stop();
    if (socket_)
      socket_->Disconnect();
    socket_ = NULL;
    handshake_complete_ = false;
    message_tag_ = 0;
    message_size_ = 0;
    input_stream_.reset();
    output_stream_.reset();
    weak_ptr_factory_.InvalidateWeakPtrs();
  }
  
  }  // namespace gcm
  