trap_detection.rb

# 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