KJTsanaktsidis · November 9, 2023 23:29 · Nov 9, 2023
diff --git a/trap_detection.rb b/trap_detection.rb
@@ -0,0 +1,74 @@
+# frozen_string_literal: true
+
+module TrapDetection
+  def trap(signal, action = nil, &block)
+    # A block might or might not be given, and action might be absent, a string, or
+    # a callable. It's actually legal to pass both action and block to Signal.trap, but
+    # in that case the block is ignored.
+    if action.respond_to?(:call)
+      action = Thunk.new(action)
+    end
+    if action.nil? && block_given?
+      action = Thunk.new(block)
+    end
+
+    # The Signal.trap method really does differentiate between being called with
+    # a nil action and no action argument at all.
+    ret = super(signal, *[action].compact)
+
+    # trap returns the old handler; unwrap it if it's ours.
+    ret = ret.wrapped_handler if ret.is_a?(Thunk)
+    ret
+  end
+
+  def trap_context?
+    !!Thread.current.thread_variable_get(:trap_detection_signal_list)&.any?
+  end
+
+  class Thunk
+    def initialize(wrapped_handler)
+      @wrapped_handler = wrapped_handler
+    end
+
+    attr_reader :wrapped_handler
+
+    def call(signo)
+      # This must be async-signal-safe - we can't allow another signal to interrupt us
+      # whilst we're getting/setting the signal_list.
+      previous_trap_handler = nil
+      Thread.handle_interrupt(Object => :never) do
+        previous_trap_handler = Thread.current.thread_variable_get(:current_trap_handler)
+        Thread.current.thread_variable_set(:current_trap_handler, @wrapped_handler)
+      end
+      begin
+        Signal._trap(signo)
+      ensure
+        Thread.current.thread_variable_set(:current_trap_handler, previous_trap_handler)
+      end
+    end
+  end
+
+  module TrapExecutionWrapper
+    def _trap(signo)
+      Thread.current.thread_variable_get(:current_trap_handler).call(signo)
+    end
+  end
+end
+
+Signal.singleton_class.prepend TrapDetection
+Signal.singleton_class.include TrapDetection::TrapExecutionWrapper
+Kernel.prepend TrapDetection
+
+# Re-register all the signal handlers so they can be wrapped.
+# The _only_ way to get the signal handlers is to register a new handler, which returns
+# the old one. So, register a new "handler", and then re-register the default one (with wrapping
+# this time). Wrap the whole thing in handle_interrupt so that no signal can arrive whilst we have
+# registgered the default handler
+Thread.handle_interrupt(Object => :never) do
+  Signal.list.keys.each do |signame|
+    existing_handler = Signal.trap(signame, 'SIG_DFL')
+    Signal.trap(signame, existing_handler)
+  rescue ArgumentError, Errno::EINVAL
+    # Some signals cannot be trapped; that's fine.
+  end
+end
No results found