sin/main.c

#define _GNU_SOURCE
#include <ctype.h>
#include <dirent.h>
#include <errno.h>
#include <fcntl.h>
#include <limits.h>
#include <signal.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/types.h>
#include <time.h>
#include <unistd.h>

#include <systemd/sd-bus.h>

#define EXIT_OK 0
#define EXIT_ARGS 1
#define EXIT_INIT 2
#define EXIT_RUNTIME 3

#define MAX_PATTERNS 128
#define MAX_PATTERN_LEN 4096U
#define MIN_POLL_INTERVAL 0.1
#define MAX_POLL_INTERVAL 3600.0
#define DEFAULT_POLL_INTERVAL 2.0
#define CMDLINE_INITIAL_SIZE 4096U
#define CMDLINE_MAX_SIZE (1024U * 1024U)
#define DBUS_CALL_TIMEOUT_USEC UINT64_C(10000000)

static const char *valid_what[] = {"idle",
                                   "sleep",
                                   "shutdown",
                                   "handle-power-key",
                                   "handle-suspend-key",
                                   "handle-hibernate-key",
                                   "handle-lid-switch",
                                   NULL};

static const char *valid_mode[] = {"block", "delay", NULL};

static volatile sig_atomic_t running = 1;
static volatile sig_atomic_t cleanup_requested = 0;

static void handle_sig(int sig) {
  (void)sig;
  running = 0;
  cleanup_requested = 1;
}

struct Inhibitor {
  sd_bus *bus;
  int fd;
  const char *what;
  const char *who;
  const char *why;
  const char *mode;
};

static int validate_value(const char *value, const char **allowed) {
  if (!value || !allowed)
    return 0;
  for (int i = 0; allowed[i] != NULL; i++) {
    if (strcmp(value, allowed[i]) == 0)
      return 1;
  }
  return 0;
}

static int inhibitor_init(struct Inhibitor *i, const char *what,
                          const char *who, const char *why, const char *mode) {
  int r;

  if (!i || !what || !who || !why || !mode) {
    fprintf(stderr, "inhibitor_init: NULL parameter\n");
    return -EINVAL;
  }

  if (!validate_value(what, valid_what)) {
    fprintf(stderr, "Invalid --what value: %s\n", what);
    fprintf(stderr, "Valid values: idle, sleep, shutdown, handle-power-key, "
                    "handle-suspend-key, handle-hibernate-key, "
                    "handle-lid-switch\n");
    return -EINVAL;
  }

  if (!validate_value(mode, valid_mode)) {
    fprintf(stderr, "Invalid --mode value: %s\n", mode);
    fprintf(stderr, "Valid values: block, delay\n");
    return -EINVAL;
  }

  memset(i, 0, sizeof(*i));
  i->fd = -1;
  i->what = what;
  i->who = who;
  i->why = why;
  i->mode = mode;

  r = sd_bus_open_system(&i->bus);
  if (r < 0) {
    fprintf(stderr, "Failed to connect to system bus: %s\n", strerror(-r));
    return r;
  }

  // Set a reasonable timeout for D-Bus calls
  r = sd_bus_set_method_call_timeout(i->bus, DBUS_CALL_TIMEOUT_USEC);
  if (r < 0) {
    fprintf(stderr, "Warning: Failed to set D-Bus timeout: %s\n", strerror(-r));
    // Non-fatal, continue
  }

  return 0;
}

static void inhibitor_free(struct Inhibitor *i) {
  if (!i)
    return;
  if (i->fd >= 0) {
    close(i->fd);
    i->fd = -1;
  }
  if (i->bus) {
    sd_bus_flush_close_unref(i->bus);
    i->bus = NULL;
  }
}

