// Copyright 2019 The Fuchsia 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 <fcntl.h>
#include <fuchsia/netemul/guest/cpp/fidl.h>
#include <lib/fdio/directory.h>
#include <lib/fdio/fd.h>
#include <lib/sys/cpp/file_descriptor.h>
#include <lib/syslog/cpp/macros.h>
#include <unistd.h>

#include <fbl/unique_fd.h>
#include <gtest/gtest.h>

#include "src/lib/testing/predicates/status.h"
#include "src/virtualization/lib/guest_interaction/common.h"
#include "src/virtualization/lib/guest_interaction/test/integration_test_lib.h"

class GuestInteractionTestWithDiscovery : public GuestInteractionTest {
 protected:
  GuestInteractionTestWithDiscovery() {
    services().AddServiceWithLaunchInfo(
        {
            .url = kGuestDiscoveryUrl,
            .out = sys::CloneFileDescriptor(STDOUT_FILENO),
            .err = sys::CloneFileDescriptor(STDERR_FILENO),

        },
        fuchsia::netemul::guest::GuestDiscovery::Name_);
  }

  fuchsia::netemul::guest::GuestInteractionPtr guest_interaction() {
    if (!guest_discovery_.is_bound()) {
      env().ConnectToService(guest_discovery_.NewRequest());
    }
    fuchsia::netemul::guest::GuestInteractionPtr guest_interaction;
    guest_discovery_->GetGuest(std::nullopt, kGuestLabel, guest_interaction.NewRequest());
    return guest_interaction;
  }

 private:
  fuchsia::netemul::guest::GuestDiscoveryPtr guest_discovery_;
};

