Support both variants of Net::SSH::Service::Forward#open()

.gitignore

@@ -9,3 +9,6 @@
 /tmp/
 *.swp
 .DS_Store
+
+.byebug_history
+*.gem

.ruby-gemset

@@ -0,0 +1 @@
+net-ssh-gateway

lib/net/ssh/gateway.rb

@@ -99,37 +99,32 @@ def shutdown!
   #   # ...
   #   gateway.close(port)
   #
+  # This function takes variable arguments.  See comments in the case statement
+  # for details.
+  #
+  # The +local_host+ parameter specifies which network interface to bind to locally,
+  # and defaults to "127.0.0.1" if omitted.
   # If +local_port+ is not specified, the next available port will be used.
-  def open(host, port, local_port=nil)
-    ensure_open!
-
-    actual_local_port = local_port || next_port
-
-    @session_mutex.synchronize do
-      @session.forward.local(actual_local_port, host, port)
-    end
-
-    if block_given?
-      begin
-        yield actual_local_port
-      ensure
-        close(actual_local_port)
-      end
+  def open(*args)
+    case args.size
+    when 2 # open(host, port)
+      open3(*args[0,2], nil)
+    when 3 # open(host, port, local_port)
+      open3(*args)
+    when 4 # open(host, port, local_host, local_port)
+      open4(*args)
     else
-      return actual_local_port
+      raise ArgumentError, "Expecting 2..4 arguments, but got #{args.size} instead."
     end
-  rescue Errno::EADDRINUSE
-    raise if local_port # if a local port was explicitly requested, bubble the error up
-    retry
   end
 
   # Cancels port-forwarding over an open port that was previously opened via
   # #open.
-  def close(port)
+  def close(port, bind_address="127.0.0.1")
     ensure_open!
 
     @session_mutex.synchronize do
-      @session.forward.cancel_local(port)
+      @session.forward.cancel_local(port, bind_address)
     end
   end
 
@@ -193,4 +188,67 @@ def next_port
         port
       end
     end
+
+    # Opens a new port on the local host and forwards it to the given host/port
+    # via the gateway host. If a block is given, the newly allocated port
+    # number will be yielded to the block, and the port automatically closed
+    # (see #close) when the block finishes. Otherwise, the port number will be
+    # returned, and the caller is responsible for closing the port (#close).
+    #
+    #   gateway.open('host', 80) do |port|
+    #     # ...
+    #   end
+    #
+    #   port = gateway.open('host', 80)
+    #   # ...
+    #   gateway.close(port)
+    #
+    # Unlike open4(), this method always binds to the loopback network
+    # interface # of "127.0.0.1".
+    #
+    # If +local_port+ is not specified, the next available port will be used.
+    def open3(host, port, local_port=nil)
+      open4(host, port, "127.0.0.1", local_port)
+    end
+
+    # Opens a new port on the local host and forwards it to the given host/port
+    # via the gateway host. If a block is given, the newly allocated port
+    # number will be yielded to the block, and the port automatically closed
+    # (see #close) when the block finishes. Otherwise, the port number will be
+    # returned, and the caller is responsible for closing the port (#close).
+    #
+    #   gateway.open('host', 80) do |port|
+    #     # ...
+    #   end
+    #
+    #   port = gateway.open('host', 80)
+    #   # ...
+    #   gateway.close(port)
+    #
+    # Unlike open3(), this method specifies which network interface to bind to
+    # locally via the +local_host+ parameter.
+    #
+    # If +local_port+ is not specified, the next available port will be used.
+    def open4(host, port, local_host, local_port=nil)
+      ensure_open!
+
+      actual_local_port = local_port || next_port
+
+      @session_mutex.synchronize do
+        @session.forward.local(local_host, actual_local_port, host, port)
+      end
+
+      if block_given?
+        begin
+          yield actual_local_port
+        ensure
+          close(actual_local_port)
+        end
+      else
+        return actual_local_port
+      end
+    rescue Errno::EADDRINUSE
+      raise if local_port # if a local port was explicitly requested, bubble the error up
+      retry
+    end
 end

lib/net/ssh/gateway/version.rb

@@ -1,7 +1,7 @@
 module Net
   module SSH
     class Gateway
-      VERSION = "2.0.0"
+      VERSION = "2.1.0"
     end
   end
 end

net-ssh-gateway.gemspec

@@ -23,6 +23,8 @@ Gem::Specification.new do |spec|
   spec.add_development_dependency "rake", "~> 10.0"
   spec.add_development_dependency "minitest", "~> 5.8.4"
   spec.add_development_dependency "mocha"
+  spec.add_development_dependency "echoe"
+  spec.add_development_dependency "byebug"
 
   spec.add_runtime_dependency "net-ssh", ">= 2.6.5"
 end

test/net/ssh/gateway_test.rb

@@ -20,27 +20,33 @@ def test_shutdown_without_any_open_connections_should_terminate_session
   def test_open_should_start_local_ports_at_65535
     gateway_session, gateway = new_gateway
     assert_equal 65535, gateway.open("app1", 22)
-    assert_equal [65535, "app1", 22], gateway_session.forward.active_locals[65535]
+    assert_equal ["127.0.0.1", 65535, "app1", 22], gateway_session.forward.active_locals[65535]
   end
 
   def test_open_should_decrement_port_and_retry_if_ports_are_in_use
     gateway_session, gateway = new_gateway(:reserved => lambda { |n| n > 65000 })
     assert_equal 65000, gateway.open("app1", 22)
-    assert_equal [65000, "app1", 22], gateway_session.forward.active_locals[65000]
+    assert_equal ["127.0.0.1", 65000, "app1", 22], gateway_session.forward.active_locals[65000]
   end
 
   def test_open_with_explicit_local_port_should_use_that_port
     gateway_session, gateway = new_gateway
     assert_equal 8181, gateway.open("app1", 22, 8181)
-    assert_equal [8181, "app1", 22], gateway_session.forward.active_locals[8181]
+    assert_equal ["127.0.0.1", 8181, "app1", 22], gateway_session.forward.active_locals[8181]
+  end
+
+  def test_open_with_explicit_local_host_and_port_should_use_that_port
+    gateway_session, gateway = new_gateway
+    assert_equal 8181, gateway.open("app1", 22, "1.2.3.4", 8181)
+    assert_equal ["1.2.3.4", 8181, "app1", 22], gateway_session.forward.active_locals[8181]
   end
 
   def test_ssh_should_return_connection_when_no_block_is_given
     gateway_session, gateway = new_gateway
     expect_connect_to("127.0.0.1", "user", :port => 65535).returns(result = mock("session"))
     newsess = gateway.ssh("app1", "user")
     assert_equal result, newsess
-    assert_equal [65535, "app1", 22], gateway_session.forward.active_locals[65535]
+    assert_equal ["127.0.0.1", 65535, "app1", 22], gateway_session.forward.active_locals[65535]
   end
 
   def test_ssh_with_block_should_yield_session_and_then_close_port
@@ -91,9 +97,9 @@ def cancel_local(port)
         @active_locals.delete(port)
       end
 
-      def local(lport, host, rport)
+      def local(lhost, lport, host, rport)
         raise Errno::EADDRINUSE if @options[:reserved] && @options[:reserved][lport]
-        @active_locals[lport] = [lport, host, rport]
+        @active_locals[lport] = [lhost, lport, host, rport]
       end
     end