Merge pull request #3223 from cowlicks/signal-handling

Fix Unix signal handling in certbot.error_handler.ErrorHandler [minor revision]

certbot/error_handler.py

@@ -5,6 +5,7 @@ import os
 import signal
 import traceback
 
+from certbot import errors
 
 logger = logging.getLogger(__name__)
 
@@ -20,43 +21,62 @@ _SIGNALS = ([signal.SIGTERM] if os.name == "nt" else
 
 
 class ErrorHandler(object):
-    """Registers functions to be called if an exception or signal occurs.
+    """Context manager for running code that must be cleaned up on failure.
 
-    This class allows you to register functions that will be called when
-    an exception (excluding SystemExit) or signal is encountered. The
-    class works best as a context manager. For example:
+    The context manager allows you to register functions that will be called
+    when an exception (excluding SystemExit) or signal is encountered. Usage:
 
-    with ErrorHandler(cleanup_func):
+    handler = ErrorHandler(cleanup1_func, *cleanup1_args, **cleanup1_kwargs)
+    handler.register(cleanup2_func, *cleanup2_args, **cleanup2_kwargs)
+
+    with handler:
         do_something()
 
-    If an exception is raised out of do_something, cleanup_func will be
-    called. The exception is not caught by the ErrorHandler. Similarly,
-    if a signal is encountered, cleanup_func is called followed by the
-    previously registered signal handler.
+    Or for one cleanup function:
 
-    Every registered function is attempted to be run to completion
-    exactly once. If a registered function raises an exception, it is
-    logged and the next function is called. If a (different) handled
-    signal occurs while calling a registered function, it is attempted
-    to be called again by the next signal handler.
+    with ErrorHandler(func, args, kwargs):
+        do_something()
+
+    If an exception is raised out of do_something, the cleanup functions will
+    be called in last in first out order. Then the exception is raised.
+    Similarly, if a signal is encountered, the cleanup functions are called
+    followed by the previously received signal handler.
+
+    Each registered cleanup function is called exactly once. If a registered
+    function raises an exception, it is logged and the next function is called.
+    Signals received while the registered functions are executing are
+    deferred until they finish.
 
     """
     def __init__(self, func=None, *args, **kwargs):
+        self.body_executed = False
         self.funcs = []
         self.prev_handlers = {}
+        self.received_signals = []
         if func is not None:
             self.register(func, *args, **kwargs)
 
     def __enter__(self):
-        self.set_signal_handlers()
+        self.body_executed = False
+        self._set_signal_handlers()
 
     def __exit__(self, exec_type, exec_value, trace):
+        self.body_executed = True
+        retval = False
         # SystemExit is ignored to properly handle forks that don't exec
-        if exec_type not in (None, SystemExit):
+        if exec_type in (None, SystemExit):
+            return retval
+        elif exec_type is errors.SignalExit:
+            logger.debug("Encountered signals: %s", self.received_signals)
+            retval = True
+        else:
             logger.debug("Encountered exception:\n%s", "".join(
                 traceback.format_exception(exec_type, exec_value, trace)))
-            self.call_registered()
-        self.reset_signal_handlers()
+
+        self._call_registered()
+        self._reset_signal_handlers()
+        self._call_signals()
+        return retval
 
     def register(self, func, *args, **kwargs):
         """Sets func to be called with *args and **kwargs during cleanup
@@ -66,7 +86,7 @@ class ErrorHandler(object):
         """
         self.funcs.append(functools.partial(func, *args, **kwargs))
 
-    def call_registered(self):
+    def _call_registered(self):
         """Calls all registered functions"""
         logger.debug("Calling registered functions")
         while self.funcs:
@@ -77,7 +97,7 @@ class ErrorHandler(object):
                 logger.exception(error)
             self.funcs.pop()
 
-    def set_signal_handlers(self):
+    def _set_signal_handlers(self):
         """Sets signal handlers for signals in _SIGNALS."""
         for signum in _SIGNALS:
             prev_handler = signal.getsignal(signum)
@@ -86,19 +106,27 @@ class ErrorHandler(object):
                 self.prev_handlers[signum] = prev_handler
                 signal.signal(signum, self._signal_handler)
 
-    def reset_signal_handlers(self):
+    def _reset_signal_handlers(self):
         """Resets signal handlers for signals in _SIGNALS."""
         for signum in self.prev_handlers:
             signal.signal(signum, self.prev_handlers[signum])
         self.prev_handlers.clear()
 
     def _signal_handler(self, signum, unused_frame):
-        """Calls registered functions and the previous signal handler.
+        """Replacement function for handling recieved signals.
+
+        Store the recieved signal. If we are executing the code block in
+        the body of the context manager, stop by raising signal exit.
 
         :param int signum: number of current signal
 
         """
-        logger.debug("Singal %s encountered", signum)
-        self.call_registered()
-        signal.signal(signum, self.prev_handlers[signum])
-        os.kill(os.getpid(), signum)
+        self.received_signals.append(signum)
+        if not self.body_executed:
+            raise errors.SignalExit
+
+    def _call_signals(self):
+        """Finally call the deferred signals."""
+        for signum in self.received_signals:
+            logger.debug("Calling signal %s", signum)
+            os.kill(os.getpid(), signum)

certbot/errors.py

@@ -29,6 +29,10 @@ class HookCommandNotFound(Error):
     """Failed to find a hook command in the PATH."""
 
 
+class SignalExit(Error):
+    """A Unix signal was recieved while in the ErrorHandler context manager."""
+
+
 # Auth Handler Errors
 class AuthorizationError(Error):
     """Authorization error."""

certbot/tests/error_handler_test.py

@@ -1,10 +1,38 @@
 """Tests for certbot.error_handler."""
+import contextlib
+import os
 import signal
 import sys
 import unittest
 
 import mock
 
+def get_signals(signums):
+    """Get the handlers for an iterable of signums."""
+    return dict((s, signal.getsignal(s)) for s in signums)
+
+
+def set_signals(sig_handler_dict):
+    """Set the signal (keys) with the handler (values) from the input dict."""
+    for s, h in sig_handler_dict.items():
+        signal.signal(s, h)
+
+
+@contextlib.contextmanager
+def signal_receiver(signums):
+    """Context manager to catch signals"""
+    signals = []
+    prev_handlers = {}
+    prev_handlers = get_signals(signums)
+    set_signals(dict((s, lambda s, _: signals.append(s)) for s in signums))
+    yield signals
+    set_signals(prev_handlers)
+
+
+def send_signal(signum):
+    """Send the given signal"""
+    os.kill(os.getpid(), signum)
+
 
 class ErrorHandlerTest(unittest.TestCase):
     """Tests for certbot.error_handler."""
@@ -22,6 +50,38 @@ class ErrorHandlerTest(unittest.TestCase):
         self.signals = error_handler._SIGNALS
 
     def test_context_manager(self):
+        exception_raised = False
+        try:
+            with self.handler:
+                raise ValueError
+        except ValueError:
+            exception_raised = True
+
+        self.assertTrue(exception_raised)
+        self.init_func.assert_called_once_with(*self.init_args,
+                                               **self.init_kwargs)
+
+    def test_context_manager_with_signal(self):
+        init_signals = get_signals(self.signals)
+        with signal_receiver(self.signals) as signals_received:
+            with self.handler:
+                should_be_42 = 42
+                send_signal(self.signals[0])
+                should_be_42 *= 10
+
+        # check exectuion stoped when the signal was sent
+        self.assertEqual(42, should_be_42)
+        # assert signals were caught
+        self.assertEqual([self.signals[0]], signals_received)
+        # assert the error handling function was just called once
+        self.init_func.assert_called_once_with(*self.init_args,
+                                               **self.init_kwargs)
+        for signum in self.signals:
+            self.assertEqual(init_signals[signum], signal.getsignal(signum))
+
+    def test_bad_recovery(self):
+        bad_func = mock.MagicMock(side_effect=[ValueError])
+        self.handler.register(bad_func)
         try:
             with self.handler:
                 raise ValueError
@@ -29,31 +89,17 @@ class ErrorHandlerTest(unittest.TestCase):
             pass
         self.init_func.assert_called_once_with(*self.init_args,
                                                **self.init_kwargs)
+        bad_func.assert_called_once_with()
 
-    @mock.patch('certbot.error_handler.os')
-    @mock.patch('certbot.error_handler.signal')
-    def test_signal_handler(self, mock_signal, mock_os):
-        # pylint: disable=protected-access
-        mock_signal.getsignal.return_value = signal.SIG_DFL
-        self.handler.set_signal_handlers()
-        signal_handler = self.handler._signal_handler
-        for signum in self.signals:
-            mock_signal.signal.assert_any_call(signum, signal_handler)
-
-        signum = self.signals[0]
-        signal_handler(signum, None)
-        self.init_func.assert_called_once_with(*self.init_args,
-                                               **self.init_kwargs)
-        mock_os.kill.assert_called_once_with(mock_os.getpid(), signum)
-
-        self.handler.reset_signal_handlers()
-        for signum in self.signals:
-            mock_signal.signal.assert_any_call(signum, signal.SIG_DFL)
-
-    def test_bad_recovery(self):
-        bad_func = mock.MagicMock(side_effect=[ValueError])
+    def test_bad_recovery_with_signal(self):
+        sig1 = self.signals[0]
+        sig2 = self.signals[-1]
+        bad_func = mock.MagicMock(side_effect=lambda: send_signal(sig1))
         self.handler.register(bad_func)
-        self.handler.call_registered()
+        with signal_receiver(self.signals) as signals_received:
+            with self.handler:
+                send_signal(sig2)
+        self.assertEqual([sig2, sig1], signals_received)
         self.init_func.assert_called_once_with(*self.init_args,
                                                **self.init_kwargs)
         bad_func.assert_called_once_with()