src/persistent_service.cc

// Copyright 2023 Google Inc. All Rights Reserved.
//
// 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 "persistent_service.h"

#include "async_loop.h"
#include "ipc_utils.h"
#include "util.h"

#ifdef _WIN32
#include <windows.h>
#else
#include <fcntl.h>
#include <sys/stat.h>
#include <unistd.h>
#endif

#define DEBUG 0

#if DEBUG
#include <stdio.h>
#define LOG(...)                                                         \
  (fprintf(stderr, "PERSISTENT_SERVICE "), fprintf(stderr, __VA_ARGS__), \
   fprintf(stderr, "\n"))
#else  // !DEBUG
#define LOG(...) LOGGER_LOG(logger_, __VA_ARGS__)
#endif  // !DEBUG

namespace {

// Sleep for |delay_ms| milliseconds.
void SleepMilliSeconds(int delay_ms) {
#ifdef _WIN32
  if (delay_ms > 0)
    ::Sleep(static_cast<DWORD>(delay_ms));
#else
  usleep(static_cast<useconds_t>(delay_ms) * 1000);
#endif
}

struct ProcessInfo {
#ifdef _WIN32
  HANDLE process_handle;
#else
  pid_t process_pid;
#endif
};

// Name of environment variable used to specify a log file for
// server log messages.
static const char kDebugPersistentServiceLogFileEnv[] =
    "DEBUG_PERSISTENT_SERVICE_LOG_FILE";

// Start a new process, with command-line |args|.
// Uses fork()/exec() on Posix, and CreateProcess on Win32.
// Return true on success, or set |*err| and return false on failure.
bool SpawnServerProcess(const PersistentService::Config& config,
                        const Logger& logger, ProcessInfo* info,
                        std::string* err) {
  if (config.command.empty()) {
    *err = "Empty command line!";
    return false;
  }

#ifdef _WIN32
  std::string command_string;
  for (const auto& arg : config.command) {
    std::string escaped;
    GetWin32EscapedString(arg, &escaped);
    command_string += escaped;
    command_string += ' ';
  }
  // Remove trailing space if any.
  if (!command_string.empty())
    command_string.resize(command_string.size() - 1u);

  SECURITY_ATTRIBUTES security_attributes = {};
  security_attributes.nLength = sizeof(SECURITY_ATTRIBUTES);
  security_attributes.bInheritHandle = TRUE;
  // Must be inheritable so subprocesses can dup to children.
  HANDLE nul =
      CreateFileA("NUL", GENERIC_READ | GENERIC_WRITE,
                  FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE,
                  &security_attributes, OPEN_EXISTING, 0, NULL);
  if (nul == INVALID_HANDLE_VALUE)
    Win32Fatal("couldn't open nul");

  HANDLE log = nul;
  std::string log_file = config.log_file;
  if (log_file.empty()) {
    // As a debug helper, use the log file from this environment variable.
    const char* env = getenv(kDebugPersistentServiceLogFileEnv);
    if (env)
      log_file = env;
  }
  if (!log_file.empty()) {
    log =
        CreateFileA(log_file.c_str(), STANDARD_RIGHTS_WRITE | FILE_APPEND_DATA,
                    FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE,
                    &security_attributes, OPEN_ALWAYS, 0, NULL);
    if (log == INVALID_HANDLE_VALUE)
      Win32Fatal("couldn't open log file", log_file.c_str());
  }

  STARTUPINFOA startup_info = {};
  startup_info.cb = sizeof(STARTUPINFO);
  startup_info.dwFlags = STARTF_USESTDHANDLES;
  startup_info.hStdInput = nul;
  startup_info.hStdOutput = log;
  startup_info.hStdError = log;
  PROCESS_INFORMATION process_info = {};

  // Ninja handles ctrl-c, except for subprocesses in console pools.
  DWORD process_flags = CREATE_NEW_PROCESS_GROUP;

  {
    // TODO(digit): Create new environment variable block and add/replace
    // the value in it, without touching the parent env.
    for (const auto& pair : config.env_vars) {
      const char* varname = pair.first.c_str();
      const char* value = pair.second.c_str();
      SetEnvironmentVariable(varname, value);
    }
  }

  // Do not prepend 'cmd /c' on Windows, this breaks command
  // lines greater than 8,191 chars.
  if (!CreateProcessA(NULL, (char*)command_string.c_str(), NULL, NULL,
                      /* inherit handles */ TRUE, process_flags, NULL, NULL,
                      &startup_info, &process_info)) {
    DWORD error = GetLastError();
    if (error == ERROR_FILE_NOT_FOUND) {
      CloseHandle(nul);
      if (log != nul)
        CloseHandle(log);
      // child_ is already NULL;
      *err =
          "CreateProcess failed: The system cannot find the file "
          "specified: [" +
          command_string + "]";
      return false;
    } else {
      fprintf(stderr, "\nCreateProcess failed. Command attempted:\n\"%s\"\n",
              command_string.c_str());
      const char* hint = NULL;
      // ERROR_INVALID_PARAMETER means the command line was formatted
      // incorrectly. This can be caused by a command line being too long or
      // leading whitespace in the command. Give extra context for this case.
      if (error == ERROR_INVALID_PARAMETER) {
        hint = "is the command line too long?";
      }
      Win32Fatal("CreateProcess", hint);
    }
  }

  CloseHandle(nul);
  if (log != nul)
    CloseHandle(log);

  CloseHandle(process_info.hThread);
  return true;
#else   // !_WIN32
  // Build arguments array for future exec() call.
  std::vector<char*> exec_args;
  exec_args.reserve(config.command.size() + 1);
  for (const auto& arg : config.command) {
    exec_args.push_back(const_cast<char*>(arg.data()));
  }
  exec_args.push_back(nullptr);

  pid_t process = fork();
  if (process < 0) {
    *err = "fork failed()!";
    return false;
  }

  if (process == 0) {
    // Create new session to not receive signals from parent process group.
    if (setsid() < 0) {
      fprintf(stderr, "ERROR: setsid() failed: %s\n", strerror(errno));
      exit(1);
    }

    // Change current working directory.
    if (!config.working_dir.empty()) {
      const char* work_dir = config.working_dir.c_str();
      if (chdir(work_dir) < 0)
        ErrnoFatal("chdir", work_dir);
    }

    // Redirect stdin to /dev/null and stdout/stderr to a log file if
    // PERSISTENT_LOG_FILE is set in the environment, or to
    // /dev/null otherwise.
    int null_fd = open("/dev/null", O_RDWR);
    if (null_fd < 0) {
      fprintf(stderr, "ERROR: open(/dev/null) failed: %s\n", strerror(errno));
      exit(1);
    }
    int log_fd = null_fd;
    std::string log_file = config.log_file;
    if (log_file.empty()) {
      // As a debug helper, use the log file from this environment variable.
      const char* env = getenv(kDebugPersistentServiceLogFileEnv);
      if (env)
        log_file = env;
    }
    if (!log_file.empty()) {
      log_fd = open(log_file.c_str(), O_WRONLY | O_APPEND | O_CREAT, 0755);
      if (log_fd < 0) {
        fprintf(stderr, "ERROR: open(%s) failed: %s\n", log_file.c_str(),
                strerror(errno));
        exit(1);
      }
    }

    fflush(stdout);
    fflush(stderr);

    dup2(null_fd, 0);
    dup2(log_fd, 1);
    dup2(log_fd, 2);

    // Set extra environment variables.
    for (const auto& pair : config.env_vars) {
      const char* varname = pair.first.c_str();
      const char* value = pair.second.c_str();
      setenv(varname, value, 1);
    }

    fprintf(stderr, "\n\nSTARTING NEW PERSISTENT SERVER: %s\n",
            config.command[0].c_str());
    execv(config.command[0].c_str(), exec_args.data());
    fprintf(stderr, "ERROR: exec() failed: %s\n", strerror(errno));
    exit(1);
  }
  // In parent process, do not do anything.
  info->process_pid = process;
  return true;
#endif  // !_WIN32
}

// Type of commands received from the server.
enum ServerCommandType {
  kServerCommandTypeStop,
  kServerCommandTypeGetPid,
  kServerCommandTypeClientQuery,
};

std::string ToDebugString(ServerCommandType command_type) {
  const char* type_str;
  switch (command_type) {
  case kServerCommandTypeStop:
    type_str = "Stop";
    break;
  case kServerCommandTypeGetPid:
    type_str = "GetPid";
    break;
  case kServerCommandTypeClientQuery:
    type_str = "ClientQuery";
    break;
  default:
    type_str = "UNKNOWN??";
  }
  return StringFormat("%s(%d)", type_str, command_type);
}

}  // namespace

