Forward the local registry port over the OS ssh client

lib/kamal/cli/build/port_forwarding.rb

@@ -1,16 +1,25 @@
-require "concurrent/atomic/count_down_latch"
+require "open3"
+require "timeout"
 
 class Kamal::Cli::Build::PortForwarding
+  READY_TOKEN = "kamal-port-forward-ready"
+  READY_TIMEOUT = 30
+  TEARDOWN_GRACE = 5
+
   attr_reader :hosts, :port, :ssh_options
 
   def initialize(hosts, port, **ssh_options)
     @hosts = hosts
     @port = port
     @ssh_options = ssh_options
+
+    if Array(ssh_options[:key_data]).any?
+      raise "ssh key_data is not supported when forwarding the local registry port; use keys or an ssh agent instead (key_data is deprecated and will be removed in Kamal 3)"
+    end
   end
 
   def forward
-    @done = false
+    @children = []
     forward_ports
 
     yield
@@ -19,48 +28,145 @@ def forward
   end
 
   private
-    def stop
-      @done = true
-      @threads.to_a.each(&:join)
+    # net-ssh's in-process reverse forwarding deadlocks when a large amount of
+    # data flows back through the tunnel — e.g. pulling an image with a layer
+    # bigger than ~27 MiB from the local registry. The single ssh.loop pump
+    # starves while the main thread drives the pull, the forwarded channel's
+    # window is exhausted, and the transfer wedges. Carry the tunnel over the
+    # OS ssh client instead, which handles bulk transfer reliably. See #1886.
+    def forward_ports
+      hosts.each { |host| @children << start_forward(host) }
+
+      # The tunnels establish in parallel once spawned, so READY_TIMEOUT is a
+      # single deadline shared across all hosts, not a per-host allowance.
+      deadline = Process.clock_gettime(Process::CLOCK_MONOTONIC) + READY_TIMEOUT
+      @children.each { |child| wait_until_ready(child, deadline) }
     end
 
-    def forward_ports
-      ready = Concurrent::CountDownLatch.new(hosts.size)
-
-      @threads = hosts.map do |host|
-        Thread.new do
-          begin
-            Net::SSH.start(host, ssh_options[:user], **ssh_options.except(:user)) do |ssh|
-              ssh.forward.remote(port, "localhost", port, "127.0.0.1") do |remote_port, bind_address|
-                if remote_port == :error
-                  raise "Failed to establish port forward on #{host}"
-                else
-                  ready.count_down
-                end
-              end
-
-              ssh.loop(0.1) do
-                if @done
-                  ssh.forward.cancel_remote(port, "127.0.0.1")
-                  break
-                else
-                  true
-                end
-              end
-            end
-          rescue Exception => e
-            error "Error setting up port forwarding to #{host}: #{e.class}: #{e.message}"
-            error e.backtrace.join("\n")
-
-            raise
-          end
+    def start_forward(host)
+      stdin, output, wait_thread = Open3.popen2e(*ssh_command(host))
+      { host: host, stdin: stdin, output: output, wait_thread: wait_thread }
+    end
+
+    # With ExitOnForwardFailure=yes the remote command only runs once the
+    # reverse forward is established, so receiving READY proves the tunnel is up.
+    def wait_until_ready(child, deadline)
+      remaining = deadline - Process.clock_gettime(Process::CLOCK_MONOTONIC)
+      raise Timeout::Error if remaining <= 0
+
+      Timeout.timeout(remaining) do
+        while (line = child[:output].gets)
+          return if line.strip == READY_TOKEN
         end