// Call into the guest interaction fidl service to push a test bash script, run
// the bash script, and pull back the results generated by the script.  The
// script emits output to stdout and stderr and accepts input from stdin which
// is used to generate the output file.
TEST_F(GuestInteractionTestWithDiscovery, FidlExecScriptTest) {
  fuchsia::netemul::guest::GuestInteractionPtr gis = guest_interaction();
  std::optional<zx_status_t> gis_status;
  gis.set_error_handler([&gis_status](zx_status_t status) { gis_status = status; });

  // Push the bash script to the guest
  FX_LOGS(INFO) << "Sending script to guest";
  fidl::InterfaceHandle<fuchsia::io::File> put_file;
  ASSERT_OK(fdio_open(kTestScriptSource,
                      static_cast<uint32_t>(fuchsia::io::OpenFlags::RIGHT_READABLE),
                      put_file.NewRequest().TakeChannel().release()));

  std::optional<zx_status_t> put_status;
  gis->PutFile(std::move(put_file), kGuestScriptDestination,
               [&put_status](zx_status_t status) { put_status = status; });
  RunLoopUntil(
      [&gis_status, &put_status]() { return gis_status.has_value() || put_status.has_value(); });
  ASSERT_FALSE(gis_status.has_value()) << zx_status_get_string(gis_status.value());
  ASSERT_TRUE(put_status.has_value());
  ASSERT_OK(put_status.value());

  // Run the bash script in the guest.  The script will write to stdout and
  // stderr.  The script will also block waiting to receive input from stdin.
  zx::socket stdin_writer, stdin_reader;
  ASSERT_OK(zx::socket::create(0, &stdin_writer, &stdin_reader));

  zx::socket stdout_writer, stdout_reader;
  ASSERT_OK(zx::socket::create(0, &stdout_writer, &stdout_reader));

  zx::socket stderr_writer, stderr_reader;
  ASSERT_OK(zx::socket::create(0, &stderr_writer, &stderr_reader));

  // Run the bash script on the guest.
  FX_LOGS(INFO) << "Running script in guest";
  std::string command = std::string("/bin/sh ");
  command.append(kGuestScriptDestination);
  std::vector<fuchsia::netemul::guest::EnvironmentVariable> env_vars = {
      {"STDOUT_STRING", kTestStdout}, {"STDERR_STRING", kTestStderr}};
  std::string std_out;
  std::string std_err;
  std::optional<zx_status_t> exec_started;
  std::optional<zx_status_t> exec_terminated;

  fuchsia::netemul::guest::CommandListenerPtr listener;
  listener.events().OnStarted = [&](zx_status_t status) {
    // Once the subprocess has started, write to stdin.
    std::string to_write = std::string(::kTestScriptInput);
    uint32_t bytes_written = 0;

    while (bytes_written < to_write.size()) {
      size_t curr_bytes_written;
      zx_status_t write_status =
          stdin_writer.write(0, &(to_write.c_str()[bytes_written]), to_write.size() - bytes_written,
                             &curr_bytes_written);
      if (write_status != ZX_OK) {
        return;
      }

      bytes_written += curr_bytes_written;
    }
    stdin_writer.reset();
    exec_started = status;
  };
  listener.events().OnTerminated = [&](zx_status_t status, int32_t exit_code) {
    // When the process terminates, read from stdout and stderr.
    exec_terminated = status;

    size_t bytes_read = 0;
    char read_buf[100];
    while (true) {
      zx_status_t read_status = stdout_reader.read(0, read_buf, 100, &bytes_read);
      if (read_status != ZX_OK) {
        break;
      }
      std_out += std::string(read_buf, bytes_read);
    }
    while (true) {
      zx_status_t read_status = stderr_reader.read(0, read_buf, 100, &bytes_read);
      if (read_status != ZX_OK) {
        break;
      }
      std_err += std::string(read_buf, bytes_read);
    }
  };

  gis->ExecuteCommand(command, env_vars, std::move(stdin_reader), std::move(stdout_writer),
                      std::move(stderr_writer), listener.NewRequest());

  // Wait for the subprocess to start.
  RunLoopUntil([&gis_status, &exec_started]() {
    return gis_status.has_value() || exec_started.has_value();
  });
  ASSERT_FALSE(gis_status.has_value()) << zx_status_get_string(gis_status.value());
  ASSERT_TRUE(exec_started.has_value());
  ASSERT_OK(exec_started.value());

  // Wait for the bash script to finish.
  RunLoopUntil([&gis_status, &exec_terminated]() {
    return gis_status.has_value() || exec_terminated.has_value();
  });
  ASSERT_FALSE(gis_status.has_value()) << zx_status_get_string(gis_status.value());
  ASSERT_TRUE(exec_terminated.has_value());
  ASSERT_OK(exec_terminated.value());

  // Validate the stdout and stderr.
  ASSERT_EQ(std_out, (std::string(kTestStdout) + std::string("\n")));
  ASSERT_EQ(std_err, (std::string(kTestStderr) + std::string("\n")));

  // The bash script will create a file with contents that were written to
  // stdin.  Pull this file back and inspect its contents.
  fidl::InterfaceHandle<fuchsia::io::File> get_file;
  ASSERT_OK(fdio_open(
      kHostOuputCopyLocation,
      static_cast<uint32_t>(fuchsia::io::OpenFlags::RIGHT_WRITABLE |
                            fuchsia::io::OpenFlags::CREATE | fuchsia::io::OpenFlags::TRUNCATE),
      get_file.NewRequest().TakeChannel().release()));

  std::optional<zx_status_t> get_status;
  gis->GetFile(kGuestFileOutputLocation, std::move(get_file),
               [&get_status](zx_status_t status) { get_status = status; });
  RunLoopUntil(
      [&gis_status, &get_status]() { return gis_status.has_value() || get_status.has_value(); });
  ASSERT_FALSE(gis_status.has_value()) << zx_status_get_string(gis_status.value());
  ASSERT_TRUE(get_status.has_value());
  ASSERT_OK(get_status.value());

  // Verify the contents that were communicated through stdin.
  std::string output_string;
  {
    fbl::unique_fd fd;
    ASSERT_TRUE(fd = fbl::unique_fd(open(kHostOuputCopyLocation, O_RDONLY))) << strerror(errno);

    char output_buf[100];
    while (true) {
      ssize_t read_size = read(fd.get(), output_buf, sizeof(output_buf));
      ASSERT_GE(read_size, 0) << strerror(errno);
      if (read_size == 0) {
        break;
      }
      output_string.append(std::string(output_buf, read_size));
    }
  }

  ASSERT_EQ(output_string, std::string(kTestScriptInput));
}

