codesign_parallel.c

/*
 * codesign_parallel - Parallel wrapper around macOS codesign binary
 * Distributes code signing work across multiple CPU cores for improved throughput
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdbool.h>
#include <unistd.h>
#include <signal.h>
#include <errno.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <sys/sysctl.h>
#include <spawn.h>

// Yeah, hardcoded, because it is hard to imagine it being anywhere else on macOS
#define CODESIGN_PATH "/usr/bin/codesign"

static pid_t *worker_pids = NULL;
static int worker_count = 0;
static volatile bool is_signal_received = false;

static bool should_fail_fast = false;
static bool should_be_verbose = false;
static int parallel_count;

// Arguments to pass to codesign
char **codesign_argv;
int codesign_argc;

// List of binaries to sign
char **binaries;
int binaries_count;

// Per-worker assignment structure
typedef struct {
  int start_index;
  int count;
} worker_assignment;

void help() {
  fprintf(stderr, "Usage: codesign_parallel [codesign_parallel options] [codesign options] -- BINARY1 [BINARY2 ...]\n");
}

static int get_cpu_count(void) {
  int mib[2] = { CTL_HW, HW_NCPU };
  int ncpu;
  size_t len = sizeof(ncpu);

  if (sysctl(mib, 2, &ncpu, &len, NULL, 0) == -1) {
    // sysctl failed?!
    return 1;
  }

  return ncpu > 0 ? ncpu : 1;
}

// Check if verbose flag is present in arguments and set should_be_verbose
static void check_if_verbose_specified(int argc, char **argv) {
  for (int i = 0; i < argc; i++) {
    if (strncmp(argv[i], "-v", 2) == 0 || strncmp(argv[i], "--v", 3) == 0) {
      should_be_verbose = true;
      return;
    }
  }
}

static int parse_codesign_args(char **argv, int delimiter_index) {
  codesign_argv = malloc(sizeof(char *) * (delimiter_index + 1));
  if (!codesign_argv) {
    fprintf(stderr, "codesign_parallel: error: memory allocation failed\n");
    return -1;
  }

  // Parse arguments before delimiter
  codesign_argc = 0;

  for (int i = 1; i < delimiter_index; i++) {
    if (strncmp(argv[i], "--parallel=", 11) == 0) {
      parallel_count = atoi(argv[i] + 11);
      if (parallel_count <= 0) {
        fprintf(stderr, "codesign_parallel: error: invalid --parallel value\n");
        free(codesign_argv);
        return -1;
      }

    } else if (strcmp(argv[i], "--fail-fast") == 0) {
      should_fail_fast = true;

    } else {
      // Pass through to codesign
      codesign_argv[codesign_argc++] = argv[i];
    }
  }

  codesign_argv[codesign_argc] = NULL;

  return 0;
}

static int parse_args(int argc, char **argv) {
  for (int i = 1; i < argc; i++) {
    if (strcmp(argv[i], "--help") == 0 || strcmp(argv[i], "-h") == 0) {
      help();
      return -1;
    }
  }

  // Find the delimiter
  int delimiter_index = -1;
  for (int i = 1; i < argc; i++) {
    if (strcmp(argv[i], "--") == 0) {
      delimiter_index = i;
      break;
    }
  }

  if (delimiter_index == -1) {
    fprintf(stderr, "codesign_parallel: error: missing '--' delimiter before binary list\n");
    help();
    return -1;
  }

  // Count binaries after delimiter
  binaries_count = argc - delimiter_index - 1;
  if (binaries_count == 0) {
    fprintf(stderr, "codesign_parallel: warning: no binaries specified\n");
    return 0;
  }

  // Set binaries pointer
  binaries = &argv[delimiter_index + 1];

  return parse_codesign_args(argv, delimiter_index);
}

static int validate_codesign_binary(void) {
  if (access(CODESIGN_PATH, X_OK) != 0) {
    fprintf(stderr, "codesign_parallel: error: %s does not exists or is not executable: %s\n", CODESIGN_PATH, strerror(errno));
    return -1;
  }

  return 0;
}

static void send_signal_to_all_workers(int sig) {
  for (int i = 0; i < worker_count; i++) {
    if (worker_pids[i] > 0) {
      kill(worker_pids[i], sig);
    }
  }
}

static void signal_handler(int sig) {
  is_signal_received = true;
  send_signal_to_all_workers(sig);
}

// Calculate contiguous chunk assignments for workers
static worker_assignment *calculate_assignments(int binaries_count, int cpu_count) {
  worker_count = cpu_count < binaries_count ? cpu_count : binaries_count;

  int count = binaries_count / worker_count;
  int reminder = binaries_count % worker_count;

  // Print summary if verbose
  if (should_be_verbose) {
    fprintf(stderr, "codesign_parallel: Signing %d binaries across %d workers (%d per worker", binaries_count, worker_count, count);
    if (reminder > 0) {
      fprintf(stderr, " + %d extra for the last worker)", reminder);
    } else {
      fprintf(stderr, ")");
    }
    fprintf(stderr, "\n");
  }

  worker_assignment *assignments = malloc(sizeof(worker_assignment) * worker_count);
  if (!assignments) {
    fprintf(stderr, "codesign_parallel: error: memory allocation failed\n");
    return NULL;
  }

  int current_index = 0;

  for (int i = 0; i < worker_count; i++) {
    assignments[i].start_index = current_index;
    assignments[i].count = count;

    if (i == worker_count - 1) {
      assignments[i].count += reminder;
    }

    if (current_index + assignments[i].count > binaries_count) {
      assignments[i].count = binaries_count - current_index;
    }

    current_index += assignments[i].count;
  }

  return assignments;
}

// Spawn a single worker process
// The worker signs binaries from start_index to start_index + count - 1
static pid_t spawn_worker(int start_index, int count) {
  extern char **environ;

  // Build argv for codesign: codesign [args...] binary1 binary2 ... NULL
  int total_argc = 1 + codesign_argc + count + 1;
  char **child_argv = malloc(sizeof(char *) * total_argc);
  if (!child_argv) {
    fprintf(stderr, "codesign_parallel: error: memory allocation failed\n");
    return -1;
  }

  int idx = 0;
  child_argv[idx++] = CODESIGN_PATH;

  // Add codesign arguments
  for (int i = 0; i < codesign_argc; i++) {
    child_argv[idx++] = codesign_argv[i];
  }

  // Add binaries for this worker
  for (int i = 0; i < count; i++) {
    child_argv[idx++] = binaries[start_index + i];
  }

  child_argv[idx] = NULL;

  // Need to dump the argv for debugging?
  // for (int i = 0; i < idx; i++) {
  //   printf("%s ", child_argv[i]);
  // }
  // printf("\n");

  // Spawn!
  pid_t pid;
  int ret = posix_spawn(&pid, CODESIGN_PATH, NULL, NULL, child_argv, environ);

  free(child_argv);

  if (ret != 0) {
    fprintf(stderr, "codesign_parallel: error: failed to spawn worker: %s\n", strerror(ret));
    return -1;
  }

  return pid;
}

static int spawn_workers(const worker_assignment *assignments) {
  worker_pids = calloc(worker_count, sizeof(pid_t));
  if (!worker_pids) {
    fprintf(stderr, "codesign_parallel: error: memory allocation failed\n");
    return -1;
  }

  for (int i = 0; i < worker_count; i++) {
    pid_t pid = spawn_worker(assignments[i].start_index, assignments[i].count);
    if (pid < 0) {
      // Kill already spawned workers
      send_signal_to_all_workers(SIGTERM);
      return -1;
    }
    worker_pids[i] = pid;
  }

  return 0;
}

static int wait_for_workers_to_finish() {
  int failed_worker_count = 0;
  int remaining_worker_count = worker_count;

  while (remaining_worker_count > 0) {
    int status;
    pid_t pid = wait(&status);

    if (pid < 0) {
      if (errno == EINTR) {
        // Interrupted by signal, continue waiting
        continue;
      }

      // No more children
      break;
    }

    remaining_worker_count--;

    // Check exit status
    bool is_worker_failed = false;
    if (WIFEXITED(status)) {
      if (WEXITSTATUS(status) != 0) {
        is_worker_failed = true;
      }

    } else if (WIFSIGNALED(status)) {
      is_worker_failed = true;
    }

    if (is_worker_failed) {
      failed_worker_count = 1;
      if (should_fail_fast && remaining_worker_count > 0) {
        // Kill remaining workers
        send_signal_to_all_workers(SIGTERM);
      }
    }

    // Mark this PID as done
    for (int i = 0; i < worker_count; i++) {
      if (worker_pids[i] == pid) {
        worker_pids[i] = 0;
        break;
      }
    }
  }

  return failed_worker_count;
}

static void setup_signal_handlers(void) {
  struct sigaction sa;
  memset(&sa, 0, sizeof(sa));
  sa.sa_handler = signal_handler;
  sigemptyset(&sa.sa_mask);
  sa.sa_flags = 0;

  sigaction(SIGINT, &sa, NULL);
  sigaction(SIGTERM, &sa, NULL);
}

int main(int argc, char **argv) {
  if (argc < 2) {
    help();
    return 1;
  }

  if (parse_args(argc, argv) != 0) {
    return 1;
  }

  check_if_verbose_specified(codesign_argc, codesign_argv);

  if (binaries_count == 0) {
    return 0;
  }

  if (validate_codesign_binary() != 0) {
    return 1;
  }

  int cpu_count = parallel_count > 0 ? parallel_count : get_cpu_count();

  worker_assignment *assignments = calculate_assignments(binaries_count, cpu_count);

  setup_signal_handlers();

  if (spawn_workers(assignments) != 0) {
    return 1;
  }

  // Not that we actually need to free this memory, given the fact that we shall exit soon anyway,
  // but let'd be a good citizen.
  free(assignments);

  int ret = wait_for_workers_to_finish();

  // Check if we received a signal
  if (is_signal_received) {
    return 1;
  }

  return ret;
}