src/throttle.c

   1 /*
   2
   3 throttle.c -- throttle utility to make use of taskd.pl throttle services
   4 Copyright (C) 2015 Kyle J. McKay.  All rights reserved.
   5
   6 This program is free software; you can redistribute it and/or
   7 modify it under the terms of the GNU General Public License
   8 as published by the Free Software Foundation; either version 2
   9 of the License, or (at your option) any later version.
  10
  11 This program is distributed in the hope that it will be useful,
  12 but WITHOUT ANY WARRANTY; without even the implied warranty of
  13 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  14 GNU General Public License for more details.
  15
  16 You should have received a copy of the GNU General Public License
  17 along with this program; if not, write to the Free Software
  18 Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
  19
  20 */
  21
  22 /*
  23   This utility provides a means to make use of taskd.pl's throttle services.
  24   It is simply started with options specifying the throttle class, the item
  25   description and an optional supplementary message to included in case the
  26   request is throttled.
  27
  28   If the request is allowed to proceed, the non-option arguments specify a
  29   command and arguments to be execvp'd.  If the request is throttled, a result
  30   suitable for a CGI is written to stdout with a 503 status.  Alternatively
  31   on failure a non-zero exit code can be returned in lieu of CGI output.
  32 */
  33
  34 #ifndef SOCKET_FILE
  35 #error SOCKET_FILE must be defined to full path to taskd.pl socket
  36 #endif
  37
  38 #ifndef _POSIX_C_SOURCE
  39 #define _POSIX_C_SOURCE 200809
  40 #endif
  41
  42 #include <errno.h>
  43 #include <fcntl.h>
  44 #include <stddef.h>
  45 #include <stdio.h>
  46 #include <stdlib.h>
  47 #include <string.h>
  48 #include <unistd.h>
  49 #include <signal.h>
  50 #include <sys/types.h>
  51 #include <sys/socket.h>
  52 #include <sys/time.h>
  53 #include <sys/select.h>
  54 #include <sys/un.h>
  55
  56 #ifndef FD_COPY
  57 #define FD_COPY(src,dst) ((*dst)=(*src))
  58 #endif
  59
  60 #define USAGE \
  61 "Usage: throttle [-t] [-p] [-e] -c <class> -d <desc> [-m <msg>] command [arg]...\n" \
  62 "  -h                show this help\n" \
  63 "  -e                exit with error (3 = exec fail, 10 = connect fail,\n" \
  64 "                    11 = request fail, 20 = request error, 21 = throttled)\n" \
  65 "                    on failure but do not output anything to standard output\n" \
  66 "  -t                throttle on connect failure (default is proceed)\n" \
  67 "  -p                proceed on throttle request error (default is throttle)\n" \
  68 "  -c <class>        required throttle class\n" \
  69 "  -d <desc>         required description of item in class\n" \
  70 "  -m <msg>          supplementary message to include with 503 error\n" \
  71 "  command [arg]...  command and arguments to execvp on proceed\n"
  72
  73 static int ok_connect_fail = 1;
  74 static int ok_request_fail = 0;
  75 static int err_on_failure = 0;
  76
  77 static void
  78 die(
  79     const char *msg)
  80 {
  81     fprintf(stderr, "fatal: %s\n", msg ? msg : "error");
  82     exit(EXIT_FAILURE);
  83 }
  84
  85 static void
  86 diefailure(
  87     const char *msg)
  88 {
  89     fprintf(stderr, "fatal: %s: %s\n", msg ? msg : "failed", strerror(errno));
  90     exit(EXIT_FAILURE);
  91 }
  92
  93 static void
  94 dieusage(
  95     const char *msg)
  96 {
  97     fprintf(stderr, "fatal: %s\n", msg ? msg : "error");
  98     fprintf(stderr, "%s", USAGE);
  99     exit(EXIT_FAILURE);
 100 }
 101
 102 static int
 103 isalphach(
 104     char ch)
 105 {
 106     return ('A' <= ch && ch <= 'Z') ||
 107         ('a' <= ch && ch <= 'z');
 108 }
 109
 110 static size_t
 111 get_escxml_size(
 112     const char *str)
 113 {
 114     size_t ans = 0;
 115     char ch;
 116
 117     if (!str)
 118         return ans;
 119     while ((ch = *str++) != 0) {
 120         switch (ch) {
 121             case '<':
 122             case '>':
 123                 ans += 4;
 124                 break;
 125             case '&':
 126                 ans += 5;
 127                 break;
 128             case '\n':
 129                 if (*str == '\n') {
 130                     while (*++str == '\n')
 131                         ;
 132                     ans += 10; /* strlen("\n</p>\n<p>\n") */
 133                     break;
 134                 }
 135                 /* fall through */
 136             default:
 137                 ++ans;
 138                 break;
 139         }
 140     }
 141     return ans;
 142 }
 143
 144 static char *
 145 xml_escape(
 146     const char *msg)
 147 {
 148     size_t bl;
 149     char *ans;
 150     char ch;
 151     char *p;
 152
 153     if (!msg)
 154         return NULL;
 155     bl = get_escxml_size(msg);
 156     ans = (char *)malloc(bl + 1);
 157     if (!ans)
 158         diefailure("malloc failed");
 159     p = ans;
 160     while ((ch = *msg++) != 0) {
 161         switch(ch) {
 162             case '<':
 163                 strncpy(p, "&lt;", 4);
 164                 p += 4;
 165                 break;
 166             case '>':
 167                 strncpy(p, "&gt;", 4);
 168                 p += 4;
 169                 break;
 170             case '&':
 171                 strncpy(p, "&amp;", 5);
 172                 p += 5;
 173                 break;
 174             case '\n':
 175                 if (*msg == '\n') {
 176                     while (*++msg == '\n')
 177                         ;
 178                     strncpy(p, "\n</p>\n<p>\n", 10);
 179                     p += 10;
 180                     break;
 181                 }
 182                 /* fall through */
 183             default:
 184                 *p++ = ch;
 185                 break;
 186         }
 187     }
 188     *p = 0;
 189     return ans;
 190 }
 191
 192 static sig_atomic_t piped = 0;
 193 static int wakeup[2] = {-1, -1};
 194 #define wakeup_r (wakeup[0])
 195 #define wakeup_w (wakeup[1])
 196
 197 static void
 198 sigpipe(
 199     int sig)
 200 {
 201     ssize_t ignore;
 202
 203     (void)sig;
 204     piped = 1;
 205     ignore = write(wakeup_w, "!", 1);
 206     (void)ignore;
 207 }
 208
 209 static void
 210 setcloexec(
 211     int fd)
 212 {
 213     if (fcntl(fd, F_SETFD, FD_CLOEXEC))
 214         diefailure("fcntl F_SETFD failed");
 215 }
 216
 217 static void
 218 setnonblock(
 219     int fd)
 220 {
 221     int flags = fcntl(fd, F_GETFL, 0);
 222     if (flags == -1)
 223         diefailure("fcntl F_GETFL failed");
 224     if (fcntl(fd, F_SETFL, flags | O_NONBLOCK) == -1)
 225         diefailure("fcntl F_SETFL O_NONBLOCK failed");
 226 }
 227
 228 static void
 229 setblock(
 230     int fd)
 231 {
 232     int flags = fcntl(fd, F_GETFL, 0);
 233     if (flags == -1)
 234         diefailure("fcntl F_GETFL failed");
 235     if (fcntl(fd, F_GETFL, flags & ~O_NONBLOCK) == -1)
 236         diefailure("fcntl F_SETFL !O_NONBLOCK failed");
 237 }
 238
 239 static ssize_t
 240 writeall(
 241     int fd,
 242     const void *_buf,
 243     size_t nbyte)
 244 {
 245     const char *buf = (const char *)_buf;
 246     size_t off = 0;
 247     ssize_t written;
 248     size_t count = nbyte;
 249
 250     if (!nbyte)
 251         return 0;
 252     do {
 253         written = write(fd, buf+off, count);
 254         if (written > 0) {
 255             count -= (size_t)written;
 256             off += (size_t)written;
 257         }
 258     } while(!piped && count && (written > 0 || errno == EINTR));
 259     if (!count && !piped)
 260         return nbyte;
 261     if (!written)
 262         errno = EIO;
 263     if (piped)
 264         errno = EPIPE;
 265     return -1;
 266 }
 267
 268 #define GENERIC_MESSAGE \
 269 "The requested service is temporarily unavailable, please try again later."
 270
 271 void
 272 send_failure(
 273     int err, const char *msg)
 274 {
 275     if (err_on_failure)
 276         exit(err);
 277     printf("%s\r\n", "Status: 503 Service Temporarily Unavailable");
 278     printf("%s\r\n", "Expires: Fri, 01 Jan 1980 00:00:00 GMT");
 279     printf("%s\r\n", "Pragma: no-cache");
 280     printf("%s\r\n", "Cache-Control: no-cache,max-age=0,must-revalidate");
 281     printf("%s\r\n", "Content-Type: text/html; charset=utf-8");
 282     printf("%s\r\n", "");
 283     printf("%s\n%s\n%s\n", "<p>", GENERIC_MESSAGE, "</p>");
 284     if (msg && *msg)
 285         printf("%s\n%s\n%s\n", "<p>", msg, "</p>");
 286     fflush(stdout);
 287     exit(0);
 288 }
 289
 290 int
 291 main(
 292     int argc,
 293     char *argv[])
 294 {
 295     int ch;
 296     const char *classname = NULL;
 297     const char *desc = NULL;
 298     char *escmsg = NULL;
 299     char **args;
 300     union {
 301         struct sockaddr sa;
 302         struct sockaddr_un sun;
 303     } su;
 304     int fd, fd2;
 305     char request[1024];
 306
 307     while ((ch = getopt(argc, argv, "htpec:d:m:")) != -1) {
 308         switch(ch) {
 309             case 'h':
 310                 printf("%s", USAGE);
 311                 return 0;
 312                 break;
 313             case 't':
 314                 ok_connect_fail = 0;
 315                 break;
 316             case 'p':
 317                 ok_request_fail = 1;
 318                 break;
 319             case 'e':
 320                 err_on_failure = 1;
 321                 break;
 322             case 'c':
 323                 if (!optarg || !*optarg)
 324                     dieusage("-c: invalid class name");
 325                 classname = optarg;
 326                 break;
 327             case 'd':
 328                 if (!optarg || !*optarg)
 329                     dieusage("-c: invalid description");
 330                 desc = optarg;
 331                 break;
 332             case 'm':
 333                 if (!optarg || !*optarg)
 334                     dieusage("-m: invalid message");
 335                 if (escmsg)
 336                     free(escmsg);
 337                 escmsg = xml_escape(optarg);
 338                 break;
 339             case '?':
 340             default:
 341                 dieusage("invalid option");
 342                 break;
 343         }
 344     }
 345     if (!classname || !desc)
 346         dieusage("-c <class> and -d <desc> are required");
 347     if (strlen(classname) > 64)
 348         die("classname must be 64 characters or less in length");
 349     if (strlen(desc) > 512)
 350         die("description must be 512 characters or less in length");
 351     argv += optind;
 352     argc -= optind;
 353     if (!argc || !*argv)
 354         dieusage("no command specified");
 355     args = (char **)malloc((argc + 1) * sizeof(char *));
 356     if (!args)
 357         diefailure("malloc failed");
 358     memcpy(args, argv, argc * sizeof(char *));
 359     args[argc] = NULL;
 360
 361     if (pipe(wakeup))
 362         diefailure("pipe failed");
 363     setcloexec(wakeup_w);
 364     setcloexec(wakeup_r);
 365     setnonblock(wakeup_w);
 366     setnonblock(wakeup_r);
 367
 368     memset(&su.sun, 0, sizeof(su.sun));
 369     su.sun.sun_family = AF_UNIX;
 370     strncpy(su.sun.sun_path, SOCKET_FILE, sizeof(su.sun.sun_path)-1);
 371     fd2 = socket(AF_UNIX, SOCK_STREAM, 0);
 372     if (fd2 < 0)
 373         diefailure("socket failed");
 374     fd = fcntl(fd2, F_DUPFD, 16);
 375     if (fd < 0)
 376         diefailure("fcntl F_DUPFD failed");
 377     close(fd2);
 378     if (connect(fd, &su.sa, sizeof(su.sun))) {
 379         if (errno == ECONNREFUSED || errno == EACCES ||
 380             errno == ENOENT || errno == ENOTDIR) {
 381
 382             if (!ok_connect_fail)
 383                 send_failure(10, escmsg);
 384             else
 385                 close(fd), fd=-1;
 386         } else {
 387             diefailure("connect failed");
 388         }
 389     }
 390
 391     signal(SIGPIPE, sigpipe);
 392
 393     if (fd >= 0) {
 394         snprintf(request, sizeof(request), "throttle %u %s %s\n",
 395             (unsigned)getpid(), classname, desc);
 396         if (writeall(fd, request, strlen(request)) == -1) {
 397             if (!ok_request_fail)
 398                 send_failure(11, escmsg);
 399             else
 400                 close(fd), fd=-1;
 401         }
 402     }
 403
 404     if (fd >= 0) {
 405         fd_set fds;
 406         int nfds = wakeup_r + 1;
 407         size_t off = 0;
 408         int requesterr = 0;
 409         int proceed = 0;
 410         int throttled = 0;
 411
 412         if (fd > wakeup_r)
 413             nfds = fd + 1;
 414         setnonblock(fd);
 415         FD_ZERO(&fds);
 416         FD_SET(wakeup_r, &fds);
 417         FD_SET(fd, &fds);
 418         while (!piped && !requesterr && !proceed && !throttled) {
 419             fd_set rset, eset;
 420             struct timeval tv;
 421             ssize_t readcnt;
 422
 423             FD_COPY(&fds, &rset);
 424             FD_COPY(&fds, &eset);
 425             tv.tv_sec = 30;
 426             tv.tv_usec = 0;
 427             select(nfds, &rset, NULL, &eset, &tv);
 428             do {
 429                 readcnt = read(fd, request + off, sizeof(request) - 1 - off);
 430                 if (readcnt > 0) {
 431                     char *nl;
 432                     off += (size_t)readcnt;
 433                     request[off] = 0;
 434                     if ((nl = strchr(request, '\n')) != NULL) {
 435                         if (!strcmp(request, "proceed\n"))
 436                             proceed = 1;
 437                         else if (nl > request && isalphach(request[0]))
 438                             throttled = 1;
 439                         else
 440                             requesterr = 1;
 441                     }
 442                 }
 443             } while (readcnt > 0 && !proceed && !throttled && !requesterr);
 444             if ((!readcnt && !proceed && !throttled) ||
 445                 (readcnt == -1 && errno != EINTR && errno != EAGAIN) ||
 446                 (off + 1 >= sizeof(request) && !proceed && !throttled)) {
 447
 448                 requesterr = 1;
 449             }
 450             setblock(fd);
 451             if (writeall(fd, "keepalive\n", 10) == -1)
 452                 piped = 1;
 453             setnonblock(fd);
 454         }
 455         if (requesterr && !ok_request_fail)
 456             send_failure(20, escmsg);
 457         else if (throttled)
 458             send_failure(21, escmsg);
 459     }
 460
 461     close(wakeup_r);
 462     close(wakeup_w);
 463     signal(SIGPIPE, SIG_DFL);
 464
 465     execvp(args[0], args);
 466     send_failure(3, escmsg);
 467     return EXIT_FAILURE;
 468 }