+
+        raise "Failed to establish port forward on #{child[:host]} (ssh exited #{child[:wait_thread].value.exitstatus})"
+      end
+    rescue Timeout::Error
+      raise "Timed out waiting for port forwarding to be established on #{child[:host]}"
+    end
+
+    def stop
+      Array(@children).each do |child|
+        close_io child[:stdin]
+        terminate child[:wait_thread]
+        close_io child[:output]
+      end
+
+      @children = []
+    end
+
+    def close_io(io)
+      io.close unless io.closed?
+    rescue IOError
+      # already closed
+    end
+
+    # Closing stdin sends EOF to the remote `cat`, so ssh usually exits on its
+    # own; TERM then KILL are a backstop.
+    def terminate(wait_thread)
+      # Once the wait thread has reaped the process, the PID may be recycled —
+      # don't signal it.
+      return unless wait_thread.alive?
+
+      Process.kill "TERM", wait_thread.pid
+      return if wait_thread.join(TEARDOWN_GRACE)
+
+      Process.kill "KILL", wait_thread.pid
+      wait_thread.join(TEARDOWN_GRACE)
+    rescue Errno::ESRCH, Errno::EPERM
+      # process already gone
+    end
+
+    def ssh_command(host)
+      [
+        "ssh", "-T",
+        "-o", "ExitOnForwardFailure=yes",
+        # BatchMode fails fast instead of prompting — a background tunnel can't
+        # service prompts. accept-new keeps net-ssh's default host key policy:
+        # accept unknown keys, reject changed ones.
+        "-o", "BatchMode=yes",
+        "-o", "StrictHostKeyChecking=accept-new",
+        *keepalive_options,
+        *port_option,
+        *key_options,
+        *config_option,
+        *proxy_option,
+        "-R", "127.0.0.1:#{port}:localhost:#{port}",
+        destination(host),
+        "echo #{READY_TOKEN} && exec cat"
+      ]
+    end
+
+    def destination(host)
+      ssh_options[:user] ? "#{ssh_options[:user]}@#{host}" : host.to_s
+    end
+
+    def port_option
+      ssh_options[:port] ? [ "-p", ssh_options[:port].to_s ] : []
+    end
+
+    def keepalive_options
+      return [] if ssh_options[:keepalive] == false
+
+      [ "-o", "ServerAliveInterval=#{ssh_options[:keepalive_interval] || 15}", "-o", "ServerAliveCountMax=4" ]
+    end
+
+    def key_options
+      options = Array(ssh_options[:keys]).flat_map { |key| [ "-i", key.to_s ] }
+      options += [ "-o", "IdentitiesOnly=yes" ] if ssh_options[:keys_only]
+      options += [ "-o", "ForwardAgent=yes" ] if ssh_options[:forward_agent]
+      options
+    end
+
+    def config_option
+      case (config = ssh_options[:config])
+      when nil, true
+        [] # ssh reads its default config files, matching net-ssh's default
+      when false
+        [ "-F", "/dev/null" ] # ignore all config files
+      else
+        config_file_option Array(config)
       end
+    end
 
-      raise "Timed out waiting for port forwarding to be established" unless ready.wait(30)
+    def config_file_option(paths)
+      # OpenSSH honors only the last -F flag, so multiple config files can't
+      # be passed through faithfully.
+      if paths.size > 1
+        raise "Multiple ssh config files (#{paths.join(", ")}) are not supported when forwarding the local registry port; combine them into one file, e.g. with Include"
+      end
+
+      [ "-F", (paths.first || "/dev/null").to_s ]
     end
 
-    def error(message)
-      SSHKit.config.output.error(message)
+    def proxy_option
+      case (proxy = ssh_options[:proxy])
+      when Net::SSH::Proxy::Jump
+        [ "-J", proxy.jump_proxies ]
+      when Net::SSH::Proxy::Command
+        [ "-o", "ProxyCommand=#{proxy.command_line_template}" ]
+      else
+        []
+      end
     end
 end

test/cli/build/port_forwarding_test.rb

