// Copyright 2015 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 "build/build_config.h"

#include <stddef.h>
#include <stdint.h>

#include <limits>
#include <set>

#include "base/run_loop.h"
#include "ipc/attachment_broker.h"
#include "ipc/brokerable_attachment.h"
#include "ipc/ipc_channel_reader.h"
#include "ipc/placeholder_brokerable_attachment.h"
#include "testing/gtest/include/gtest/gtest.h"

// Whether IPC::Message::FindNext() can determine message size for
// partial messages. The condition is from FindNext() implementation.
#if USE_ATTACHMENT_BROKER
#define MESSAGE_FINDNEXT_PARTIAL 0
#else
#define MESSAGE_FINDNEXT_PARTIAL 1
#endif

namespace IPC {
namespace internal {

namespace {

#if USE_ATTACHMENT_BROKER

class MockAttachment : public BrokerableAttachment {
 public:
  MockAttachment() {}
  MockAttachment(BrokerableAttachment::AttachmentId id)
      : BrokerableAttachment(id) {}

#if defined(OS_POSIX)
  base::PlatformFile TakePlatformFile() override {
    return base::PlatformFile();
  }
#endif  // OS_POSIX

  BrokerableType GetBrokerableType() const override { return WIN_HANDLE; }

 private:
  ~MockAttachment() override {}
};

class MockAttachmentBroker : public AttachmentBroker {
 public:
  typedef std::set<scoped_refptr<BrokerableAttachment>> AttachmentSet;

  bool SendAttachmentToProcess(
      const scoped_refptr<BrokerableAttachment>& attachment,
      base::ProcessId destination_process) override {
    return false;
  }

  bool OnMessageReceived(const Message& message) override { return false; }

  void AddAttachment(scoped_refptr<BrokerableAttachment> attachment) {
    get_attachments()->push_back(attachment);
    NotifyObservers(attachment->GetIdentifier());
  }
};

#endif  // USE_ATTACHMENT_BROKER

class MockChannelReader : public ChannelReader {
 public:
  MockChannelReader()
      : ChannelReader(nullptr), last_dispatched_message_(nullptr) {}

  ReadState ReadData(char* buffer, int buffer_len, int* bytes_read) override {
    if (data_.empty())
      return READ_PENDING;

    size_t read_len = std::min(static_cast<size_t>(buffer_len), data_.size());
    memcpy(buffer, data_.data(), read_len);
    *bytes_read = static_cast<int>(read_len);
    data_.erase(0, read_len);
    return READ_SUCCEEDED;
  }

  bool ShouldDispatchInputMessage(Message* msg) override { return true; }

  bool GetNonBrokeredAttachments(Message* msg) override { return true; }

  bool DidEmptyInputBuffers() override { return true; }

  void HandleInternalMessage(const Message& msg) override {}

  void DispatchMessage(Message* m) override { last_dispatched_message_ = m; }

  base::ProcessId GetSenderPID() override { return base::kNullProcessId; }

  bool IsAttachmentBrokerEndpoint() override { return false; }

  AttachmentBroker* GetAttachmentBroker() override { return broker_; }

  // This instance takes ownership of |m|.
  void AddMessageForDispatch(Message* m) {
    get_queued_messages()->push_back(m);
  }

  Message* get_last_dispatched_message() { return last_dispatched_message_; }

  void set_broker(AttachmentBroker* broker) { broker_ = broker; }

  void AppendData(const void* data, size_t size) {
    data_.append(static_cast<const char*>(data), size);
  }

  void AppendMessageData(const Message& message) {
    AppendData(message.data(), message.size());
  }