static int inhibitor_acquire(struct Inhibitor *i) {
  if (!i)
    return -EINVAL;

  if (i->fd >= 0)
    return 0;

  if (!i->bus) {
    fprintf(stderr, "inhibitor_acquire: bus not initialized\n");
    return -EINVAL;
  }

  sd_bus_message *m = NULL;
  sd_bus_error error = SD_BUS_ERROR_NULL;
  int r;

  r = sd_bus_call_method(i->bus, "org.freedesktop.login1",
                         "/org/freedesktop/login1",
                         "org.freedesktop.login1.Manager", "Inhibit", &error,
                         &m, "ssss", i->what, i->who, i->why, i->mode);
  if (r < 0) {
    fprintf(stderr, "Inhibit call failed: %s\n",
            error.message ? error.message : strerror(-r));
    sd_bus_error_free(&error);
    return r;
  }
  sd_bus_error_free(&error);

  int received_fd = -1;
  r = sd_bus_message_read(m, "h", &received_fd);
  if (r < 0) {
    fprintf(stderr, "Failed to read fd from Inhibit reply: %s\n", strerror(-r));
    sd_bus_message_unref(m);
    return r;
  }

  if (received_fd < 0) {
    fprintf(stderr, "Received invalid fd from Inhibit: %d\n", received_fd);
    sd_bus_message_unref(m);
    return -EIO;
  }

  // Duplicate FD to ensure we own it and can close it independently
  int dupfd = fcntl(received_fd, F_DUPFD_CLOEXEC, 3);
  if (dupfd < 0) {
    // Fallback without CLOEXEC if F_DUPFD_CLOEXEC not supported
    dupfd = dup(received_fd);
    if (dupfd < 0) {
      int saved_errno = errno;
      fprintf(stderr, "dup failed: %s\n", strerror(errno));
      sd_bus_message_unref(m);
      return -saved_errno;
    }
    int flags = fcntl(dupfd, F_GETFD);
    if (flags >= 0) {
      (void)fcntl(dupfd, F_SETFD, flags | FD_CLOEXEC);
    }
  }

  i->fd = dupfd;
  sd_bus_message_unref(m);
  return 0;
}

static void inhibitor_release(struct Inhibitor *i) {
  if (!i)
    return;
  if (i->fd >= 0) {
    close(i->fd);
    i->fd = -1;
  }
}

/* Read /proc/<pid>/cmdline into a heap-allocated, NUL-terminated string.
 * Caller must free(). Returns NULL on error or empty cmdline.
 */
static char *read_cmdline(pid_t pid) {
  char path[64];
  int n = snprintf(path, sizeof(path), "/proc/%d/cmdline", (int)pid);
  if (n < 0 || (size_t)n >= sizeof(path)) {
    return NULL;
  }

  int fd = open(path, O_RDONLY | O_CLOEXEC);
  if (fd < 0)
    return NULL;

  size_t cap = CMDLINE_INITIAL_SIZE;
  char *buf = malloc(cap);
  if (!buf) {
    close(fd);
    return NULL;
  }

  size_t len = 0;
  for (;;) {
    if (len >= CMDLINE_MAX_SIZE) {
      // Prevent unbounded growth
      free(buf);
      close(fd);
      return NULL;
    }

    ssize_t r = read(fd, buf + len, cap - len);
    if (r < 0) {
      if (errno == EINTR)
        continue;
      free(buf);
      close(fd);
      return NULL;
    }
    if (r == 0)
      break;

    len += (size_t)r;
    if (len >= cap) {
      size_t new_cap = cap * 2U;
      if (new_cap > CMDLINE_MAX_SIZE)
        new_cap = CMDLINE_MAX_SIZE;
      if (new_cap <= cap) {
        break;
      }
      char *nb = realloc(buf, new_cap);
      if (!nb) {
        free(buf);
        close(fd);
        return NULL;
      }
      buf = nb;
      cap = new_cap;
    }
  }
  close(fd);

  if (len == 0) {
    free(buf);
    return NULL;
  }

  // cmdline is NUL-separated; make a printable space-separated string
  for (size_t j = 0; j < len; ++j) {
    if (buf[j] == '\0')
      buf[j] = ' ';
  }

  if (len >= cap) {
    char *nb = realloc(buf, len + 1);
    if (!nb) {
      free(buf);
      return NULL;
    }
    buf = nb;
  }
  buf[len] = '\0';
  return buf;
}

// Return 1 if any process' cmdline contains any of the patterns.
static int any_process_running(char **patterns, int npats) {
  if (!patterns || npats <= 0)
    return 0;

  DIR *d = opendir("/proc");
  if (!d) {
    fprintf(stderr, "opendir /proc failed: %s\n", strerror(errno));
    return 0;
  }

  struct dirent *ent;
  int found = 0;

  while ((ent = readdir(d)) != NULL && !found) {
    const char *name = ent->d_name;

    // Skip non-numeric entries
    if (!isdigit((unsigned char)name[0]))
      continue;

    // Verify all characters are digits
    int all_digits = 1;
    for (const char *p = name; *p; ++p) {
      if (!isdigit((unsigned char)*p)) {
        all_digits = 0;
        break;
      }
    }
    if (!all_digits)
      continue;

    // Convert to pid
    char *endptr;
    long pid_long = strtol(name, &endptr, 10);
    if (*endptr != '\0' || pid_long <= 0 || pid_long > INT_MAX)
      continue;

    pid_t pid = (pid_t)pid_long;
    char *cmd = read_cmdline(pid);
    if (!cmd)
      continue;

    // Check all patterns
    for (int i = 0; i < npats; ++i) {
      if (strstr(cmd, patterns[i]) != NULL) {
        found = 1;
        break;
      }
    }
    free(cmd);
  }
  closedir(d);
  return found;
}

