/*
 * Copyright (C) 2020 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.
 */

#ifndef INCLUDE_PERFETTO_EXT_BASE_SUBPROCESS_H_
#define INCLUDE_PERFETTO_EXT_BASE_SUBPROCESS_H_

#include <condition_variable>
#include <functional>
#include <initializer_list>
#include <mutex>
#include <optional>
#include <string>
#include <thread>
#include <vector>

#include "perfetto/base/build_config.h"
#include "perfetto/base/logging.h"
#include "perfetto/base/platform_handle.h"
#include "perfetto/base/proc_utils.h"
#include "perfetto/ext/base/event_fd.h"
#include "perfetto/ext/base/pipe.h"
#include "perfetto/ext/base/scoped_file.h"

namespace perfetto {
namespace base {

// Handles creation and lifecycle management of subprocesses, taking care of
// all subtleties involved in handling processes on UNIX.
// This class allows to deal with macro two use-cases:
// 1) fork() + exec() equivalent: for spawning a brand new process image.
//    This happens when |args.exec_cmd| is not empty.
//    This is safe to use even in a multi-threaded environment.
// 2) fork(): for spawning a process and running a function.
//    This happens when |args.posix_entrypoint_for_testing| is not empty.
//    This is intended only for tests as it is extremely subtle.
//    This mode must be used with extreme care. Before the entrypoint is
//    invoked all file descriptors other than stdin/out/err and the ones
//    specified in |args.preserve_fds| will be closed, to avoid each process
//    retaining a dupe of other subprocesses pipes. This however means that
//    any non trivial calls (including logging) must be avoided as they might
//    refer to FDs that are now closed. The entrypoint should really be used
//    just to signal a pipe or similar for synchronizing sequencing in tests.

//
// This class allows to control stdin/out/err pipe redirection and takes care
// of keeping all the pipes pumped (stdin) / drained (stdout/err), in a similar
// fashion of python's subprocess.Communicate()
// stdin: is always piped and closed once the |args.input| buffer is written.
// stdout/err can be either:
//   - dup()ed onto the parent process stdout/err.
//   - redirected onto /dev/null.
//   - piped onto a buffer (see output() method). There is only one output
//     buffer in total. If both stdout and stderr are set to kBuffer mode, they
//     will be merged onto the same. There doesn't seem any use case where they
//     are needed distinctly.
//
// Some caveats worth mentioning:
// - It always waitpid()s, to avoid leaving zombies around. If the process is
//   not terminated by the time the destructor is reached, the dtor will
//   send a SIGKILL and wait for the termination.
// - After fork()-ing it will close all file descriptors, preserving only
//   stdin/out/err and the fds listed in |args.preserve_fds|.
// - On Linux/Android, the child process will be SIGKILL-ed if the calling
//   thread exists, even if the Subprocess is std::move()-d onto another thread.
//   This happens by virtue PR_SET_PDEATHSIG, which is used to avoid that
//   child processes are leaked in the case of a crash of the parent (frequent
//   in tests). However, the child process might still be leaked if execing
//   a setuid/setgid binary (see man 2 prctl).
//
// Usage:
// base::Subprocess p({"/bin/cat", "-"});
// (or equivalently:
//     base::Subprocess p;
//     p.args.exec_cmd.push_back("/bin/cat");
//     p.args.exec_cmd.push_back("-");
//  )
// p.args.stdout_mode = base::Subprocess::kBuffer;
// p.args.stderr_mode = base::Subprocess::kInherit;
// p.args.input = "stdin contents";
// p.Call();
// (or equivalently:
//     p.Start();
//     p.Wait();
// )
// EXPECT_EQ(p.status(), base::Subprocess::kTerminated);
// EXPECT_EQ(p.returncode(), 0);
class Subprocess {
 public:
  enum Status {
    kNotStarted = 0,  // Before calling Start() or Call().
    kRunning,         // After calling Start(), before Wait().
    kTerminated,      // The subprocess terminated, either successfully or not.
                      // This includes crashes or other signals on UNIX.
  };

  enum class OutputMode {
    kInherit = 0,  // Inherit's the caller process stdout/stderr.
    kDevNull,      // dup() onto /dev/null.
    kBuffer,       // dup() onto a pipe and move it into the output() buffer.
    kFd,           // dup() onto the passed args.fd.
  };

  enum class InputMode {
    kBuffer = 0,  // dup() onto a pipe and write args.input on it.
    kDevNull,     // dup() onto /dev/null.
  };

  // Input arguments for configuring the subprocess behavior.
  struct Args {
    Args(std::initializer_list<std::string> _cmd = {}) : exec_cmd(_cmd) {}
    Args(Args&&) noexcept;
    Args& operator=(Args&&);
    // If non-empty this will cause an exec() when Start()/Call() are called.
    std::vector<std::string> exec_cmd;

#if !PERFETTO_BUILDFLAG(PERFETTO_OS_WIN)
    // If non-empty, it changes the argv[0] argument passed to exec. If
    // unset, argv[0] == exec_cmd[0]. This is to handle cases like:
    // exec_cmd = {"/proc/self/exec"}, argv0: "my_custom_test_override".
    std::string posix_argv0_override_for_testing;