 private:
  Message* last_dispatched_message_;
  AttachmentBroker* broker_;
  std::string data_;
};

class ExposedMessage: public Message {
 public:
  using Message::Header;
  using Message::header;
};

// Payload that makes messages large
const size_t LargePayloadSize = Channel::kMaximumReadBufferSize * 3 / 2;

}  // namespace

#if USE_ATTACHMENT_BROKER

TEST(ChannelReaderTest, AttachmentAlreadyBrokered) {
  MockAttachmentBroker broker;
  MockChannelReader reader;
  reader.set_broker(&broker);
  scoped_refptr<MockAttachment> attachment(new MockAttachment);
  broker.AddAttachment(attachment);

  Message* m = new Message;
  PlaceholderBrokerableAttachment* needs_brokering_attachment =
      new PlaceholderBrokerableAttachment(attachment->GetIdentifier());
  EXPECT_TRUE(m->WriteAttachment(needs_brokering_attachment));
  reader.AddMessageForDispatch(m);
  EXPECT_EQ(ChannelReader::DISPATCH_FINISHED, reader.DispatchMessages());
  EXPECT_EQ(m, reader.get_last_dispatched_message());
}

TEST(ChannelReaderTest, AttachmentNotYetBrokered) {
  scoped_ptr<base::MessageLoop> message_loop(new base::MessageLoopForIO());

  MockAttachmentBroker broker;
  MockChannelReader reader;
  reader.set_broker(&broker);
  scoped_refptr<MockAttachment> attachment(new MockAttachment);

  Message* m = new Message;
  PlaceholderBrokerableAttachment* needs_brokering_attachment =
      new PlaceholderBrokerableAttachment(attachment->GetIdentifier());
  EXPECT_TRUE(m->WriteAttachment(needs_brokering_attachment));
  reader.AddMessageForDispatch(m);
  EXPECT_EQ(ChannelReader::DISPATCH_WAITING_ON_BROKER,
            reader.DispatchMessages());
  EXPECT_EQ(nullptr, reader.get_last_dispatched_message());

  broker.AddAttachment(attachment);
  base::RunLoop run_loop;
  run_loop.RunUntilIdle();

  EXPECT_EQ(m, reader.get_last_dispatched_message());
}

#endif  // USE_ATTACHMENT_BROKER

#if !USE_ATTACHMENT_BROKER

// We can determine message size from its header (and hence resize the buffer)
// only when attachment broker is not used, see IPC::Message::FindNext().

TEST(ChannelReaderTest, ResizeOverflowBuffer) {
  MockChannelReader reader;

  ExposedMessage::Header header = {};

  header.payload_size = 128 * 1024;
  EXPECT_LT(reader.input_overflow_buf_.capacity(), header.payload_size);
  EXPECT_TRUE(reader.TranslateInputData(
      reinterpret_cast<const char*>(&header), sizeof(header)));

  // Once message header is available we resize overflow buffer to
  // fit the entire message.
  EXPECT_GE(reader.input_overflow_buf_.capacity(), header.payload_size);
}

TEST(ChannelReaderTest, InvalidMessageSize) {
  MockChannelReader reader;

  ExposedMessage::Header header = {};

  size_t capacity_before = reader.input_overflow_buf_.capacity();

  // Message is slightly larger than maximum allowed size
  header.payload_size = Channel::kMaximumMessageSize + 1;
  EXPECT_FALSE(reader.TranslateInputData(
      reinterpret_cast<const char*>(&header), sizeof(header)));
  EXPECT_LE(reader.input_overflow_buf_.capacity(), capacity_before);

  // Payload size is negative, overflow is detected by Pickle::PeekNext()
  header.payload_size = static_cast<uint32_t>(-1);
  EXPECT_FALSE(reader.TranslateInputData(
      reinterpret_cast<const char*>(&header), sizeof(header)));
  EXPECT_LE(reader.input_overflow_buf_.capacity(), capacity_before);

  // Payload size is maximum int32_t value
  header.payload_size = std::numeric_limits<int32_t>::max();
  EXPECT_FALSE(reader.TranslateInputData(
      reinterpret_cast<const char*>(&header), sizeof(header)));
  EXPECT_LE(reader.input_overflow_buf_.capacity(), capacity_before);
}

#endif  // !USE_ATTACHMENT_BROKER

TEST(ChannelReaderTest, TrimBuffer) {
  // ChannelReader uses std::string as a buffer, and calls reserve()
  // to trim it to kMaximumReadBufferSize. However, an implementation
  // is free to actually reserve a larger amount.
  size_t trimmed_buffer_size;
  {
    std::string buf;
    buf.reserve(Channel::kMaximumReadBufferSize);
    trimmed_buffer_size = buf.capacity();
  }

  // Buffer is trimmed after message is processed.
  {
    MockChannelReader reader;

    Message message;
    message.WriteString(std::string(LargePayloadSize, 'X'));

    // Sanity check
    EXPECT_TRUE(message.size() > trimmed_buffer_size);

    // Initially buffer is small
    EXPECT_LE(reader.input_overflow_buf_.capacity(), trimmed_buffer_size);

    // Write and process large message
    reader.AppendMessageData(message);
    EXPECT_EQ(ChannelReader::DISPATCH_FINISHED,
              reader.ProcessIncomingMessages());

    // After processing large message buffer is trimmed
    EXPECT_EQ(reader.input_overflow_buf_.capacity(), trimmed_buffer_size);
  }

  // Buffer is trimmed only after entire message is processed.
  {
    MockChannelReader reader;

    ExposedMessage message;
    message.WriteString(std::string(LargePayloadSize, 'X'));

    // Write and process message header
    reader.AppendData(message.header(), sizeof(ExposedMessage::Header));
    EXPECT_EQ(ChannelReader::DISPATCH_FINISHED,
              reader.ProcessIncomingMessages());

#if MESSAGE_FINDNEXT_PARTIAL
    // We determined message size for the message from its header, so
    // we resized the buffer to fit.
    EXPECT_GE(reader.input_overflow_buf_.capacity(), message.size());
#else
    // We couldn't determine message size, so we didn't resize the buffer.
#endif

    // Write and process payload
    reader.AppendData(message.payload(), message.payload_size());
    EXPECT_EQ(ChannelReader::DISPATCH_FINISHED,
              reader.ProcessIncomingMessages());

    // But once we process the message, we trim the buffer
    EXPECT_EQ(reader.input_overflow_buf_.capacity(), trimmed_buffer_size);
  }

  // Buffer is not trimmed if the next message is also large.
  {
    MockChannelReader reader;

    // Write large message
    Message message1;
    message1.WriteString(std::string(LargePayloadSize * 2, 'X'));
    reader.AppendMessageData(message1);

    // Write header for the next large message
    ExposedMessage message2;
    message2.WriteString(std::string(LargePayloadSize, 'Y'));
    reader.AppendData(message2.header(), sizeof(ExposedMessage::Header));

    // Process messages
    EXPECT_EQ(ChannelReader::DISPATCH_FINISHED,
              reader.ProcessIncomingMessages());

#if MESSAGE_FINDNEXT_PARTIAL
    // We determined message size for the second (partial) message, so
    // we resized the buffer to fit.
    EXPECT_GE(reader.input_overflow_buf_.capacity(), message1.size());
#else
    // We couldn't determine message size for the second (partial) message,
    // so we trimmed the buffer.
    EXPECT_EQ(reader.input_overflow_buf_.capacity(), trimmed_buffer_size);
#endif
  }

  // Buffer resized appropriately if next message is larger than the first.
  // (Similar to the test above except for the order of messages.)
  {
    MockChannelReader reader;

    // Write large message
    Message message1;
    message1.WriteString(std::string(LargePayloadSize, 'Y'));
    reader.AppendMessageData(message1);

    // Write header for the next even larger message
    ExposedMessage message2;
    message2.WriteString(std::string(LargePayloadSize * 2, 'X'));
    reader.AppendData(message2.header(), sizeof(ExposedMessage::Header));

    // Process messages
    EXPECT_EQ(ChannelReader::DISPATCH_FINISHED,
              reader.ProcessIncomingMessages());

#if MESSAGE_FINDNEXT_PARTIAL
    // We determined message size for the second (partial) message, and
    // resized the buffer to fit it.
    EXPECT_GE(reader.input_overflow_buf_.capacity(), message2.size());
#else
    // We couldn't determine message size for the second (partial) message,
    // so we trimmed the buffer.
    EXPECT_EQ(reader.input_overflow_buf_.capacity(), trimmed_buffer_size);
#endif
  }

  // Buffer is not trimmed if we've just resized it to accommodate large
  // incoming message.
  {
    MockChannelReader reader;

    // Write small message
    Message message1;
    message1.WriteString(std::string(11, 'X'));
    reader.AppendMessageData(message1);

    // Write header for the next large message
    ExposedMessage message2;
    message2.WriteString(std::string(LargePayloadSize, 'Y'));
    reader.AppendData(message2.header(), sizeof(ExposedMessage::Header));

    EXPECT_EQ(ChannelReader::DISPATCH_FINISHED,
              reader.ProcessIncomingMessages());

#if MESSAGE_FINDNEXT_PARTIAL
    // We determined message size for the second (partial) message, so
    // we resized the buffer to fit.
    EXPECT_GE(reader.input_overflow_buf_.capacity(), message2.size());
#else
    // We couldn't determine size for the second (partial) message, and
    // first message was small, so we did nothing.
#endif
  }
}

}  // namespace internal
}  // namespace IPC