static void usage(FILE *f, const char *prog) {
  fprintf(f,
          "Usage: %s -n <pattern> [-n <pattern> ...] "
          "[--what=WHAT] [--mode=MODE]\n"
          "       [--who=WHO] [--why=WHY] [--poll=SECONDS]\n"
          "\n"
          "Options:\n"
          "  -n, --name PATTERN    Process command-line substring to watch "
          "(repeatable)\n"
          "  --what WHAT           What to inhibit (default: idle)\n"
          "                        Valid: idle, sleep, shutdown, "
          "handle-power-key,\n"
          "                               handle-suspend-key, "
          "handle-hibernate-key,\n"
          "                               handle-lid-switch\n"
          "  --mode MODE           Mode: block or delay (default: block)\n"
          "  --who WHO             Who string for Inhibit (default: "
          "inhibit-monitor)\n"
          "  --why WHY             Why string for Inhibit (default: prevent "
          "system idle...)\n"
          "  --poll SECONDS        Poll interval in seconds (%.1f-%.1f, "
          "default: %.1f)\n"
          "  -h, --help            Show this help\n",
          prog, MIN_POLL_INTERVAL, MAX_POLL_INTERVAL, DEFAULT_POLL_INTERVAL);
}

static int validate_string_param(const char *param, const char *name,
                                 size_t max_len) {
  if (!param || param[0] == '\0') {
    fprintf(stderr, "Error: %s cannot be empty\n", name);
    return 0;
  }
  size_t len = strlen(param);
  if (len > max_len) {
    fprintf(stderr, "Error: %s too long (max %zu characters)\n", name, max_len);
    return 0;
  }
  return 1;
}

// Safe realloc wrapper that doesn't leak on failure
static void *safe_realloc(void *ptr, size_t new_size) {
  void *new_ptr = realloc(ptr, new_size);
  if (!new_ptr && new_size > 0) {
    return NULL; // realloc failed, original pointer still valid
  }
  return new_ptr;
}

