compctld.c

   1 #define _GNU_SOURCE /* struct ucred */
   2 #include <assert.h>
   3 #include <errno.h>
   4 #include <signal.h>
   5 #include <stdio.h>
   6 #include <stdlib.h>
   7 #include <string.h>
   8 #include <sys/socket.h>
   9 #include <sys/un.h>
  10 #include <syslog.h>
  11 #include <unistd.h>
  12
  13 #include "cgroup.h"
  14 #include "common.h"
  15
  16
  17 #define begins_with(s_, a_) (!strncmp(s_, a_, sizeof(s_) - 1))
  18
  19
  20 void
  21 logperror(const char *s)
  22 {
  23         syslog(LOG_ERR, "%s: %s", s, strerror(errno));
  24 }
  25
  26 void
  27 memory_limits(size_t *minuser, size_t *mincomp, size_t *maxcomp, size_t *total)
  28 {
  29         FILE *f = fopen("/proc/meminfo", "r");
  30         char line[1024];
  31         while (fgets(line, sizeof(line), f)) {
  32                 if (begins_with("MemTotal:", line)) {
  33                         *total = 0;
  34                         sscanf(line, "MemTotal:%zu", total);
  35                         *total *= 1024;
  36                         break;
  37                 }
  38         }
  39         fclose(f);
  40
  41         *minuser = *total * split_ratio;
  42         if (*minuser < static_minfree)
  43                 *minuser = static_minfree;
  44         if (*minuser > static_maxfree)
  45                 *minuser = static_maxfree;
  46
  47         *mincomp = static_minfree;
  48         *maxcomp = *total - *minuser;
  49         if (*maxcomp < 0) *maxcomp = 0;
  50         /* maxcomp < mincomp may happen; they are used in different
  51          * settings. */
  52 }
  53
  54 size_t
  55 get_default_mem_limit(void)
  56 {
  57         size_t minuser, mincomp, maxcomp, total;
  58         memory_limits(&minuser, &mincomp, &maxcomp, &total);
  59         return maxcomp;
  60 }
  61
  62
  63 void
  64 cgroup_init(void)
  65 {
  66         if (cgroup_setup(chier, "memory") < 0)
  67                 exit(EXIT_FAILURE);
  68         int ret = cgroup_create(chier, cgroup);
  69         if (ret < 0)
  70                 exit(EXIT_FAILURE);
  71         if (ret > 0) {
  72                 /* CGroup newly created, set limit. */
  73                 if (cgroup_set_mem_limit(chier, cgroup, get_default_mem_limit()) < 0)
  74                         exit(EXIT_FAILURE);
  75         }
  76 }
  77
  78
  79 int
  80 main(int argc, char *argv[])
  81 {
  82         /* Do this while everyone can still see the error. */
  83         cgroup_init();
  84
  85         pid_t p = fork();
  86         if (p < 0) {
  87                 perror("fork");
  88                 exit(EXIT_FAILURE);
  89         }
  90         if (p > 0)
  91                 exit(EXIT_SUCCESS);
  92
  93         fclose(stderr);
  94         fclose(stdout);
  95         fclose(stdin);
  96         openlog("compctl", LOG_PID, LOG_DAEMON);
  97         cgroup_perror = logperror;
  98
  99         setsid();
 100
 101         int s = socket(AF_UNIX, SOCK_STREAM, 0);
 102         int on = 1; setsockopt(s, SOL_SOCKET, SO_PASSCRED, &on, sizeof(on));
 103
 104         struct sockaddr_un sun = { .sun_family = AF_UNIX, .sun_path = SOCKFILE };
 105         if (bind(s, (struct sockaddr *) &sun, sizeof(sun.sun_family) + strlen(sun.sun_path) + 1) < 0) {
 106                 logperror(SOCKFILE);
 107                 exit(EXIT_FAILURE);
 108         }
 109
 110         listen(s, 10);
 111
 112         int fd;
 113         while ((fd = accept(s, NULL, NULL)) >= 0) {
 114                 /* We handle only a single client at a time. This means
 115                  * that it is rather easy to write a script that will DOS
 116                  * the daemon, this is just an attack vector we ignore. */
 117                 /* TODO: alarm() to wake from stuck clients. */
 118
 119                 /* Decode the message with credentials. */
 120
 121                 struct msghdr msg;
 122                 char *errmsg;
 123                 if (recvmsg(fd, &msg, MSG_WAITALL) <= 0) {
 124                         errmsg = "recvmsg";
 125 sockerror:
 126                         logperror(errmsg);
 127                         close(fd);
 128                         continue;
 129                 }
 130                 struct ucred *cred;
 131                 struct cmsghdr *cmsg;
 132                 cmsg = CMSG_FIRSTHDR(&msg);
 133                 if (cmsg == NULL || cmsg->cmsg_len != CMSG_LEN(sizeof(struct ucred))) {
 134                         errmsg = "cmsg";
 135                         goto sockerror;
 136                 }
 137                 if (cmsg->cmsg_level != SOL_SOCKET || cmsg->cmsg_type != SCM_CREDENTIALS) {
 138                         errmsg = "cmsg designation";
 139                         goto sockerror;
 140                 }
 141                 cred = (struct ucred *) CMSG_DATA(cmsg);
 142
 143                 FILE *f = fdopen(fd, "r");
 144                 char line[1024];
 145                 fgets(line, sizeof(line), f);
 146                 size_t linelen = strlen(line);
 147                 if (linelen < 2 || strcmp(&line[linelen - 2], "\r\n")) {
 148                         syslog(LOG_WARNING, "protocol error (%s)", line);
 149                         fclose(f);
 150                         continue;
 151                 }
 152                 line[linelen - 2] = 0;
 153
 154                 /* Analyze command */
 155                 if (!strcmp("blessme", line)) {
 156                         syslog(LOG_INFO, "new computation process %d", cred->pid);
 157                         if (cgroup_add_task(chier, cgroup, cred->pid) < 0)
 158                                 fprintf(f, "0 error: %s\r\n", strerror(errno));
 159                         else
 160                                 fputs("1 blessed\r\n", f);
 161
 162                 } else if (begins_with("stop ", line)) {
 163                         pid_t pid = atoi(line + sizeof("stop "));
 164
 165                         /* Sanity check. */
 166                         if (pid < 10 || pid > 32768) {
 167                                 syslog(LOG_WARNING, "stop: invalid pid (%d)", pid);
 168                                 fputs("0 invalid pid\r\n", f);
 169                                 fclose(f);
 170                                 continue;
 171                         }
 172                         if (!cgroup_is_task_in_cgroup(chier, cgroup, pid)) {
 173                                 fputs("0 task not marked as computation\r\n", f);
 174                                 fclose(f);
 175                                 continue;
 176                         }
 177
 178                         syslog(LOG_INFO, "stopping process %d (request by pid %d uid %d)", pid, cred->pid, cred->uid);
 179                         kill(pid, SIGTERM);
 180                         /* TODO: Grace period and then kill with SIGKILL. */
 181                         fputs("1 task stopped\r\n", f);
 182
 183                 } else if (!strcmp("stopall", line)) {
 184                         pid_t *tasks;
 185                         int tasks_n = cgroup_task_list(chier, cgroup, &tasks);
 186                         if (tasks_n < 0) {
 187                                 fprintf(f, "0 error: %s\r\n", strerror(errno));
 188                                 fclose(f);
 189                                 continue;
 190                         }
 191                         for (int i = 0; i < tasks_n; i++) {
 192                                 syslog(LOG_INFO, "stopping process %d (mass request by pid %d uid %d)", tasks[i], cred->pid, cred->uid);
 193                                 kill(tasks[i], SIGTERM);
 194                         }
 195                         /* TODO: Grace period and then kill with SIGKILL. */
 196                         fprintf(f, "1 %d tasks stopped\r\n", tasks_n);
 197                         free(tasks);
 198
 199                 } else if (begins_with("limitmem ", line)) {
 200                         size_t limit = atol(line + sizeof("limitmem "));
 201                         size_t minuser, mincomp, maxcomp, total;
 202                         memory_limits(&minuser, &mincomp, &maxcomp, &total);
 203
 204                         /* Sanity check. */
 205                         if (limit < 1024 || limit > total) {
 206                                 syslog(LOG_WARNING, "limitmem: invalid limit (%zu)", limit);
 207                                 fputs("0 invalid limit value\r\n", f);
 208                                 fclose(f);
 209                                 continue;
 210                         }
 211
 212                         if (limit < mincomp) {
 213                                 fprintf(f, "-1 at least %zuM must remain available for computations.\r\n", mincomp / 1048576);
 214                                 fclose(f);
 215                                 continue;
 216                         }
 217                         if (total - limit < minuser) {
 218                                 fprintf(f, "-2 at least %zuM must remain available for users.\r\n", minuser / 1048576);
 219                                 fclose(f);
 220                                 continue;
 221                         }
 222
 223                         syslog(LOG_INFO, "setting limit %zu (request by pid %d uid %d)", limit, cred->pid, cred->uid);
 224                         if (cgroup_set_mem_limit(chier, cgroup, limit) < 0)
 225                                 fprintf(f, "0 error: %s\r\n", strerror(errno));
 226                         else
 227                                 fputs("1 limit set\r\n", f);
 228
 229                 } else {
 230                         syslog(LOG_WARNING, "invalid command (%s)", line);
 231                 }
 232
 233                 fclose(f);
 234         }
 235
 236         return EXIT_FAILURE;
 237 }