// 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 "ipc/attachment_broker_privileged_mac.h"

#include <mach/mach.h>
#include <mach/mach_vm.h>

#include <map>

#include "base/command_line.h"
#include "base/mac/mac_util.h"
#include "base/mac/mach_logging.h"
#include "base/mac/scoped_mach_port.h"
#include "base/memory/shared_memory.h"
#include "base/process/port_provider_mac.h"
#include "base/process/process_handle.h"
#include "base/sys_info.h"
#include "base/test/multiprocess_test.h"
#include "base/test/test_timeouts.h"
#include "ipc/test_util_mac.h"
#include "testing/multiprocess_func_list.h"

namespace IPC {

namespace {

static const std::string g_service_switch_name = "service_name";

// Sends a uint32_t to a mach port.
void SendUInt32(mach_port_t port, uint32_t message) {
  int message_size = sizeof(uint32_t);
  int total_size = message_size + sizeof(mach_msg_header_t);
  void* buffer = malloc(total_size);
  mach_msg_header_t* header = (mach_msg_header_t*)buffer;
  header->msgh_remote_port = port;
  header->msgh_local_port = MACH_PORT_NULL;
  header->msgh_bits = MACH_MSGH_BITS(MACH_MSG_TYPE_COPY_SEND, 0);
  header->msgh_reserved = 0;
  header->msgh_id = 0;
  header->msgh_size = total_size;
  memcpy(static_cast<char*>(buffer) + sizeof(mach_msg_header_t), &message,
         message_size);

  kern_return_t kr;
  kr = mach_msg(static_cast<mach_msg_header_t*>(buffer), MACH_SEND_MSG,
                total_size, 0, MACH_PORT_NULL, MACH_MSG_TIMEOUT_NONE,
                MACH_PORT_NULL);
  MACH_CHECK(kr == KERN_SUCCESS, kr) << "SendUInt32";
  free(buffer);
}

// Receives a uint32_t from a mach port.
uint32_t ReceiveUInt32(mach_port_t listening_port) {
  int message_size = sizeof(uint32_t);
  int total_size =
      message_size + sizeof(mach_msg_header_t) + sizeof(mach_msg_trailer_t);
  int options = MACH_RCV_MSG;
  void* buffer = malloc(total_size);

  int kr =
      mach_msg(static_cast<mach_msg_header_t*>(buffer), options, 0, total_size,
               listening_port, MACH_MSG_TIMEOUT_NONE, MACH_PORT_NULL);
  MACH_CHECK(kr == KERN_SUCCESS, kr) << "ReceiveUInt32";

  uint32_t response;
  memcpy(&response, static_cast<char*>(buffer) + sizeof(mach_msg_header_t),
         message_size);

  free(buffer);
  return response;
}

// Sets up the mach communication ports with the server. Returns a port to which
// the server will send mach objects.
// |original_name_count| is an output variable that describes the number of
// active names in this task before the task port is shared with the server.
base::mac::ScopedMachReceiveRight CommonChildProcessSetUp(
    mach_msg_type_number_t* original_name_count) {
  base::CommandLine cmd_line = *base::CommandLine::ForCurrentProcess();
  std::string service_name =
      cmd_line.GetSwitchValueASCII(g_service_switch_name);
  base::mac::ScopedMachSendRight server_port(
      LookupServer(service_name.c_str()));
  base::mac::ScopedMachReceiveRight client_port(MakeReceivingPort());

  // |server_port| is a newly allocated right which will be deallocated once
  // this method returns.
  *original_name_count = GetActiveNameCount() - 1;

  // Send the port that this process is listening on to the server.
  SendMachPort(server_port.get(), client_port.get(), MACH_MSG_TYPE_MAKE_SEND);

  // Send the task port for this process.
  SendMachPort(server_port.get(), mach_task_self(), MACH_MSG_TYPE_COPY_SEND);
  return client_port;
}

// Creates a new shared memory region populated with 'a'.
scoped_ptr<base::SharedMemory> CreateAndPopulateSharedMemoryHandle(
    size_t size) {
  base::SharedMemoryHandle shm(size);
  scoped_ptr<base::SharedMemory> shared_memory(
      new base::SharedMemory(shm, false));
  shared_memory->Map(size);
  memset(shared_memory->memory(), 'a', size);
  return shared_memory;
}

// Create a shared memory region from a memory object. The returned object takes
// ownership of |memory_object|.
scoped_ptr<base::SharedMemory> MapMemoryObject(mach_port_t memory_object,
                                               size_t size) {
  base::SharedMemoryHandle shm(memory_object, size, base::GetCurrentProcId());
  scoped_ptr<base::SharedMemory> shared_memory(
      new base::SharedMemory(shm, false));
  shared_memory->Map(size);
  return shared_memory;
}

class MockPortProvider : public base::PortProvider {
 public:
  MockPortProvider() {}
  ~MockPortProvider() override {}
  mach_port_t TaskForPid(base::ProcessHandle process) const override {
    return MACH_PORT_NULL;
  }
};

}  // namespace

class AttachmentBrokerPrivilegedMacMultiProcessTest
    : public base::MultiProcessTest {
 public:
  AttachmentBrokerPrivilegedMacMultiProcessTest() {}

  base::CommandLine MakeCmdLine(const std::string& procname) override {
    base::CommandLine command_line = MultiProcessTest::MakeCmdLine(procname);
    // Pass the service name to the child process.
    command_line.AppendSwitchASCII(g_service_switch_name, service_name_);
    return command_line;
  }

  void SetUpChild(const std::string& name) {
    // Make a random service name so that this test doesn't conflict with other
    // similar tests.
    service_name_ = CreateRandomServiceName();
    server_port_.reset(BecomeMachServer(service_name_.c_str()).release());
    child_process_ = SpawnChild(name);
    client_port_.reset(ReceiveMachPort(server_port_.get()).release());
    client_task_port_.reset(ReceiveMachPort(server_port_.get()).release());
  }

  static const int s_memory_size = 99999;

 protected:
  std::string service_name_;

  // A port on which the main process listens for mach messages from the child
  // process.
  base::mac::ScopedMachReceiveRight server_port_;

  // A port on which the child process listens for mach messages from the main
  // process.
  base::mac::ScopedMachSendRight client_port_;

  // Child process's task port.
  base::mac::ScopedMachSendRight client_task_port_;

  // Dummy port provider.
  MockPortProvider port_provider_;

  base::Process child_process_;
  DISALLOW_COPY_AND_ASSIGN(AttachmentBrokerPrivilegedMacMultiProcessTest);
};

// The attachment broker inserts a right for a memory object into the
// destination task.
TEST_F(AttachmentBrokerPrivilegedMacMultiProcessTest, InsertRight) {
  // Mach-based SharedMemory isn't support on OSX 10.6.
  if (base::mac::IsOSSnowLeopard())
    return;

  SetUpChild("InsertRightClient");
  mach_msg_type_number_t original_name_count = GetActiveNameCount();
  IPC::AttachmentBrokerPrivilegedMac broker(&port_provider_);

  // Create some shared memory.
  scoped_ptr<base::SharedMemory> shared_memory =
      CreateAndPopulateSharedMemoryHandle(s_memory_size);
  ASSERT_TRUE(shared_memory->handle().IsValid());

  // Insert the memory object into the destination task, via an intermediate
  // port.
  IncrementMachRefCount(shared_memory->handle().GetMemoryObject(),
                        MACH_PORT_RIGHT_SEND);
  mach_port_name_t inserted_memory_object = broker.CreateIntermediateMachPort(
      client_task_port_.get(), base::mac::ScopedMachSendRight(
                                   shared_memory->handle().GetMemoryObject()));
  EXPECT_NE(inserted_memory_object,
            static_cast<mach_port_name_t>(MACH_PORT_NULL));
  SendUInt32(client_port_.get(), inserted_memory_object);

  // Check that no names have been leaked.
  shared_memory.reset();
  EXPECT_EQ(original_name_count, GetActiveNameCount());

  int rv = -1;
  ASSERT_TRUE(child_process_.WaitForExitWithTimeout(
      TestTimeouts::action_timeout(), &rv));
  EXPECT_EQ(0, rv);
}

MULTIPROCESS_TEST_MAIN(InsertRightClient) {
  mach_msg_type_number_t original_name_count = 0;
  base::mac::ScopedMachReceiveRight client_port(
      CommonChildProcessSetUp(&original_name_count).release());
  base::mac::ScopedMachReceiveRight inserted_port(
      ReceiveUInt32(client_port.get()));
  base::mac::ScopedMachSendRight memory_object(
      ReceiveMachPort(inserted_port.get()));
  inserted_port.reset();

  // The server should have inserted a right into this process.
  EXPECT_EQ(original_name_count + 1, GetActiveNameCount());

  // Map the memory object and check its contents.
  scoped_ptr<base::SharedMemory> shared_memory(MapMemoryObject(
      memory_object.release(),
      AttachmentBrokerPrivilegedMacMultiProcessTest::s_memory_size));
  const char* start = static_cast<const char*>(shared_memory->memory());
  for (int i = 0;
       i < AttachmentBrokerPrivilegedMacMultiProcessTest::s_memory_size; ++i) {
    DCHECK_EQ(start[i], 'a');
  }

  // Check that no names have been leaked.
  shared_memory.reset();
  EXPECT_EQ(original_name_count, GetActiveNameCount());

  return 0;
}

// The attachment broker inserts the right for a memory object into the
// destination task twice.
TEST_F(AttachmentBrokerPrivilegedMacMultiProcessTest, InsertSameRightTwice) {
  // Mach-based SharedMemory isn't support on OSX 10.6.
  if (base::mac::IsOSSnowLeopard())
    return;

  SetUpChild("InsertSameRightTwiceClient");
  mach_msg_type_number_t original_name_count = GetActiveNameCount();
  IPC::AttachmentBrokerPrivilegedMac broker(&port_provider_);

  // Create some shared memory.
  scoped_ptr<base::SharedMemory> shared_memory =
      CreateAndPopulateSharedMemoryHandle(s_memory_size);
  ASSERT_TRUE(shared_memory->handle().IsValid());

  // Insert the memory object into the destination task, via an intermediate
  // port, twice.
  for (int i = 0; i < 2; ++i) {
    IncrementMachRefCount(shared_memory->handle().GetMemoryObject(),
                          MACH_PORT_RIGHT_SEND);
    mach_port_name_t inserted_memory_object = broker.CreateIntermediateMachPort(
        client_task_port_.get(),
        base::mac::ScopedMachSendRight(
            shared_memory->handle().GetMemoryObject()));
    EXPECT_NE(inserted_memory_object,
              static_cast<mach_port_name_t>(MACH_PORT_NULL));
    SendUInt32(client_port_.get(), inserted_memory_object);
  }

  // Check that no names have been leaked.
  shared_memory.reset();
  EXPECT_EQ(original_name_count, GetActiveNameCount());

  int rv = -1;
  ASSERT_TRUE(child_process_.WaitForExitWithTimeout(
      TestTimeouts::action_timeout(), &rv));
  EXPECT_EQ(0, rv);
}

MULTIPROCESS_TEST_MAIN(InsertSameRightTwiceClient) {
  mach_msg_type_number_t original_name_count = 0;
  base::mac::ScopedMachReceiveRight client_port(
      CommonChildProcessSetUp(&original_name_count).release());

  // Receive two memory objects.
  base::mac::ScopedMachReceiveRight inserted_port(
      ReceiveUInt32(client_port.get()));
  base::mac::ScopedMachReceiveRight inserted_port2(
      ReceiveUInt32(client_port.get()));
  base::mac::ScopedMachSendRight memory_object(
      ReceiveMachPort(inserted_port.get()));
  base::mac::ScopedMachSendRight memory_object2(
      ReceiveMachPort(inserted_port2.get()));
  inserted_port.reset();
  inserted_port2.reset();

  // Both rights are for the same Mach port, so only one new name should appear.
  EXPECT_EQ(original_name_count + 1, GetActiveNameCount());

  // Map both memory objects and check their contents.
  scoped_ptr<base::SharedMemory> shared_memory(MapMemoryObject(
      memory_object.release(),
      AttachmentBrokerPrivilegedMacMultiProcessTest::s_memory_size));
  char* start = static_cast<char*>(shared_memory->memory());
  for (int i = 0;
       i < AttachmentBrokerPrivilegedMacMultiProcessTest::s_memory_size; ++i) {
    DCHECK_EQ(start[i], 'a');
  }

  scoped_ptr<base::SharedMemory> shared_memory2(MapMemoryObject(
      memory_object2.release(),
      AttachmentBrokerPrivilegedMacMultiProcessTest::s_memory_size));
  char* start2 = static_cast<char*>(shared_memory2->memory());
  for (int i = 0;
       i < AttachmentBrokerPrivilegedMacMultiProcessTest::s_memory_size; ++i) {
    DCHECK_EQ(start2[i], 'a');
  }

  // Check that the contents of both regions are shared.
  start[0] = 'b';
  DCHECK_EQ(start2[0], 'b');

  // After releasing one shared memory region, the name count shouldn't change,
  // since another reference exists.
  shared_memory.reset();
  EXPECT_EQ(original_name_count + 1, GetActiveNameCount());

  // After releasing the second shared memory region, the name count should be
  // as if no names were ever inserted
  shared_memory2.reset();
  EXPECT_EQ(original_name_count, GetActiveNameCount());

  return 0;
}

// The attachment broker inserts the rights for two memory objects into the
// destination task.
TEST_F(AttachmentBrokerPrivilegedMacMultiProcessTest, InsertTwoRights) {
  // Mach-based SharedMemory isn't support on OSX 10.6.
  if (base::mac::IsOSSnowLeopard())
    return;

  SetUpChild("InsertTwoRightsClient");
  mach_msg_type_number_t original_name_count = GetActiveNameCount();
  IPC::AttachmentBrokerPrivilegedMac broker(&port_provider_);

  for (int i = 0; i < 2; ++i) {
    // Create some shared memory.
    scoped_ptr<base::SharedMemory> shared_memory =
        CreateAndPopulateSharedMemoryHandle(s_memory_size);
    ASSERT_TRUE(shared_memory->handle().IsValid());

    // Insert the memory object into the destination task, via an intermediate
    // port.
    IncrementMachRefCount(shared_memory->handle().GetMemoryObject(),
                          MACH_PORT_RIGHT_SEND);
    mach_port_name_t inserted_memory_object = broker.CreateIntermediateMachPort(
        client_task_port_.get(),
        base::mac::ScopedMachSendRight(
            shared_memory->handle().GetMemoryObject()));
    EXPECT_NE(inserted_memory_object,
              static_cast<mach_port_name_t>(MACH_PORT_NULL));
    SendUInt32(client_port_.get(), inserted_memory_object);
  }

  // Check that no names have been leaked.
  EXPECT_EQ(original_name_count, GetActiveNameCount());

  int rv = -1;
  ASSERT_TRUE(child_process_.WaitForExitWithTimeout(
      TestTimeouts::action_timeout(), &rv));
  EXPECT_EQ(0, rv);
}

MULTIPROCESS_TEST_MAIN(InsertTwoRightsClient) {
  mach_msg_type_number_t original_name_count = 0;
  base::mac::ScopedMachReceiveRight client_port(
      CommonChildProcessSetUp(&original_name_count).release());

  // Receive two memory objects.
  base::mac::ScopedMachReceiveRight inserted_port(
      ReceiveUInt32(client_port.get()));
  base::mac::ScopedMachReceiveRight inserted_port2(
      ReceiveUInt32(client_port.get()));
  base::mac::ScopedMachSendRight memory_object(
      ReceiveMachPort(inserted_port.get()));
  base::mac::ScopedMachSendRight memory_object2(
      ReceiveMachPort(inserted_port2.get()));
  inserted_port.reset();
  inserted_port2.reset();

  // There should be two new names to reflect the two new shared memory regions.
  EXPECT_EQ(original_name_count + 2, GetActiveNameCount());

  // Map both memory objects and check their contents.
  scoped_ptr<base::SharedMemory> shared_memory(MapMemoryObject(
      memory_object.release(),
      AttachmentBrokerPrivilegedMacMultiProcessTest::s_memory_size));
  char* start = static_cast<char*>(shared_memory->memory());
  for (int i = 0;
       i < AttachmentBrokerPrivilegedMacMultiProcessTest::s_memory_size; ++i) {
    DCHECK_EQ(start[i], 'a');
  }

  scoped_ptr<base::SharedMemory> shared_memory2(MapMemoryObject(
      memory_object2.release(),
      AttachmentBrokerPrivilegedMacMultiProcessTest::s_memory_size));
  char* start2 = static_cast<char*>(shared_memory2->memory());
  for (int i = 0;
       i < AttachmentBrokerPrivilegedMacMultiProcessTest::s_memory_size; ++i) {
    DCHECK_EQ(start2[i], 'a');
  }

  // Check that the contents of both regions are not shared.
  start[0] = 'b';
  DCHECK_EQ(start2[0], 'a');

  // After releasing one shared memory region, the name count should decrement.
  shared_memory.reset();
  EXPECT_EQ(original_name_count + 1, GetActiveNameCount());
  shared_memory2.reset();
  EXPECT_EQ(original_name_count, GetActiveNameCount());

  return 0;
}

}  //  namespace IPC