cleanmysourcetree.c

   1 #define _GNU_SOURCE
   2
   3 #include <dirent.h>
   4 #include <unistd.h>
   5 #include <stdio.h>
   6 #include <err.h>
   7 #include <string.h>
   8 #include <stdbool.h>
   9 #include <signal.h>
  10 #include <poll.h>
  11 #include <errno.h>
  12 #include <sys/types.h>
  13 #include <sys/stat.h>
  14 #include <sys/inotify.h>
  15 #include <sys/wait.h>
  16 #include <sys/prctl.h>
  17 #include <uthash.h>
  18
  19 enum delstate {
  20   DELSTATE_NO, /* source file that was accessed, keep it */
  21   DELSTATE_MAYBE, /* source file that was not touched so far, delete it if it stays unused */
  22   DELSTATE_YES /* generated file in a run with delete_new = true, delete it */
  23 };
  24
  25 // describes a file or a directory that was created during compilation
  26 typedef struct {
  27   UT_hash_handle hh;
  28   enum delstate delstate;
  29   char name[]; /* path relative to project root; hashmap key */
  30 } file;
  31 file *hashed_files = NULL;
  32
  33 // describes a directory that existed before compilation
  34 typedef struct {
  35   UT_hash_handle hh;
  36   int watch_descriptor; /* hashmap key */
  37   char name[];
  38 } directory;
  39 file *hashed_dirs = NULL;
  40
  41 int inotify_fd = -1;
  42
  43 void *xmalloc(size_t len) {
  44   void *ret = malloc(len);
  45   if (ret == NULL && len != 0)
  46     err(1, "memory allocation failure");
  47   return ret;
  48 }
  49
  50 static void add_file(const char *file_path) {
  51   file *f;
  52   HASH_FIND_STR(hashed_files, file_path, f);
  53   if (f)
  54     errx(1, "we tried to add the file '%s' twice, probably because you're messing "
  55       "around in the filesystem and removing and adding files during the initial scan "
  56       "step. stop it.", file_path);
  57
  58   f = xmalloc(sizeof(*f) + strlen(file_path) + 1);
  59   strcpy(f->name, file_path);
  60   f->delstate = DELSTATE_MAYBE;
  61   HASH_ADD_STR(hashed_files, name, f);
  62 }
  63
  64 static void add_dir(const char *dir_path) {
  65   directory *d = xmalloc(sizeof(*d) + strlen(dir_path) + 1);
  66   strcpy(d->name, dir_path);
  67
  68   // Don't register watches on files to avoid running into the fs.inotify.max_user_watches limit.
  69   // Instead, register watches on all directories. There shouldn't be too many of those.
  70   d->watch_descriptor = inotify_add_watch(inotify_fd, dir_path,
  71     IN_OPEN | IN_CREATE | IN_MOVED_TO | IN_EXCL_UNLINK | IN_ONLYDIR);
  72   if (d->watch_descriptor == -1)
  73     err(1, "unable to add inotify watch for '%s'", dir_path);
  74
  75   directory *d_existing;
  76   HASH_FIND_INT(hashed_dirs, &d->watch_descriptor, d_existing);
  77   if (d_existing)
  78     errx(1, "the kernel says we watched the same directory twice. whatever you're doing "
  79       "to make the filesystem behave that way, stop it.");
  80   HASH_ADD_INT(hashed_dirs, watch_descriptor, d);
  81 }
  82
  83 static void add_files_recursive(const char *current_path) {
  84   DIR *d = opendir((current_path[0] == '\0') ? "." : current_path);
  85   if (d == NULL)
  86     err(1, "unable to open directory '%s'", current_path);
  87
  88   struct dirent entry;
  89   struct dirent *r_entry;
  90   while (1) {
  91     if (readdir_r(d, &entry, &r_entry))
  92       errx(1, "readdir_r failed");
  93     if (r_entry == NULL)
  94       return;
  95     if (strcmp(entry.d_name, ".") == 0 || strcmp(entry.d_name, "..") == 0)
  96       continue;
  97
  98     char file_path[strlen(current_path) + 1 + strlen(entry.d_name) + 1];
  99     if (current_path[0] == '\0') {
 100       strcpy(file_path, entry.d_name);
 101     } else {
 102       sprintf(file_path, "%s/%s", current_path, entry.d_name);
 103     }
 104
 105     struct stat st;
 106     if (lstat(file_path, &st))
 107       err(1, "unable to stat '%s'", file_path);
 108     switch (st.st_mode & S_IFMT) {
 109       case S_IFREG:
 110         add_file(file_path);
 111         break;
 112       case S_IFDIR:
 113         // First recurse, then add the watch. This avoids spamming ourselves with irrelevant
 114         // directory open events.
 115         add_files_recursive(file_path);
 116         add_dir(file_path);
 117         break;
 118       default:
 119         break;
 120     }
 121   }
 122 }
 123
 124 volatile bool child_quit = false;
 125
 126 void handle_sigchld(int n) {
 127   if (waitpid(-1, NULL, WNOHANG) == -1)
 128     err(1, "wait for child failed");
 129   child_quit = true;
 130   puts("processing remaining events...");
 131 }
 132
 133 void usage(void) {
 134   puts("invocation: cleanmysourcetree <deletenew/keepnew> [<target directory>]");
 135   puts("deletenew will cause files and directories created by "
 136     "the compilation process to be deleted.");
 137   exit(1);
 138 }
 139
 140 int main(int argc, char **argv) {
 141   bool delete_new;
 142
 143   if (argc == 3) {
 144     if (chdir(argv[2]))
 145       err(1, "unable to chdir to specified directory");
 146   } else if (argc != 2) {
 147     usage();
 148   }
 149   if (strcmp(argv[1], "deletenew") == 0) {
 150     delete_new = true;
 151   } else if (strcmp(argv[1], "keepnew") == 0) {
 152     delete_new = false;
 153   } else {
 154     usage();
 155   }
 156
 157   // collect filenames, init watches
 158   inotify_fd = inotify_init1(IN_CLOEXEC | IN_NONBLOCK);
 159   if (inotify_fd == -1)
 160     err(1, "unable to open inotify fd");
 161   add_files_recursive("");
 162
 163   sigset_t sigchild_mask;
 164   sigemptyset(&sigchild_mask);
 165   sigaddset(&sigchild_mask, SIGCHLD);
 166   if (sigprocmask(SIG_SETMASK, &sigchild_mask, NULL))
 167     err(1, "sigprocmask failed");
 168   if (signal(SIGCHLD, handle_sigchld) == SIG_ERR)
 169     err(1, "unable to register signal handler");
 170
 171   pid_t child = fork();
 172   if (child == -1)
 173     err(1, "can't fork");
 174   if (child == 0) {
 175     prctl(PR_SET_PDEATHSIG, SIGTERM); /* stupid racy API :/ */
 176     puts("dropping you into an interactive shell now. compile the project, then "
 177       "exit the shell.");
 178     // uhhh... yeah.
 179     system("$SHELL");
 180     exit(0);
 181   }
 182
 183   // process inotify events until the child is dead
 184   // and there are no more pending events
 185   struct pollfd fds[1] = { { .fd = inotify_fd, .events = POLLIN } };
 186   struct timespec zero_timeout_ts = { .tv_sec = 0, .tv_nsec = 0};
 187   sigset_t empty_mask;
 188   sigemptyset(&empty_mask);
 189   while (1) {
 190     int r = ppoll(fds, 1, child_quit ? &zero_timeout_ts : NULL, &empty_mask);
 191     if (r == -1 && errno == EINTR)
 192       continue;
 193     if (r == -1)
 194       err(1, "ppoll failed");
 195     if (r == 0)
 196       break;
 197
 198     char buf[20 * (sizeof(struct inotify_event) + NAME_MAX + 1)];
 199     ssize_t read_res = read(inotify_fd, buf, sizeof(buf));
 200     if (read_res == -1)
 201       err(1, "read from inotify fd failed");
 202     if (read_res < sizeof(struct inotify_event))
 203       errx(1, "short/empty read from inotify");
 204     struct inotify_event *e = (void*)buf;
 205     while ((char*)e != buf + read_res) {
 206       if (e->mask & IN_Q_OVERFLOW)
 207         errx(1, "inotify queue overflow detected");
 208       if (e->len != 0 && (e->mask & (IN_OPEN | IN_CREATE | IN_MOVED_TO))) {
 209         directory *dir;
 210         HASH_FIND_INT(hashed_dirs, &e->wd, dir);
 211         if (dir == NULL)
 212           errx(1, "unable to find dir by inotify wd, bug!");
 213         char path[strlen(dir->name) + e->len];
 214         sprintf(path, "%s/%s", dir->name, e->name);
 215         file *f;
 216         HASH_FIND_STR(hashed_files, path, f);
 217         // if f is NULL, this is a file/directory generated during compilation
 218         // or a pre-existing directory
 219         if (f != NULL) {
 220           if (f->delstate == DELSTATE_MAYBE)
 221             f->delstate = DELSTATE_NO;
 222         } else if (delete_new && (e->mask & (IN_CREATE | IN_MOVED_TO))) {
 223           f = xmalloc(sizeof(*f) + strlen(path) + 1);
 224           strcpy(f->name, path);
 225           f->delstate = DELSTATE_YES;
 226           HASH_ADD_STR(hashed_files, name, f);
 227         }
 228       }
 229       // step to next event in buffer. yuck.
 230       e = (struct inotify_event *)((char *)e + sizeof(struct inotify_event) + e->len);
 231     }
 232   }
 233
 234   // reset signal handling
 235   signal(SIGCHLD, SIG_DFL);
 236   if (sigprocmask(SIG_SETMASK, &empty_mask, NULL))
 237     err(1, "sigprocmask failed");
 238
 239   close(inotify_fd);
 240   puts("inotify event collection phase is over, deleting stuff...");
 241
 242   // if we want to delete generated files and folders, we haven't seen
 243   // any events for files in generated folders. therefore, to delete
 244   // those folders, they need to be rm -rf'ed. I don't want to write
 245   // logic for that manually, so just execute rm.
 246   for (file *f = hashed_files; f != NULL; f = f->hh.next) {
 247     if (f->delstate == DELSTATE_NO) continue;
 248
 249     pid_t rm_pid = fork();
 250     if (rm_pid == -1)
 251       err(1, "unable to fork for rm");
 252     if (rm_pid == 0) {
 253       execlp("rm", "rm", "-rfv", "--", f->name, NULL);
 254       err(1, "unable to invoke rm");
 255     }
 256     if (wait(NULL) != rm_pid)
 257       err(1, "waiting for rm failed");
 258   }
 259   puts("cleanup complete");
 260   return 0;
 261 }