///////////////////////////////////////////////////////////////////////////
///
///   C L I E N T   S I D E
///

PersistentService::Client::Client(const std::string& service_name)
    : service_name_(service_name) {}

PersistentService::Client::~Client() = default;

void PersistentService::Client::SetLogger(Logger logger) {
  logger_ = std::move(logger);
}

bool PersistentService::Client::HasServer() const {
  return IpcService::HasServer(service_name_);
}

int PersistentService::Client::GetServerPid() const {
  std::string err;
  ScopedHandle client = RawConnect(&err);
  if (!client)
    return -1;

  uint8_t request_type = kServerCommandTypeGetPid;
  int server_pid = -1;
  if (!RemoteWrite(request_type, client, &err) ||
      !RemoteRead(server_pid, client, &err)) {
    return -1;
  }
  return server_pid;
}

bool PersistentService::Client::StopServer(std::string* err) const {
  ScopedHandle client = RawConnect(err);
  if (!client)
    return false;

  uint8_t request_type = kServerCommandTypeStop;
  if (!client.Write(&request_type, sizeof(request_type), err)) {
    *err = StringFormat("Could not stop server: %s", err->c_str());
    return false;
  }

  return true;
}

bool PersistentService::Client::WaitForServerShutdown() {
  std::string error;
  for (int try_count = 0; try_count < 20; ++try_count) {
    if (!HasServer()) {
      LOG("WaitForServerShutdown: no server, try_count=%d", try_count);
      return true;
    }
    LOG("WaitForServerShutdown: server still running, sleeping 100ms, "
        "try_count=%d",
        try_count);

    SleepMilliSeconds(100);
  }
  LOG("WaitForServerShutdown: server still running, abandoning!");
  return false;
}

