/*
 * Copyright (C) 2018 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "common/libs/utils/subprocess.h"

#include <errno.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>

#include <map>
#include <set>

#include <glog/logging.h>

#include "common/libs/fs/shared_buf.h"

namespace {

// If a redirected-to file descriptor was already closed, it's possible that
// some inherited file descriptor duped to this file descriptor and the redirect
// would override that. This function makes sure that doesn't happen.
bool validate_redirects(
    const std::map<cvd::Subprocess::StdIOChannel, int>& redirects,
    const std::map<cvd::SharedFD, int>& inherited_fds) {
  // Add the redirected IO channels to a set as integers. This allows converting
  // the enum values into integers instead of the other way around.
  std::set<int> int_redirects;
  for (const auto& entry: redirects) {
    int_redirects.insert(static_cast<int>(entry.first));
  }
  for (const auto& entry: inherited_fds) {
    auto dupped_fd = entry.second;
    if (int_redirects.count(dupped_fd)) {
      LOG(ERROR) << "Requested redirect of fd(" << dupped_fd
                 << ") conflicts with inherited FD.";
      return false;
    }
  }
  return true;
}

void do_redirects(const std::map<cvd::Subprocess::StdIOChannel, int>& redirects) {
  for (const auto& entry: redirects) {
    auto std_channel = static_cast<int>(entry.first);
    auto fd = entry.second;
    TEMP_FAILURE_RETRY(dup2(fd, std_channel));
  }
}

cvd::Subprocess subprocess_impl(
    const char* const* command,
    const char* const* envp,
    const std::map<cvd::Subprocess::StdIOChannel, int>& redirects,
    const std::map<cvd::SharedFD, int>& inherited_fds,
    bool with_control_socket) {
  // The parent socket will get closed on the child on the call to exec, the
  // child socket will be closed on the parent when this function returns and no
  // references to the fd are left
  cvd::SharedFD parent_socket, child_socket;
  if (with_control_socket) {
    if (!cvd::SharedFD::SocketPair(AF_LOCAL, SOCK_STREAM, 0, &parent_socket,
                                   &child_socket)) {
      LOG(ERROR) << "Unable to create control socket pair: " << strerror(errno);
      return cvd::Subprocess(-1, {});
    }
    // Remove FD_CLOEXEC from the child socket, ensure the parent has it
    child_socket->Fcntl(F_SETFD, 0);
    parent_socket->Fcntl(F_SETFD, FD_CLOEXEC);
  }

  if (!validate_redirects(redirects, inherited_fds)) {
    return cvd::Subprocess(-1, {});
  }

  pid_t pid = fork();
  if (!pid) {
    do_redirects(redirects);
    int rval;
    // If envp is NULL, the current process's environment is used as the
    // environment of the child process. To force an empty emvironment for
    // the child process pass the address of a pointer to NULL
    if (envp == NULL) {
      rval = execv(command[0], const_cast<char* const*>(command));
    } else {
      rval = execve(command[0],
                    const_cast<char* const*>(command),
                    const_cast<char* const*>(envp));
    }
    // No need for an if: if exec worked it wouldn't have returned
    LOG(ERROR) << "exec of " << command[0] << " failed (" << strerror(errno)
               << ")";
    exit(rval);
  }
  if (pid == -1) {
    LOG(ERROR) << "fork failed (" << strerror(errno) << ")";
  }
  LOG(INFO) << "Started (pid: " << pid << "): " << command[0];
  int i = 1;
  while (command[i]) {
    LOG(INFO) << command[i++];
  }
  return cvd::Subprocess(pid, parent_socket);
}

std::vector<const char*> ToCharPointers(
    const std::vector<std::string>& vect) {
  std::vector<const char*> ret = {};
  for (const auto& str : vect) {
    ret.push_back(str.c_str());
  }
  ret.push_back(NULL);
  return ret;
}
}  // namespace
namespace cvd {

Subprocess::Subprocess(Subprocess&& subprocess)
    : pid_(subprocess.pid_),
      started_(subprocess.started_),
      control_socket_(subprocess.control_socket_) {
  // Make sure the moved object no longer controls this subprocess
  subprocess.pid_ = -1;
  subprocess.started_ = false;
  subprocess.control_socket_ = SharedFD();
}

Subprocess& Subprocess::operator=(Subprocess&& other) {
  pid_ = other.pid_;
  started_ = other.started_;
  control_socket_ = other.control_socket_;

  other.pid_ = -1;
  other.started_ = false;
  other.control_socket_ = SharedFD();
  return *this;
}

int Subprocess::Wait() {
  if (pid_ < 0) {
    LOG(ERROR)
        << "Attempt to wait on invalid pid(has it been waited on already?): "
        << pid_;
    return -1;
  }
  int wstatus = 0;
  auto pid = pid_; // Wait will set pid_ to -1 after waiting
  auto wait_ret = Wait(&wstatus, 0);
  if (wait_ret < 0) {
    LOG(ERROR) << "Error on call to waitpid: " << strerror(errno);
    return wait_ret;
  }
  int retval = 0;
  if (WIFEXITED(wstatus)) {
    retval = WEXITSTATUS(wstatus);
    if (retval) {
      LOG(ERROR) << "Subprocess " << pid
                 << " exited with error code: " << retval;
    }
  } else if (WIFSIGNALED(wstatus)) {
    LOG(ERROR) << "Subprocess " << pid
               << " was interrupted by a signal: " << WTERMSIG(wstatus);
    retval = -1;
  }
  return retval;
}
pid_t Subprocess::Wait(int* wstatus, int options) {
  if (pid_ < 0) {
    LOG(ERROR)
        << "Attempt to wait on invalid pid(has it been waited on already?): "
        << pid_;
    return -1;
  }
  auto retval = waitpid(pid_, wstatus, options);
  // We don't want to wait twice for the same process
  pid_ = -1;
  return retval;
}

Command::ParameterBuilder::~ParameterBuilder() {
  Build();
}
void Command::ParameterBuilder::Build()  {
  auto param = stream_.str();
  stream_ = std::stringstream();
  if (param.size()) {
    cmd_->AddParameter(param);
  }
}

Command::~Command() {
  // Close all inherited file descriptors
  for(const auto& entry: inherited_fds_) {
    close(entry.second);
  }
  // Close all redirected file descriptors
  for (const auto& entry: redirects_) {
    close(entry.second);
  }
}

bool Command::BuildParameter(std::stringstream* stream, SharedFD shared_fd) {
  int fd;
  if (inherited_fds_.count(shared_fd)) {
    fd = inherited_fds_[shared_fd];
  } else {
    fd = shared_fd->UNMANAGED_Dup();
    if (fd < 0) {
      return false;
    }
    inherited_fds_[shared_fd] = fd;
  }
  *stream << fd;
  return true;
}

bool Command::RedirectStdIO(cvd::Subprocess::StdIOChannel channel,
                            cvd::SharedFD shared_fd){
  if (!shared_fd->IsOpen()) {
    return false;
  }
  if (redirects_.count(channel)) {
    LOG(ERROR) << "Attempted multiple redirections of fd: "
               << static_cast<int>(channel);
    return false;
  }
  auto dup_fd = shared_fd->UNMANAGED_Dup();
  if (dup_fd < 0) {
    return false;
  }
  redirects_[channel] = dup_fd;
  return true;
}

Subprocess Command::Start(bool with_control_socket) const {
  auto cmd = ToCharPointers(command_);
  if (use_parent_env_) {
    return subprocess_impl(cmd.data(), nullptr, redirects_, inherited_fds_,
                           with_control_socket);
  } else {
    auto envp = ToCharPointers(env_);
    return subprocess_impl(cmd.data(), envp.data(), redirects_, inherited_fds_,
                           with_control_socket);
  }
}

int execute(const std::vector<std::string>& command,
            const std::vector<std::string>& env) {
  Command cmd(command[0]);
  for (size_t i = 1; i < command.size(); ++i) {
    cmd.AddParameter(command[i]);
  }
  cmd.SetEnvironment(env);
  auto subprocess = cmd.Start();
  if (!subprocess.Started()) {
    return -1;
  }
  return subprocess.Wait();
}
int execute(const std::vector<std::string>& command) {
  Command cmd(command[0]);
  for (size_t i = 1; i < command.size(); ++i) {
    cmd.AddParameter(command[i]);
  }
  auto subprocess = cmd.Start();
  if (!subprocess.Started()) {
    return -1;
  }
  return subprocess.Wait();
}

int execute_capture_output(const std::vector<std::string>& command,
                           std::string* output) {
  Command cmd(command[0]);
  for (size_t i = 1; i < command.size(); ++i) {
    cmd.AddParameter(command[i]);
  }
  cvd::SharedFD pipe_read, pipe_write;
  cvd::SharedFD::Pipe(&pipe_read, &pipe_write);
  cmd.RedirectStdIO(cvd::Subprocess::StdIOChannel::kStdOut, pipe_write);
  cmd.RedirectStdIO(cvd::Subprocess::StdIOChannel::kStdErr, pipe_write);

  auto subprocess = cmd.Start();
  if (!subprocess.Started()) {
    return -1;
  }
  {
    pipe_write->Close();
    // Force the destructor to run by moving it into a smaller scope.
    // This is necessary to close the write end of the pipe.
    cvd::Command forceDelete = std::move(cmd);
  }

  int read = cvd::ReadAll(pipe_read, output);
  if (read < 0) {
    LOG(FATAL) << "Could not read from pipe in execute_capture_output";
  }
  return subprocess.Wait();
}

}  // namespace cvd