    // If non-empty this will be invoked on the fork()-ed child process, after
    // stdin/out/err has been redirected and all other file descriptor are
    // closed. It is valid to specify both |exec_cmd| AND
    // |posix_entrypoint_for_testing|. In this case the latter will be invoked
    // just before the exec() call, but after having closed all fds % stdin/o/e.
    // This is for synchronization barriers in tests.
    std::function<void()> posix_entrypoint_for_testing;

    // When set, will will move the process to the given process group. If set
    // and zero, it will create a new process group. Effectively this calls
    // setpgid(0 /*self_pid*/, posix_proc_group_id).
    // This can be used to avoid that subprocesses receive CTRL-C from the
    // terminal, while still living in the same session.
    std::optional<pid_t> posix_proc_group_id{};
#endif

    // If non-empty, replaces the environment passed to exec().
    std::vector<std::string> env;

    // The file descriptors in this list will not be closed.
    std::vector<int> preserve_fds;

    // The data to push in the child process stdin, if input_mode ==
    // InputMode::kBuffer.
    std::string input;

    InputMode stdin_mode = InputMode::kBuffer;
    OutputMode stdout_mode = OutputMode::kInherit;
    OutputMode stderr_mode = OutputMode::kInherit;

    base::ScopedPlatformHandle out_fd;

    // Returns " ".join(exec_cmd), quoting arguments.
    std::string GetCmdString() const;
  };

  struct ResourceUsage {
    uint32_t cpu_utime_ms = 0;
    uint32_t cpu_stime_ms = 0;
    uint32_t max_rss_kb = 0;
    uint32_t min_page_faults = 0;
    uint32_t maj_page_faults = 0;
    uint32_t vol_ctx_switch = 0;
    uint32_t invol_ctx_switch = 0;

    uint32_t cpu_time_ms() const { return cpu_utime_ms + cpu_stime_ms; }
  };

  explicit Subprocess(std::initializer_list<std::string> exec_cmd = {});
  Subprocess(Subprocess&&) noexcept;
  Subprocess& operator=(Subprocess&&);
  ~Subprocess();  // It will KillAndWaitForTermination() if still alive.

  // Starts the subprocess but doesn't wait for its termination. The caller
  // is expected to either call Wait() or Poll() after this call.
  void Start();

  // Wait for process termination. Can be called more than once.
  // Args:
  //   |timeout_ms| = 0: wait indefinitely.
  //   |timeout_ms| > 0: wait for at most |timeout_ms|.
  // Returns:
  //  True: The process terminated. See status() and returncode().
  //  False: Timeout reached, the process is still running. In this case the
  //         process will be left in the kRunning state.
  bool Wait(int timeout_ms = 0);

  // Equivalent of Start() + Wait();
  // Returns true if the process exited cleanly with return code 0. False in
  // any othe case.
  bool Call(int timeout_ms = 0);

  Status Poll();

  // Sends a signal (SIGKILL if not specified) and wait for process termination.
  void KillAndWaitForTermination(int sig_num = 0);

  PlatformProcessId pid() const { return s_->pid; }

  // The accessors below are updated only after a call to Poll(), Wait() or
  // KillAndWaitForTermination().
  // In most cases you want to call Poll() rather than these accessors.

  Status status() const { return s_->status; }
  int returncode() const { return s_->returncode; }
  bool timed_out() const { return s_->timed_out; }

  // This contains both stdout and stderr (if the corresponding _mode ==
  // OutputMode::kBuffer). It's non-const so the caller can std::move() it.
  std::string& output() { return s_->output; }
  const std::string& output() const { return s_->output; }

  const ResourceUsage& posix_rusage() const { return *s_->rusage; }

  Args args;

 private:
  // The signal/exit code used when killing the process in case of a timeout.
  static const int kTimeoutSignal;

  Subprocess(const Subprocess&) = delete;
  Subprocess& operator=(const Subprocess&) = delete;

  // This is to deal robustly with the move operators, without having to
  // manually maintain member-wise move instructions.
  struct MovableState {
    base::Pipe stdin_pipe;
    base::Pipe stdouterr_pipe;
    PlatformProcessId pid;
    Status status = kNotStarted;
    int returncode = -1;
    std::string output;  // Stdin+stderr. Only when OutputMode::kBuffer.
    std::unique_ptr<ResourceUsage> rusage{new ResourceUsage()};
    bool timed_out = false;
#if PERFETTO_BUILDFLAG(PERFETTO_OS_WIN)
    std::thread stdouterr_thread;
    std::thread stdin_thread;
    ScopedPlatformHandle win_proc_handle;
    ScopedPlatformHandle win_thread_handle;

    base::EventFd stdouterr_done_event;
    std::mutex mutex;  // Protects locked_outerr_buf and the two pipes.
    std::string locked_outerr_buf;
#else
    base::Pipe exit_status_pipe;
    size_t input_written = 0;
    std::thread waitpid_thread;
#endif
  };

#if PERFETTO_BUILDFLAG(PERFETTO_OS_WIN)
  static void StdinThread(MovableState*, std::string input);
  static void StdoutErrThread(MovableState*);
#else
  void TryPushStdin();
  void TryReadStdoutAndErr();
  void TryReadExitStatus();
  bool PollInternal(int poll_timeout_ms);
#endif

  std::unique_ptr<MovableState> s_;
};

}  // namespace base
}  // namespace perfetto

#endif  // INCLUDE_PERFETTO_EXT_BASE_SUBPROCESS_H_