@@ -0,0 +1,117 @@
+require "test_helper"
+
+class PortForwardingTest < ActiveSupport::TestCase
+  test "forwards the local registry port over the OS ssh client" do
+    command = ssh_command(user: "root", port: 22)
+
+    assert_equal "ssh", command.first
+    assert_includes_sequence command, [ "-R", "127.0.0.1:5000:localhost:5000" ]
+    assert_equal "root@1.1.1.1", command[-2]
+  end
+
+  test "does not pass -N so the readiness command runs" do
+    # -N tells ssh not to execute a remote command, which would prevent the
+    # READY handshake from ever firing and hang every deploy.
+    command = ssh_command(user: "root", port: 22)
+
+    assert_not_includes command, "-N"
+    assert_includes_sequence command, [ "-o", "ExitOnForwardFailure=yes" ]
+    assert_match(/\Aecho \S+ && exec cat\z/, command.last)
+  end
+
+  test "maps ssh options to ssh flags" do
+    command = ssh_command(
+      user: "app", port: 2222, keys: [ "/k1", "/k2" ], keys_only: true,
+      config: "/my/ssh_config", forward_agent: true, keepalive_interval: 45
+    )
+
+    assert_equal "app@1.1.1.1", command[-2]
+    assert_includes_sequence command, [ "-p", "2222" ]
+    assert_includes_sequence command, [ "-i", "/k1" ]
+    assert_includes_sequence command, [ "-i", "/k2" ]
+    assert_includes_sequence command, [ "-o", "IdentitiesOnly=yes" ]
+    assert_includes_sequence command, [ "-o", "ForwardAgent=yes" ]
+    assert_includes_sequence command, [ "-F", "/my/ssh_config" ]
+    assert_includes_sequence command, [ "-o", "ServerAliveInterval=45" ]
+  end
+
+  test "omits keepalives when keepalive is disabled" do
+    assert_empty ssh_command(keepalive: false).grep(/ServerAlive/)
+    assert_includes_sequence ssh_command(keepalive: true), [ "-o", "ServerAliveCountMax=4" ]
+  end
+
+  test "fails fast instead of prompting, but accepts unknown host keys like net-ssh" do
+    command = ssh_command
+
+    assert_includes_sequence command, [ "-o", "BatchMode=yes" ]
+    assert_includes_sequence command, [ "-o", "StrictHostKeyChecking=accept-new" ]
+  end
+
+  test "reads the default ssh config files when config is unset or true" do
+    assert_not_includes ssh_command, "-F"
+    assert_not_includes ssh_command(config: true), "-F"
+  end
+
+  test "ignores ssh config files when config is false" do
+    assert_includes_sequence ssh_command(config: false), [ "-F", "/dev/null" ]
+  end
+
+  test "maps a config file array like a single path" do
+    assert_includes_sequence ssh_command(config: [ "/my/ssh_config" ]), [ "-F", "/my/ssh_config" ]
+    assert_includes_sequence ssh_command(config: []), [ "-F", "/dev/null" ]
+  end
+
+  test "rejects multiple ssh config files" do
+    error = assert_raises(RuntimeError) { ssh_command(config: [ "/a", "/b" ]) }
+    assert_match(/Multiple ssh config files/, error.message)
+  end
+
+  test "rejects inline key_data instead of silently ignoring it" do
+    error = assert_raises(RuntimeError) do
+      Kamal::Cli::Build::PortForwarding.new([ "1.1.1.1" ], 5000, key_data: [ "-----BEGIN OPENSSH PRIVATE KEY-----" ])
+    end
+
+    assert_match(/key_data/, error.message)
+  end
+
+  test "waits for readiness against a shared deadline" do
+    forwarding = Kamal::Cli::Build::PortForwarding.new([ "1.1.1.1" ], 5000)
+    reader, writer = IO.pipe
+
+    error = assert_raises(RuntimeError) do
+      forwarding.send(:wait_until_ready, { host: "1.1.1.1", output: reader }, monotonic_now - 1)
+    end
+    assert_equal "Timed out waiting for port forwarding to be established on 1.1.1.1", error.message
+
+    writer.puts Kamal::Cli::Build::PortForwarding::READY_TOKEN
+    forwarding.send(:wait_until_ready, { host: "1.1.1.1", output: reader }, monotonic_now + 5)
+  ensure
+    [ reader, writer ].each(&:close)
+  end
+
+  test "maps a jump proxy to -J" do
+    command = ssh_command(proxy: Net::SSH::Proxy::Jump.new("root@bastion"))
+
+    assert_includes_sequence command, [ "-J", "root@bastion" ]
+  end
+
+  test "maps a proxy command to ProxyCommand" do
+    command = ssh_command(proxy: Net::SSH::Proxy::Command.new("connect -S relay %h %p"))
+
+    assert_includes_sequence command, [ "-o", "ProxyCommand=connect -S relay %h %p" ]
+  end
+
+  private
+    def ssh_command(**ssh_options)
+      Kamal::Cli::Build::PortForwarding.new([ "1.1.1.1" ], 5000, **ssh_options).send(:ssh_command, "1.1.1.1")
+    end
+
+    def monotonic_now
+      Process.clock_gettime(Process::CLOCK_MONOTONIC)
+    end
+
+    def assert_includes_sequence(array, subarray)
+      found = array.each_cons(subarray.size).include?(subarray)
+      assert found, "expected #{array.inspect}\nto contain consecutive #{subarray.inspect}"
+    end
+end