Fix relay scheduler cleanup between runs

lib/rex/io/fiber_scheduler.rb

@@ -15,29 +15,25 @@ def initialize
       @ready = []
       @pending = []
       @blocking = 0
-      @urgent = Rex::Compat.pipe
+      @urgent = nil
       @mutex = Mutex.new
     end
 
     # This allows the fiber to be scheduled in the #run thread from another thread
     def schedule_fiber(&block)
       @mutex.synchronize do
+        @urgent ||= Rex::Compat.pipe
         @pending << block
+        # Wake up the scheduler while holding the same lock used by #reset! so
+        # the signaling pipe cannot be replaced between lookup and write.
+        @urgent.last.write_nonblock('.', exception: false)
       end
-      # Wake up the scheduler
-      @urgent.last.write_nonblock('.') rescue nil
     end
 
     def run
       while @readable.any? or @writable.any? or @waiting.any? or @blocking.positive? or @ready.any? or @pending.any?
         # Start any pending fibers
-        pending_blocks = []
-        @mutex.synchronize do
-          pending_blocks = @pending.dup
-          @pending.clear
-        end
-
-        pending_blocks.each do |block|
+        while (block = @mutex.synchronize { @pending.shift })
           fiber = Fiber.new(blocking: false, &block)
           fiber.resume
         end
@@ -75,8 +71,7 @@ def run
           end
         end
 
-        ready, @ready = @ready, []
-        ready.each do |fiber|
+        while (fiber = @mutex.synchronize { @ready.shift })
           fiber.resume
         end
       end
@@ -118,19 +113,40 @@ def block(blocker, timeout = nil)
     end
 
     def unblock(blocker, fiber)
-      @ready << fiber
-      io = @urgent.last
-      io.write_nonblock('.')
+      @mutex.synchronize do
+        return false unless @urgent
+
+        @ready << fiber
+        @urgent.last.write_nonblock('.', exception: false)
+      end
     end
 
     def close
-      run
-      @urgent.each(&:close)
-      @urgent = nil
+      first_error = nil
+
+      loop do
+        begin
+          run
+          break
+        rescue StandardError => e
+          first_error ||= e
+          break unless work_pending?
+        end
+      end
+
+      raise first_error if first_error
+    ensure
+      close_urgent_pipe
     end
 
     def closed?
-      @urgent.nil?
+      @mutex.synchronize { @urgent.nil? }
+    end
+
+    def work_pending?
+      @mutex.synchronize do
+        @readable.any? || @writable.any? || @waiting.any? || @blocking.positive? || @ready.any? || @pending.any?
+      end
     end
 
     def reset!
@@ -141,7 +157,7 @@ def reset!
         @ready.clear
         @pending.clear
         @blocking = 0
-        @urgent = Rex::Compat.pipe
+        close_urgent_pipe_without_lock
       end
     end
 
@@ -153,6 +169,15 @@ def fiber(&block)
 
     private
 
+    def close_urgent_pipe
+      @mutex.synchronize { close_urgent_pipe_without_lock }
+    end
+
+    def close_urgent_pipe_without_lock
+      @urgent&.each { |io| io.close unless io.closed? }
+      @urgent = nil
+    end
+
     def current_time
       Process.clock_gettime(Process::CLOCK_MONOTONIC)
     end
@@ -165,4 +190,4 @@ def cleanup_closed_ios
   end
 
 end
-end
+end

lib/rex/io/relay_manager.rb

@@ -11,6 +11,10 @@ class RelayManager
   def initialize
     @thread = nil
     @scheduler = FiberScheduler.new
+    @mutex = Mutex.new
+    @teardown_condition = ConditionVariable.new
+    @tearing_down = false
+    @running = false
   end
 
   # Add a IO relay to the manager. This will start relaying data from the source
@@ -26,11 +30,24 @@ def initialize
   # @param name [String] A human-friendly name for the relay used in debug output.
   # @param on_exit [#call] A callback to be invoked when sink can no longer be read from.
   def add_relay(sock, sink: nil, name: nil, on_exit: nil)
-    @scheduler.schedule_fiber do
-      relay_fiber(sock, sink, name, on_exit: on_exit)
+    @mutex.synchronize do
+      # An external caller that arrives during worker teardown waits until the
+      # scheduler is reusable. A callback on the worker itself must remain
+      # nonblocking; the scheduler's final run will process its accepted work.
+      while @tearing_down && Thread.current != @thread
+        @teardown_condition.wait(@mutex)
+      end
+
+      start_gate = start unless running?
+
+      begin
+        @scheduler.schedule_fiber do
+          relay_fiber(sock, sink, name, on_exit: on_exit)
+        end
+      ensure
+        start_gate << true if start_gate
+      end
     end
-
-    start unless running?
   end
 
   # Write all data to the specified IO. This is intended to be used in scenarios
@@ -51,26 +68,66 @@ def self.io_write_all(io, data)
   private
 
   def running?
-    @thread && @thread.alive?
+    @running
   end
 
   def start
-    return false if running?
+    return nil if running?
 
-    @thread = Thread.new { run }
-    true
+    start_gate = Queue.new
+    @running = true
+    @thread = Thread.new do
+      start_gate.pop
+      run
+    end
+    start_gate
+  rescue StandardError
+    @running = false
+    raise
   end
 
   def run
+    teardown_started = false
     old_scheduler = Fiber.scheduler
     # A fiber scheduler can be set per-thread
     Fiber.set_scheduler(@scheduler)
 
-    # Run the scheduler (blocks here)
-    @scheduler.run
+    loop do
+      # Run the scheduler (blocks here)
+      @scheduler.run
+
+      # #add_relay schedules and tests @running while holding this same lock.
+      # If it won the race after the scheduler observed an empty queue, process
+      # its accepted work before committing this worker to teardown.
+      teardown_started = @mutex.synchronize do
+        if @scheduler.work_pending?
+          false
+        else
+          @tearing_down = true
+          true
+        end
+      end
+      break if teardown_started
+    end
   ensure
-    Fiber.set_scheduler(old_scheduler)
-    @scheduler.reset!
+    # On an exceptional exit, commit to teardown before removing the scheduler.
+    # External submissions then wait; callbacks on this worker remain able to
+    # schedule work for the scheduler's final run.
+    @mutex.synchronize { @tearing_down = true } unless teardown_started
+
+    begin
+      Fiber.set_scheduler(old_scheduler)
+    ensure
+      begin
+        @scheduler.reset!
+      ensure
+        @mutex.synchronize do
+          @running = false
+          @tearing_down = false
+          @teardown_condition.broadcast
+        end
+      end
+    end
   end
 
   def relay_fiber(sock, sink, name, on_exit: nil)
@@ -111,4 +168,4 @@ def write_to_sink(sink, data)
 end
 
 end
-end
+end