int main(int argc, char **argv) {
  double poll = DEFAULT_POLL_INTERVAL;
  const char *what = "idle";
  const char *who = "inhibit-monitor";
  const char *why = "prevent system idle while important process runs";
  const char *mode = "block";

  char **patterns = NULL;
  int npats = 0;
  int exit_code = EXIT_OK;

  // Oh Clap how I miss you so...
  for (int i = 1; i < argc; ++i) {
    if ((strcmp(argv[i], "-n") == 0 || strcmp(argv[i], "--name") == 0)) {
      if (i + 1 >= argc) {
        fprintf(stderr, "Error: %s requires an argument\n", argv[i]);
        usage(stderr, argv[0]);
        free(patterns);
        return EXIT_ARGS;
      }
      i++;

      if (npats >= MAX_PATTERNS) {
        fprintf(stderr, "Error: Too many patterns (max %d)\n", MAX_PATTERNS);
        free(patterns);
        return EXIT_ARGS;
      }

      if (!validate_string_param(argv[i], "pattern", MAX_PATTERN_LEN)) {
        free(patterns);
        return EXIT_ARGS;
      }

      char **new_patterns =
          safe_realloc(patterns, (size_t)(npats + 1) * sizeof(char *));
      if (!new_patterns) {
        fprintf(stderr, "Error: Memory allocation failed\n");
        free(patterns);
        return EXIT_INIT;
      }
      patterns = new_patterns;
      patterns[npats++] = argv[i];

    } else if (strncmp(argv[i], "--what=", 7U) == 0) {
      what = argv[i] + 7;
      if (!validate_string_param(what, "--what", 256U)) {
        free(patterns);
        return EXIT_ARGS;
      }

    } else if (strncmp(argv[i], "--mode=", 7U) == 0) {
      mode = argv[i] + 7;
      if (!validate_string_param(mode, "--mode", 256U)) {
        free(patterns);
        return EXIT_ARGS;
      }

    } else if (strncmp(argv[i], "--who=", 6U) == 0) {
      who = argv[i] + 6;
      if (!validate_string_param(who, "--who", 256U)) {
        free(patterns);
        return EXIT_ARGS;
      }

    } else if (strncmp(argv[i], "--why=", 6U) == 0) {
      why = argv[i] + 6;
      if (!validate_string_param(why, "--why", 1024U)) {
        free(patterns);
        return EXIT_ARGS;
      }

    } else if (strncmp(argv[i], "--poll=", 7U) == 0) {
      char *endptr;
      errno = 0;
      poll = strtod(argv[i] + 7, &endptr);

      if (errno != 0 || *endptr != '\0' || endptr == argv[i] + 7) {
        fprintf(stderr, "Error: Invalid --poll value: %s\n", argv[i] + 7);
        free(patterns);
        return EXIT_ARGS;
      }

      if (poll < MIN_POLL_INTERVAL || poll > MAX_POLL_INTERVAL) {
        fprintf(stderr, "Error: --poll must be between %.1f and %.1f seconds\n",
                MIN_POLL_INTERVAL, MAX_POLL_INTERVAL);
        free(patterns);
        return EXIT_ARGS;
      }

    } else if (strcmp(argv[i], "--help") == 0 || strcmp(argv[i], "-h") == 0) {
      usage(stdout, argv[0]);
      free(patterns);
      return EXIT_OK;

    } else if (argv[i][0] == '-' && argv[i][1] == 'n' && argv[i][2] != '\0') {
      // Support "-npattern" form
      if (npats >= MAX_PATTERNS) {
        fprintf(stderr, "Error: Too many patterns (max %d)\n", MAX_PATTERNS);
        free(patterns);
        return EXIT_ARGS;
      }

      if (!validate_string_param(argv[i] + 2, "pattern", MAX_PATTERN_LEN)) {
        free(patterns);
        return EXIT_ARGS;
      }

      char **new_patterns =
          safe_realloc(patterns, (size_t)(npats + 1) * sizeof(char *));
      if (!new_patterns) {
        fprintf(stderr, "Error: Memory allocation failed\n");
        free(patterns);
        return EXIT_INIT;
      }
      patterns = new_patterns;
      patterns[npats++] = argv[i] + 2;

    } else {
      fprintf(stderr, "Error: Unknown argument: %s\n", argv[i]);
      usage(stderr, argv[0]);
      free(patterns);
      return EXIT_ARGS;
    }
  }

  if (npats == 0) {
    fprintf(stderr, "Error: At least one -n/--name pattern is required\n");
    usage(stderr, argv[0]);
    free(patterns);
    return EXIT_ARGS;
  }

  struct sigaction sa;
  memset(&sa, 0, sizeof(sa));
  sa.sa_handler = handle_sig;
  sigemptyset(&sa.sa_mask);
  sa.sa_flags = SA_RESTART;

  if (sigaction(SIGINT, &sa, NULL) < 0) {
    fprintf(stderr, "Warning: Failed to set SIGINT handler: %s\n",
            strerror(errno));
  }
  if (sigaction(SIGTERM, &sa, NULL) < 0) {
    fprintf(stderr, "Warning: Failed to set SIGTERM handler: %s\n",
            strerror(errno));
  }

  struct Inhibitor inh;
  if (inhibitor_init(&inh, what, who, why, mode) < 0) {
    free(patterns);
    return EXIT_INIT;
  }

  // Main monitoring loop
  while (running) {
    int found = any_process_running(patterns, npats);

    if (found) {
      if (inh.fd < 0) {
        fprintf(stderr, "Target process found - acquiring inhibitor\n");
        if (inhibitor_acquire(&inh) < 0) {
          fprintf(
              stderr,
              "Failed to acquire inhibitor; retrying after poll interval\n");
          exit_code = EXIT_RUNTIME;
        } else {
          fprintf(stderr, "Inhibitor acquired (fd=%d)\n", inh.fd);
        }
      }

      // Keep inhibitor active while target(s) present
      while (running && any_process_running(patterns, npats)) {
        struct timespec ts;
        ts.tv_sec = (time_t)poll;
        ts.tv_nsec = (long)((poll - (time_t)poll) * 1e9);

        // Use nanosleep with proper interrupt handling
        while (nanosleep(&ts, &ts) < 0) {
          if (errno != EINTR || !running)
            break;
        }
      }

      if (inh.fd >= 0 && running) {
        fprintf(stderr, "Target exited - releasing inhibitor\n");
        inhibitor_release(&inh);
      }
    } else {
      // No target process; sleep and check again
      struct timespec ts;
      ts.tv_sec = (time_t)poll;
      ts.tv_nsec = (long)((poll - (time_t)poll) * 1e9);

      while (nanosleep(&ts, &ts) < 0) {
        if (errno != EINTR || !running)
          break;
      }
    }
  }

  // Cleanup
  if (cleanup_requested) {
    fprintf(stderr, "Shutting down gracefully...\n");
  }

  inhibitor_free(&inh);
  free(patterns);

  return exit_code;
}