projlist: improve text browser project list formatting
[girocco.git] / src / throttle.c
blobe43ebc9b230a5ebff463f2358e373bace1262da5
1 /*
3 throttle.c -- throttle utility to make use of taskd.pl throttle services
4 Copyright (C) 2015 Kyle J. McKay. All rights reserved.
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.
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.
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.
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.
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.
34 #ifndef SOCKET_FILE
35 #error SOCKET_FILE must be defined to full path to taskd.pl socket
36 #endif
38 #ifndef _POSIX_C_SOURCE
39 #define _POSIX_C_SOURCE 200809
40 #endif
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>
56 #ifndef FD_COPY
57 #define FD_COPY(src,dst) ((*dst)=(*src))
58 #endif
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"
73 static int ok_connect_fail = 1;
74 static int ok_request_fail = 0;
75 static int err_on_failure = 0;
77 static void
78 die(
79 const char *msg)
81 fprintf(stderr, "fatal: %s\n", msg ? msg : "error");
82 exit(EXIT_FAILURE);
85 static void
86 diefailure(
87 const char *msg)
89 fprintf(stderr, "fatal: %s: %s\n", msg ? msg : "failed", strerror(errno));
90 exit(EXIT_FAILURE);
93 static void
94 dieusage(
95 const char *msg)
97 fprintf(stderr, "fatal: %s\n", msg ? msg : "error");
98 fprintf(stderr, "%s", USAGE);
99 exit(EXIT_FAILURE);
102 static int
103 isalphach(
104 char ch)
106 return ('A' <= ch && ch <= 'Z') ||
107 ('a' <= ch && ch <= 'z');
110 static size_t
111 get_escxml_size(
112 const char *str)
114 size_t ans = 0;
115 char ch;
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')
132 ans += 10; /* strlen("\n</p>\n<p>\n") */
133 break;
135 /* fall through */
136 default:
137 ++ans;
138 break;
141 return ans;
144 static char *
145 xml_escape(
146 const char *msg)
148 size_t bl;
149 char *ans;
150 char ch;
151 char *p;
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')
178 strncpy(p, "\n</p>\n<p>\n", 10);
179 p += 10;
180 break;
182 /* fall through */
183 default:
184 *p++ = ch;
185 break;
188 *p = 0;
189 return ans;
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])
197 static void
198 sigpipe(
199 int sig)
201 ssize_t ignore;
203 (void)sig;
204 piped = 1;
205 ignore = write(wakeup_w, "!", 1);
206 (void)ignore;
209 static void
210 setcloexec(
211 int fd)
213 if (fcntl(fd, F_SETFD, FD_CLOEXEC))
214 diefailure("fcntl F_SETFD failed");
217 static void
218 setnonblock(
219 int fd)
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");
228 static void
229 setblock(
230 int fd)
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");
239 static ssize_t
240 writeall(
241 int fd,
242 const void *_buf,
243 size_t nbyte)
245 const char *buf = (const char *)_buf;
246 size_t off = 0;
247 ssize_t written;
248 size_t count = nbyte;
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;
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;
268 #define GENERIC_MESSAGE \
269 "The requested service is temporarily unavailable, please try again later."
271 void
272 send_failure(
273 int err, const char *msg)
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);
291 main(
292 int argc,
293 char *argv[])
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];
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;
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;
361 if (pipe(wakeup))
362 diefailure("pipe failed");
363 setcloexec(wakeup_w);
364 setcloexec(wakeup_r);
365 setnonblock(wakeup_w);
366 setnonblock(wakeup_r);
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) {
382 if (!ok_connect_fail)
383 send_failure(10, escmsg);
384 else
385 close(fd), fd=-1;
386 } else {
387 diefailure("connect failed");
391 signal(SIGPIPE, sigpipe);
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;
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;
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;
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;
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)) {
448 requesterr = 1;
450 setblock(fd);
451 if (writeall(fd, "keepalive\n", 10) == -1)
452 piped = 1;
453 setnonblock(fd);
455 if (requesterr && !ok_request_fail)
456 send_failure(20, escmsg);
457 else if (throttled)
458 send_failure(21, escmsg);
461 close(wakeup_r);
462 close(wakeup_w);
463 signal(SIGPIPE, SIG_DFL);
465 execvp(args[0], args);
466 send_failure(3, escmsg);
467 return EXIT_FAILURE;