Merge pull request #83 from grnhse/translate-sigterm

Add '--signal' option to replace SIGTERM

README.md

@@ -106,6 +106,23 @@ completely transparent; you can even string multiple together (like `dumb-init
 dumb-init echo 'oh, hi'`).
 
 
+### Signal rewriting
+
+dumb-init allows rewriting incoming signals before proxying them. This is
+useful in cases where you have a Docker supervisor (like Mesos or Kubernates)
+which always sends a standard signal (e.g. SIGTERM). Some apps require a
+different stop signal in order to do graceful cleanup.
+
+For example, to rewrite the signal SIGTERM (number 15) to SIGQUIT (number 3),
+just add `--rewrite 15:3` on the command line.
+
+One caveat with this feature: for job control signals (SIGTSTP, SIGTTIN,
+SIGTTOU), dumb-init will always suspend itself after receiving the signal, even
+if you rewrite it to something else. Additionally, if in setsid mode, dumb-init
+will always forward SIGSTOP instead, since the original signals have no effect
+unless the child has handlers for them.
+
+
 ## Installing inside Docker containers
 
 You have a few options for using `dumb-init`:

dumb-init.c

@@ -30,11 +30,28 @@
     } \
 } while (0)
 
+// Signals we care about are numbered from 1 to 31, inclusive.
+// (32 and above are real-time signals.)
+#define MAXSIG 31
+
+// Indices are one-indexed (signal 1 is at index 1). Index zero is unused.
+int signal_rewrite[MAXSIG + 1] = {0};
+
 pid_t child_pid = -1;
 char debug = 0;
 char use_setsid = 1;
 
+int translate_signal(int signum) {
+    if (signum <= 0 || signum > MAXSIG) {
+        return signum;
+    } else {
+        int translated = signal_rewrite[signum];
+        return translated == 0 ? signum : translated;
+    }
+}
+
 void forward_signal(int signum) {
+    signum = translate_signal(signum);
     kill(use_setsid ? -child_pid : child_pid, signum);
     DEBUG("Forwarded signal %d to children.\n", signum);
 }
@@ -125,6 +142,7 @@ void print_help(char *argv[]) {
         "   -c, --single-child   Run in single-child mode.\n"
         "                        In this mode, signals are only proxied to the\n"
         "                        direct child and not any of its descendants.\n"
+        "   -r, --rewrite s:r    Rewrite received signal s to new signal r before proxying.\n"
         "   -v, --verbose        Print debugging information to stderr.\n"
         "   -h, --help           Print this help message and exit.\n"
         "   -V, --version        Print the current version and exit.\n"
@@ -135,16 +153,41 @@ void print_help(char *argv[]) {
     );
 }
 