ScopedHandle PersistentService::Client::Connect(const Config& config,
                                                std::string* error) {
  bool try_again = true;

  while (true) {
    // This is subtle, the connection attempt, or sending info
    // can fail when the server is currently shutting down, so
    // do not exit immediately before retrying at least once.
    //
    // Similarly, receiving the version check string could also
    // fail if the server shut down too fast before the client
    // receives the answer. So also ensure we try once again
    // if this happens.
    do {
      LOG("Trying to connect to server.");
      ScopedHandle client = ConnectOrStartServer(config, error);
      if (!client)
        break;

      LOG("Sending version info to server.");
      if (!RemoteWrite(config.version_info, client, error)) {
        *error = "Could not send version info: " + *error;
        break;
      }

      // Receive the |compatible| flag that indicates that the server is
      // compatible with the current build plan.
      std::string version_check;
      if (!RemoteRead(version_check, client, error)) {
        *error = "Could not read version check result: " + *error;
        break;
      }

      if (version_check.empty()) {
        // Good, return the handle now.
        return client;
      }

      // The server was not compatible with the current client.
      // Assume it exited, and start another one at least once.
      LOG("Incompatible server version: %s", version_check.c_str());

    } while (0);

    if (!try_again) {
      // Already tried once, so report failure.
      LOG("Failed to connect to or start server!");
      return {};
    }

    error->clear();
    try_again = false;
    LOG("Waiting for incompatible server shutdown.");
    if (!WaitForServerShutdown()) {
      *error = "Could not shutdown incompatible server!?!";
      return {};
    }
  }
}

ScopedHandle PersistentService::Client::RawConnect(std::string* err) const {
  bool did_timeout = false;
  return IpcService::ConnectTo(service_name_, -1, &did_timeout, err);
}

ScopedHandle PersistentService::Client::ConnectOrStartServer(
    const Config& config, std::string* err) const {
  int retry_count = 5;
  int retry_delay_ms = 10;
  bool server_started = false;
  ProcessInfo info = {};
  while (true) {
    // Try to connect to the server first.
    ScopedHandle client = RawConnect(err);
    if (client) {
      LOG("Got client connection to server!");
      uint8_t request_type = kServerCommandTypeClientQuery;
      if (!client.Write(&request_type, sizeof(request_type), err)) {
        LOG("ERROR: Could not write query type, did server disconnect?: %s",
            err->c_str());
        client.Close();
      } else {
        LOG("Sent query type");
      }
      return client;
    }

    if (!server_started) {
      LOG("No initial connection. Spawning server");
      // Spawn a server if one wasn't already started.
      if (!SpawnServerProcess(config, logger_, &info, err)) {
        LOG("Could not spawn server: %s", err->c_str());
        return {};
      }
      server_started = true;

    } else if (retry_count == 0) {
      LOG("Failure to connect to server, exiting retry loop: %s", err->c_str());
      return {};
    } else {
      --retry_count;
      if (retry_delay_ms < 1024)
        retry_delay_ms *= 2;
    }

    LOG("Waiting for %ld milliseconds", (long)retry_delay_ms);
    SleepMilliSeconds(retry_delay_ms);
  }
}