TEST_F(GuestInteractionTestWithDiscovery, FidlPutGetTest) {
  fuchsia::netemul::guest::GuestInteractionPtr gis = guest_interaction();

  // Write a file of gibberish that the test can send over to the guest.
  constexpr char test_file[] = "/tmp/test_file.txt";
  constexpr char guest_destination[] = "/root/new/directory/test_file.txt";
  constexpr char host_verification_file[] = "/tmp/verification_file.txt";

  std::string file_contents;
  for (int i = 0; i < 2 * CHUNK_SIZE; i++) {
    file_contents.push_back(static_cast<char>(i % (('z' - 'A') + 'A')));
  }
  {
    fbl::unique_fd fd;
    ASSERT_TRUE(fd = fbl::unique_fd(open(test_file, O_WRONLY | O_TRUNC | O_CREAT)))
        << strerror(errno);
    uint32_t bytes_written = 0;
    while (bytes_written < file_contents.size()) {
      ssize_t write_size = write(fd.get(), file_contents.c_str() + bytes_written,
                                 file_contents.size() - bytes_written);
      ASSERT_GT(write_size, 0);
      bytes_written += write_size;
    }
  }

  // Push the file to the guest
  fidl::InterfaceHandle<fuchsia::io::File> put_file;
  ASSERT_OK(fdio_open(test_file, static_cast<uint32_t>(fuchsia::io::OpenFlags::RIGHT_READABLE),
                      put_file.NewRequest().TakeChannel().release()));

  std::optional<zx_status_t> status;
  gis.set_error_handler([&status](zx_status_t error_status) { status = error_status; });
  gis->PutFile(std::move(put_file), guest_destination,
               [&status](zx_status_t put_result) { status = put_result; });
  RunLoopUntil([&status]() { return status.has_value(); });
  ASSERT_TRUE(status.has_value());
  ASSERT_OK(status.value());

  // Pull back the guest's copy of the file and ensure the contents match those
  // from the file generated above.
  fidl::InterfaceHandle<fuchsia::io::File> get_file;
  ASSERT_OK(fdio_open(
      host_verification_file,
      static_cast<uint32_t>(fuchsia::io::OpenFlags::RIGHT_WRITABLE |
                            fuchsia::io::OpenFlags::CREATE | fuchsia::io::OpenFlags::TRUNCATE),
      get_file.NewRequest().TakeChannel().release()));

  status.reset();
  gis->GetFile(guest_destination, std::move(get_file),
               [&status](zx_status_t get_result) { status = get_result; });
  RunLoopUntil([&status]() { return status.has_value(); });
  ASSERT_TRUE(status.has_value());
  ASSERT_OK(status.value());

  // Verify the contents that were communicated through stdin.
  std::string output_string;
  {
    fbl::unique_fd fd;
    ASSERT_TRUE(fd = fbl::unique_fd(open(host_verification_file, O_RDONLY))) << strerror(errno);

    char output_buf[100];
    while (true) {
      ssize_t read_size = read(fd.get(), output_buf, sizeof(output_buf));
      ASSERT_GE(read_size, 0) << strerror(errno);
      if (read_size == 0) {
        break;
      }
      output_string.append(std::string(output_buf, read_size));
    }
  }

  ASSERT_EQ(output_string, file_contents);
}