+void print_rewrite_signum_help() {
+    fprintf(
+        stderr,
+        "Usage: -r option takes <signum>:<signum>, where <signum> "
+        "is between 1 and %d.\n"
+        "This option can be specified multiple times.\n"
+        "Use --help for full usage.\n",
+        MAXSIG
+    );
+    exit(1);
+}
+
+void parse_rewrite_signum(char *arg) {
+    int signum, replacement;
+    if (
+        sscanf(arg, "%d:%d", &signum, &replacement) == 2 &&
+        (signum >= 1 && signum <= MAXSIG) &&
+        (replacement >= 1 && replacement <= MAXSIG)
+    ) {
+        signal_rewrite[signum] = replacement;
+    } else {
+        print_rewrite_signum_help();
+    }
+}
 
 char **parse_command(int argc, char *argv[]) {
     int opt;
     struct option long_options[] = {
-        {"help",         no_argument, NULL, 'h'},
-        {"single-child", no_argument, NULL, 'c'},
-        {"verbose",      no_argument, NULL, 'v'},
-        {"version",      no_argument, NULL, 'V'},
+        {"help",         no_argument,       NULL, 'h'},
+        {"single-child", no_argument,       NULL, 'c'},
+        {"rewrite",      required_argument, NULL, 'r'},
+        {"verbose",      no_argument,       NULL, 'v'},
+        {"version",      no_argument,       NULL, 'V'},
     };
-    while ((opt = getopt_long(argc, argv, "+hvVc", long_options, NULL)) != -1) {
+    while ((opt = getopt_long(argc, argv, "+hvVcr:", long_options, NULL)) != -1) {
         switch (opt) {
             case 'h':
                 print_help(argv);
@@ -158,6 +201,9 @@ char **parse_command(int argc, char *argv[]) {
             case 'c':
                 use_setsid = 0;
                 break;
+            case 'r':
+                parse_rewrite_signum(optarg);
+                break;
             default:
                 exit(1);
         }

tests/cli_test.py

@@ -40,6 +40,7 @@ def test_help_message(flag, both_debug_modes, both_setsid_modes, current_version
         b'   -c, --single-child   Run in single-child mode.\n'
         b'                        In this mode, signals are only proxied to the\n'
         b'                        direct child and not any of its descendants.\n'
+        b'   -r, --rewrite s:r    Rewrite received signal s to new signal r before proxying.\n'
         b'   -v, --verbose        Print debugging information to stderr.\n'
         b'   -h, --help           Print this help message and exit.\n'
         b'   -V, --version        Print the current version and exit.\n'
@@ -96,3 +97,31 @@ def test_verbose_and_single_child(flag1, flag2):
         ),
         stderr,
     )
+
+
+@pytest.mark.parametrize('extra_args', [
+    ('-r',),
+    ('-r', ''),
+    ('-r', 'herp'),
+    ('-r', 'herp:derp'),
+    ('-r', '15'),
+    ('-r', '15::12'),
+    ('-r', '15:derp'),
+    ('-r', '15:12', '-r'),
+    ('-r', '15:12', '-r', '0'),
+    ('-r', '15:12', '-r', '1:32'),
+])
+@pytest.mark.usefixtures('both_debug_modes', 'both_setsid_modes')
+def test_rewrite_errors(extra_args):
+    proc = Popen(
+        ('dumb-init',) + extra_args + ('echo', 'oh,', 'hi'),
+        stdout=PIPE, stderr=PIPE,
+    )
+    stdout, stderr = proc.communicate()
+    assert proc.returncode == 1
+    assert stderr == (
+        b'Usage: -r option takes <signum>:<signum>, where <signum> '
+        b'is between 1 and 31.\n'
+        b'This option can be specified multiple times.\n'
+        b'Use --help for full usage.\n'
+    )

tests/proxies_signals_test.py

@@ -2,25 +2,84 @@
 import re
 import signal
 import sys
+from contextlib import contextmanager
+from itertools import chain
 from subprocess import PIPE
 from subprocess import Popen
 
+import pytest
+
 from tests.lib.testing import NORMAL_SIGNALS
 from tests.lib.testing import pid_tree
 
 
-def test_proxies_signals(both_debug_modes, both_setsid_modes):
-    """Ensure dumb-init proxies regular signals to its child."""
+@contextmanager
+def _print_signals(args=()):
+    """Start print_signals and return dumb-init process."""
     proc = Popen(
-        ('dumb-init', sys.executable, '-m', 'tests.lib.print_signals'),
+        (
+            ('dumb-init',) +
+            tuple(args) +
+            (sys.executable, '-m', 'tests.lib.print_signals')
+        ),
         stdout=PIPE,
     )
-
     assert re.match(b'^ready \(pid: (?:[0-9]+)\)\n$', proc.stdout.readline())
 
-    for signum in NORMAL_SIGNALS:
-        proc.send_signal(signum)
-        assert proc.stdout.readline() == '{0}\n'.format(signum).encode('ascii')
+    yield proc
 
     for pid in pid_tree(proc.pid):
         os.kill(pid, signal.SIGKILL)
+
+
+@pytest.mark.usefixtures('both_debug_modes', 'both_setsid_modes')
+def test_proxies_signals():
+    """Ensure dumb-init proxies regular signals to its child."""
+    with _print_signals() as proc:
+        for signum in NORMAL_SIGNALS:
+            proc.send_signal(signum)
+            assert proc.stdout.readline() == '{0}\n'.format(signum).encode('ascii')
+
+
+def _rewrite_map_to_args(rewrite_map):
+    return chain.from_iterable(
+        ('-r', '{0}:{1}'.format(src, dst)) for src, dst in rewrite_map.items()
+    )
+
+
+@pytest.mark.parametrize('rewrite_map,sequence,expected', [
+    (
+        {},
+        [signal.SIGTERM, signal.SIGQUIT, signal.SIGCONT, signal.SIGINT],
+        [signal.SIGTERM, signal.SIGQUIT, signal.SIGCONT, signal.SIGINT],
+    ),
+
+    (
+        {signal.SIGTERM: signal.SIGINT},
+        [signal.SIGTERM, signal.SIGQUIT, signal.SIGCONT, signal.SIGINT],
+        [signal.SIGINT, signal.SIGQUIT, signal.SIGCONT, signal.SIGINT],
+    ),
+
+    (
+        {
+            signal.SIGTERM: signal.SIGINT,
+            signal.SIGINT: signal.SIGTERM,
+            signal.SIGQUIT: signal.SIGQUIT,
+        },
+        [signal.SIGTERM, signal.SIGQUIT, signal.SIGCONT, signal.SIGINT],
+        [signal.SIGINT, signal.SIGQUIT, signal.SIGCONT, signal.SIGTERM],
+    ),
+
+    (
+        {1: 31, 31: 1},
+        [1, 31],
+        [31, 1],
+    ),
+])
+@pytest.mark.usefixtures('both_debug_modes', 'both_setsid_modes')
+def test_proxies_signals_with_rewrite(rewrite_map, sequence, expected):
+    """Ensure dumb-init can rewrite signals."""
+    with _print_signals(_rewrite_map_to_args(rewrite_map)) as proc:
+        for send, expect_receive in zip(sequence, expected):
+            proc.send_signal(send)
+            assert proc.stdout.readline() == '{0}\n'.format(expect_receive).encode('ascii')