sapling/contrib/chg/util.c

/*
 * Utility functions
 *
 * Copyright (c) 2011 Yuya Nishihara <yuya@tcha.org>
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2 or any later version.
 */

#include <errno.h>
#include <fcntl.h>
#include <signal.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/time.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>

#include "util.h"

static int colorenabled = 0;

static inline void fsetcolor(FILE* fp, const char* code) {
  if (!colorenabled)
    return;
  fprintf(fp, "\033[%sm", code);
}

static void vabortmsgerrno(int no, const char* fmt, va_list args) {
  fsetcolor(stderr, "1;31");
  fputs("chg: abort: ", stderr);
  vfprintf(stderr, fmt, args);
  if (no != 0)
    fprintf(stderr, " (errno = %d, %s)", no, strerror(no));
  fsetcolor(stderr, "");
  fputc('\n', stderr);
  exit(255);
}

void abortmsg(const char* fmt, ...) {
  va_list args;
  va_start(args, fmt);
  vabortmsgerrno(0, fmt, args);
  va_end(args);
}

void abortmsgerrno(const char* fmt, ...) {
  int no = errno;
  va_list args;
  va_start(args, fmt);
  vabortmsgerrno(no, fmt, args);
  va_end(args);
}

static int debugmsgenabled = 0;
static double debugstart = 0;

static double now() {
  struct timeval t;
  gettimeofday(&t, NULL);
  return t.tv_usec / 1e6 + t.tv_sec;
}

double chg_now() {
  return now();
}

void enablecolor(void) {
  colorenabled = 1;
}

void enabledebugmsg(void) {
  debugmsgenabled = 1;
  debugstart = now();
}

void debugmsg(const char* fmt, ...) {
  if (!debugmsgenabled)
    return;

  va_list args;
  va_start(args, fmt);
  fsetcolor(stderr, "1;30");
  fprintf(stderr, "chg: debug: %4.6f ", now() - debugstart);
  vfprintf(stderr, fmt, args);
  fsetcolor(stderr, "");
  fputc('\n', stderr);
  va_end(args);
}

void fchdirx(int dirfd) {
  int r = fchdir(dirfd);
  if (r == -1)
    abortmsgerrno("failed to fchdir");
}

void fsetcloexec(int fd) {
  int flags = fcntl(fd, F_GETFD);
  if (flags < 0)
    abortmsgerrno("cannot get flags of fd %d", fd);
  if (fcntl(fd, F_SETFD, flags | FD_CLOEXEC) < 0)
    abortmsgerrno("cannot set flags of fd %d", fd);
}

void* chg_mallocx(size_t size) {
  void* result = malloc(size);
  if (!result)
    abortmsg("failed to malloc");
  return result;
}

void* chg_reallocx(void* ptr, size_t size) {
  void* result = realloc(ptr, size);
  if (!result)
    abortmsg("failed to realloc");
  return result;
}

/*
 * Execute a shell command in mostly the same manner as system(), with the
 * give environment variables, after chdir to the given cwd. Returns a status
 * code compatible with the Python subprocess module.
 */
int runshellcmd(const char* cmd, const char* envp[], const char* cwd) {
  enum { F_SIGINT = 1, F_SIGQUIT = 2, F_SIGMASK = 4, F_WAITPID = 8 };
  unsigned int doneflags = 0;
  int status = 0;
  struct sigaction newsa, oldsaint, oldsaquit;
  sigset_t oldmask;

  /* block or mask signals just as system() does */
  memset(&newsa, 0, sizeof(newsa));
  newsa.sa_handler = SIG_IGN;
  newsa.sa_flags = 0;
  if (sigemptyset(&newsa.sa_mask) < 0)
    goto done;
  if (sigaction(SIGINT, &newsa, &oldsaint) < 0)
    goto done;
  doneflags |= F_SIGINT;
  if (sigaction(SIGQUIT, &newsa, &oldsaquit) < 0)
    goto done;
  doneflags |= F_SIGQUIT;

  if (sigaddset(&newsa.sa_mask, SIGCHLD) < 0)
    goto done;
  if (sigprocmask(SIG_BLOCK, &newsa.sa_mask, &oldmask) < 0)
    goto done;
  doneflags |= F_SIGMASK;

  pid_t pid = fork();
  if (pid < 0)
    goto done;
  if (pid == 0) {
    sigaction(SIGINT, &oldsaint, NULL);
    sigaction(SIGQUIT, &oldsaquit, NULL);
    sigprocmask(SIG_SETMASK, &oldmask, NULL);
    if (cwd && chdir(cwd) < 0)
      _exit(127);
    const char* argv[] = {"sh", "-c", cmd, NULL};
    if (envp) {
      execve("/bin/sh", (char**)argv, (char**)envp);
    } else {
      execv("/bin/sh", (char**)argv);
    }
    _exit(127);
  } else {
    if (waitpid(pid, &status, 0) < 0)
      goto done;
    doneflags |= F_WAITPID;
  }

done:
  if (doneflags & F_SIGINT)
    sigaction(SIGINT, &oldsaint, NULL);
  if (doneflags & F_SIGQUIT)
    sigaction(SIGQUIT, &oldsaquit, NULL);
  if (doneflags & F_SIGMASK)
    sigprocmask(SIG_SETMASK, &oldmask, NULL);

  /* no way to report other errors, use 127 (= shell termination) */
  if (!(doneflags & F_WAITPID))
    return 127;
  if (WIFEXITED(status))
    return WEXITSTATUS(status);
  if (WIFSIGNALED(status))
    return -WTERMSIG(status);
  return 127;
}