From 4a59e597398e990a68fe52c1305e4f2c2b4d4a7e Mon Sep 17 00:00:00 2001
From: Jean Boussier <jean.boussier@gmail.com>
Date: Tue, 1 Aug 2023 12:30:59 +0200
Subject: [PATCH] Pitchfork::Info.close_all_ios!

---
 lib/pitchfork/info.rb      | 37 +++++++++++++++++++++++--------------
 lib/pitchfork/worker.rb    |  1 +
 test/unit/test_children.rb | 17 +++++++++++------
 test/unit/test_info.rb     | 37 +++++++++++++++++++++++++++++++++++++
 4 files changed, 72 insertions(+), 20 deletions(-)
 create mode 100644 test/unit/test_info.rb

diff --git a/lib/pitchfork/info.rb b/lib/pitchfork/info.rb
index 85f20cf8..c4a54858 100644
--- a/lib/pitchfork/info.rb
+++ b/lib/pitchfork/info.rb
@@ -12,7 +12,8 @@ class << self
       attr_accessor :workers_count
 
       def keep_io(io)
-        @kept_ios[io] = io if io && !io.to_io.closed?
+        raise ArgumentError, "#{io.inspect} doesn't respond to :to_io" unless io.respond_to?(:to_io)
+        @kept_ios[io] = io
         io
       end
 
@@ -20,23 +21,31 @@ def keep_ios(ios)
         ios.each { |io| keep_io(io) }
       end
 
-      def close_all_fds!
-        ignored_fds = [$stdin.to_i, $stdout.to_i, $stderr.to_i]
+      def close_all_ios!
+        ignored_ios = [$stdin, $stdout, $stderr]
+
         @kept_ios.each_value do |io_like|
-          if io = io_like&.to_io
-            ignored_fds << io.to_i unless io.closed?
-          end
+          ignored_ios << (io_like.is_a?(IO) ? io_like : io_like.to_io)
         end
 
-        all_fds = Dir.children("/dev/fd").map(&:to_i)
-        all_fds -= ignored_fds
+        ObjectSpace.each_object(IO) do |io|
+          closed = begin
+            io.closed?
+          rescue IOError
+            true
+          end
 
-        all_fds.each do |fd|
-          IO.for_fd(fd).close
-        rescue ArgumentError
-          # RubyVM internal file descriptor, leave it alone
-        rescue Errno::EBADF
-          # Likely a race condition
+          if !closed && io.autoclose? && !ignored_ios.include?(io)
+            if io.is_a?(TCPSocket)
+              # If we inherited a TCP Socket, calling #close directly could send FIN or RST.
+              # So we first reopen /dev/null to avoid that.
+              io.reopen(File::NULL)
+            end
+            begin
+              io.close
+            rescue Errno::EBADF
+            end
+          end
         end
       end
 
diff --git a/lib/pitchfork/worker.rb b/lib/pitchfork/worker.rb
index d2b1d341..5884bc57 100644
--- a/lib/pitchfork/worker.rb
+++ b/lib/pitchfork/worker.rb
@@ -208,6 +208,7 @@ def after_fork_in_child
     private
 
     def pipe=(socket)
+      raise ArgumentError, "pipe can't be nil" unless socket
       Info.keep_io(socket)
       @master = MessageSocket.new(socket)
     end
diff --git a/test/unit/test_children.rb b/test/unit/test_children.rb
index 34808f67..85b96df1 100644
--- a/test/unit/test_children.rb
+++ b/test/unit/test_children.rb
@@ -18,19 +18,21 @@ def test_register
     end
 
     def test_message_worker_spawned
+      pipe = IO.pipe.last
       worker = Worker.new(0)
       @children.register(worker)
       assert_predicate @children, :pending_workers?
 
-      @children.update(Message::WorkerSpawned.new(0, 42, 0, nil))
+      @children.update(Message::WorkerSpawned.new(0, 42, 0, pipe))
       refute_predicate @children, :pending_workers?
       assert_equal 42, worker.pid
       assert_equal [worker], @children.workers
     end
 
     def test_message_mold_spawned
+      pipe = IO.pipe.last
       assert_nil @children.mold
-      @children.update(Message::MoldSpawned.new(nil, 42, 1, nil))
+      @children.update(Message::MoldSpawned.new(nil, 42, 1, pipe))
 
       assert_nil @children.mold
       assert_equal 0, @children.molds.size
@@ -40,8 +42,9 @@ def test_message_mold_spawned
     end
 
     def test_message_mold_ready
+      pipe = IO.pipe.last
       assert_nil @children.mold
-      @children.update(Message::MoldSpawned.new(nil, 42, 1, nil))
+      @children.update(Message::MoldSpawned.new(nil, 42, 1, pipe))
       mold = @children.update(Message::MoldReady.new(nil, 42, 1))
 
       assert_equal mold, @children.mold
@@ -52,26 +55,28 @@ def test_message_mold_ready
     end
 
     def test_reap_worker
+      pipe = IO.pipe.last
       worker = Worker.new(0)
       @children.register(worker)
       assert_predicate @children, :pending_workers?
 
-      @children.update(Message::WorkerSpawned.new(0, 42, 0, nil))
+      @children.update(Message::WorkerSpawned.new(0, 42, 0, pipe))
 
       assert_equal worker, @children.reap(worker.pid)
       assert_nil @children.reap(worker.pid)
     end
 
     def test_reap_old_molds
+      pipe = IO.pipe.last
       assert_nil @children.mold
-      @children.update(Message::MoldSpawned.new(nil, 24, 0, nil))
+      @children.update(Message::MoldSpawned.new(nil, 24, 0, pipe))
       @children.update(Message::MoldReady.new(nil, 24, 0))
 
       first_mold = @children.mold
       refute_nil first_mold
       assert_equal 24, first_mold.pid
 
-      @children.update(Message::MoldSpawned.new(nil, 42, 1, nil))
+      @children.update(Message::MoldSpawned.new(nil, 42, 1, pipe))
       @children.update(Message::MoldReady.new(nil, 42, 1))
       second_mold = @children.mold
       refute_nil second_mold
diff --git a/test/unit/test_info.rb b/test/unit/test_info.rb
new file mode 100644
index 00000000..4f9fb50d
--- /dev/null
+++ b/test/unit/test_info.rb
@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+require 'test_helper'
+
+class TestInfo < Pitchfork::Test
+  def test_close_all_ios_except_marked_ones
+    r, w = IO.pipe
+
+    Pitchfork::Info.keep_io(w)
+
+    pid = Process.fork do
+      Pitchfork::Info.close_all_ios!
+
+      w.write(Marshal.dump([
+        $stdin.closed?,
+        $stdout.closed?,
+        $stderr.closed?,
+        r.closed?,
+        w.closed?
+      ]))
+      Process.exit!(0)
+    end
+
+    _, status = Process.wait2(pid)
+    assert_predicate status, :success?
+
+    info = Marshal.load(r)
+
+    assert_equal([
+      false, # $stdin
+      false, # $stdout
+      false, # $stderr
+      true, # r
+      false, # w
+    ], info)
+  end
+end