///////////////////////////////////////////////////////////////////////////
///
///   S E R V E R   S I D E
///

PersistentService::Server::Server(const std::string& service_name)
    : service_name_(service_name) {}

void PersistentService::Server::SetLogger(Logger logger) {
  logger_ = std::move(logger);
}

bool PersistentService::Server::BindService(std::string* err) {
  if (service_) {
    *err = "Server already started!";
    return false;
  }
  LOG("Trying to start local exclusive service: %s", service_name_.c_str());
  service_ = IpcService::StartServer(service_name_, err);
  if (!service_) {
    LOG("Got error %s", err->c_str());
  } else {
    LOG("Got it!");
  }
  return !!service_;
}

void PersistentService::Server::RunServerThenExit(
    const VersionCheckHandler& version_check_handler,
    const RequestHandler& request_handler) {
  std::string error;
  if (!service_ && !BindService(&error)) {
    Error("Could not start server: %s", error.c_str());
    exit(1);
  }
  LOG("Server started");
  RunServerThenExitInternal(version_check_handler, std::move(request_handler));
}

void PersistentService::Server::RunServerThenExitInternal(
    const VersionCheckHandler& version_check_handler,
    const RequestHandler& request_handler) {
  AsyncLoop& async_loop = AsyncLoop::Get();
  std::string error;

  while (true) {
    bool did_timeout = false;
    LOG("Waiting for new client connection");

    ScopedHandle client =
        service_.AcceptPeer(connection_timeout_ms_, &did_timeout, &error);

    if (!client) {
      if (did_timeout) {
        LOG("Timeout waiting for client connection");
      } else {
        LOG("Could not accept peer connection: %s", error.c_str());
      }
      break;
    }

    // Get command, which can be 'kill' or 'query' at the moment.
    uint8_t request_type_byte = 0;
    if (!client.Read(&request_type_byte, sizeof(request_type_byte), &error)) {
      LOG("Could not read request type from new client: %s", error.c_str());
      break;
    }
    auto request_type = static_cast<ServerCommandType>(request_type_byte);

    LOG("Request type from new client: %s",
        ToDebugString(request_type).c_str());

    if (request_type == kServerCommandTypeStop) {
      LOG("Client asking server to stop!");
      break;
    }
    if (request_type == kServerCommandTypeGetPid) {
      LOG("Client asking for server pid!");
#ifdef _WIN32
      int pid = static_cast<int>(GetCurrentProcessId());
#else
      int pid = getpid();
#endif
      if (!RemoteWrite(pid, client, &error)) {
        LOG("Could not send pid %d back to client!: %s", pid, error.c_str());
        break;
      }
      LOG("Sent pid %d to client. Looping", pid);
      continue;
    }
    if (request_type != kServerCommandTypeClientQuery) {
      LOG("Unknown request type: %s",
          ToDebugString(static_cast<ServerCommandType>(request_type)).c_str());
      break;
    }

    LOG("Accepted client query, checking version info");
    std::string version_info;
    if (!RemoteRead(version_info, client, &error)) {
      LOG("Could not read client version info: %s", error.c_str());
      break;
    }

    std::string version_check = version_check_handler(version_info);
    if (!RemoteWrite(version_check, client, &error)) {
      LOG("Could not write version check result to client: %s", error.c_str());
      break;
    }

    if (!version_check.empty()) {
      LOG("Incompatible client version info: %s", version_check.c_str());
      break;
    }

    LOG("Version check successful, invoking service request handler");
    if (!request_handler(std::move(client))) {
      LOG("Request handler returned false, exiting server loop");
      break;
    }

    // Clear pending interrupts if any.
    async_loop.ClearInterrupt();
  }
  LOG("Exiting!");

  // Since calling ::exit() directly here prevents the destructors
  // of the variables defined in this function from running, stop
  // the service for proper socket pid file cleanup on MacOS.
  service_.Stop();

  ::exit(0);
}