initial gateway implementation and tests

2 files changed, 319 insertions, 0 deletions

diff --git a/lib/net/ssh/gateway.rb b/lib/net/ssh/gateway.rb
new file mode 100644
index 0000000..aed6bdd
--- /dev/null
+++ b/lib/net/ssh/gateway.rb
@@ -0,0 +1,203 @@
+require 'thread'
+require 'net/ssh'
+require 'net/ssh/version'
+
+# A Gateway is an object that allows you to tunnel network connections through
+# a publicly visible host to a host hidden behind it. This is particularly
+# useful when dealing with hosts behind a firewall. One host will generally
+# be visible (and accessible) outside the firewall, while the others will all
+# be behind the firewall, and the only way to access those restricted hosts
+# is by first logging into the publicly visible host, and from thence logging
+# into the restricted ones.
+#
+# This class makes it easy to programmatically connect to these restricted
+# hosts via SSH. You can either simply forward a port from the local host to
+# the remote host, or you can open a new Net::SSH connection to the remote
+# host via a forwarded port.
+#
+#   require 'net/ssh/gateway'
+#
+#   gateway = Net::SSH::Gateway.new('host.name', 'user')
+#
+#   gateway.open('hidden.host', 80) do |port|
+#     Net::HTTP.get_print '127.0.0.1', '/path', port
+#   end
+#
+#   gateway.ssh('hidden.host', 'user') do |ssh|
+#     puts ssh.exec!("hostname")
+#   end
+#
+#   gateway.shutdown!
+#
+# Port numbers are allocated automatically, beginning at MAX_PORT and
+# decrementing on each request for a new port until MIN_PORT is reached. If
+# a port is already in use, this is detected and a different port will be
+# assigned. 
+class Net::SSH::Gateway
+  # A trivial class for representing the version of this library.
+  class Version < Net::SSH::Version
+    # The major component of the library's version
+    MAJOR = 0
+
+    # The minor component of the library's version
+    MINOR = 99
+
+    # The tiny component of the library's version
+    TINY  = 0
+
+    # The library's version as a Version instance
+    CURRENT = new(MAJOR, MINOR, TINY)
+
+    # The library's version as a String instance
+    STRING = CURRENT.to_s
+  end
+
+  # The maximum port number that the gateway will attempt to use to forward
+  # connections from.
+  MAX_PORT = 65535
+
+  # The minimum port number that the gateway will attempt to use to forward
+  # connections from.
+  MIN_PORT = 1024
+
+  # Instantiate a new Gateway object, using the given remote host as the
+  # tunnel. The arguments here are identical to those for Net::SSH.start, and
+  # are passed as given to that method to start up the gateway connection.
+  #
+  #   gateway = Net::SSH::Gateway.new('host', 'user', :password => "password")
+  def initialize(host, user, options={})
+    @session = Net::SSH.start(host, user, options)
+    @session_mutex = Mutex.new
+    @port_mutex = Mutex.new
+    @next_port = MAX_PORT
+
+    initiate_event_loop!
+  end
+
+  # Returns +true+ if the gateway is currently open and accepting connections.
+  # This will be the case unless #shutdown! has been invoked.
+  def active?
+    @active
+  end
+
+  # Shuts down the gateway by closing all forwarded ports and then closing
+  # the gateway's SSH session.
+  def shutdown!
+    return unless active?
+
+    @session_mutex.synchronize do
+      # cancel all active forward channels
+      @session.forward.active_locals.each do |lport, host, port|
+        @session.forward.cancel_local(lport)
+      end
+    end
+
+    @active = false
+    
+    @thread.join
+    @session.close
+  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)
+  def open(host, port)
+    ensure_open!
+
+    local_port = next_port
+
+    @session_mutex.synchronize do
+      @session.forward.local(local_port, host, port)
+    end
+
+    if block_given?
+      begin
+        yield local_port
+      ensure
+        close(local_port)
+      end
+    else
+      return local_port
+    end
+  rescue Errno::EADDRINUSE
+    retry
+  end
+
+  # Cancels port-forwarding over an open port that was previously opened via
+  # #open.
+  def close(port)
+    ensure_open!
+
+    @session_mutex.synchronize do
+      @session.forward.cancel_local(port)
+    end
+  end
+
+  # Forwards a new connection to the given +host+ and opens a new Net::SSH
+  # connection to that host over the forwarded port. If a block is given,
+  # the new SSH connection will be yielded to the block, and autoclosed
+  # when the block terminates. The forwarded port will be autoclosed as well.
+  # If no block was given, the new SSH connection will be returned, and it
+  # is up to the caller to terminate both the connection and the forwarded
+  # port when done.
+  #
+  #   gateway.ssh('host', 'user') do |ssh|
+  #     # ...
+  #   end
+  #
+  #   ssh = gateway.ssh('host', 'user')
+  #   # ...
+  #   ssh.close
+  #   gateway.close(ssh.transport.port)
+  def ssh(host, user, options={}, &block)
+    local_port = open(host, options[:port] || 22)
+
+    begin
+      Net::SSH.start("127.0.0.1", user, options.merge(:port => local_port), &block)
+    ensure
+      close(local_port) if block
+    end
+  end
+
+  private
+
+    # Raises a RuntimeError if the gateway is not active. This is used as a
+    # sanity check to make sure a client doesn't try to call any methods on
+    # a closed gateway.
+    def ensure_open!
+      raise "attempt to use a closed gateway" unless active?
+    end
+
+    # Fires up the gateway session's event loop within a thread, so that it
+    # can run in the background. The loop will run for as long as the gateway
+    # remains active.
+    def initiate_event_loop!
+      @active = true
+
+      @thread = Thread.new do
+        while @active
+          @session_mutex.synchronize { @session.process(0.1) }
+        end
+      end
+    end
+
+    # Grabs the next available port number and returns it.
+    def next_port
+      @port_mutex.synchronize do
+        port = @next_port
+        @next_port -= 1
+        @next_port = MAX_PORT if @next_port < MIN_PORT
+        port
+      end
+    end
+end
\ No newline at end of file
diff --git a/test/gateway_test.rb b/test/gateway_test.rb
new file mode 100644
index 0000000..6194886
--- /dev/null
+++ b/test/gateway_test.rb
@@ -0,0 +1,116 @@
+require 'test/unit'
+require 'mocha'
+require 'net/ssh/gateway'
+
+class GatewayTest < Test::Unit::TestCase
+  def teardown
+    Thread.list { |t| t.kill unless Thread.current == t }
+  end
+
+  def test_shutdown_without_any_open_connections_should_terminate_session
+    session, gateway = new_gateway
+    session.expects(:close)
+    gateway.shutdown!
+    assert !gateway.active?
+    assert session.forward.active_locals.empty?
+  end
+
+  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]
+  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]
+  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]
+  end
+
+  def test_ssh_with_block_should_yield_session_and_then_close_port
+    gateway_session, gateway = new_gateway
+    expect_connect_to("127.0.0.1", "user", :port => 65535).yields(result = mock("session"))
+    yielded = false
+    gateway.ssh("app1", "user") do |newsess|
+      yielded = true
+      assert_equal result, newsess
+    end
+    assert yielded
+    assert gateway_session.forward.active_locals.empty?
+  end
+
+  def test_shutdown_should_cancel_active_forwarded_ports
+    gateway_session, gateway = new_gateway
+    gateway.open("app1", 80)
+    assert !gateway_session.forward.active_locals.empty?
+    gateway.shutdown!
+    assert gateway_session.forward.active_locals.empty?
+  end
+
+  private
+
+    def expect_connect_to(host, user, options={})
+      Net::SSH.expects(:start).with do |real_host, real_user, real_options|
+        host == real_host &&
+        user == real_user &&
+        options[:port] == real_options[:port]
+      end
+    end
+
+    def new_gateway(options={})
+      session = MockSession.new(options)
+      expect_connect_to("test.host", "tester").returns(session)
+      [session, Net::SSH::Gateway.new("test.host", "tester")]
+    end
+
+    class MockForward
+      attr_reader :active_locals
+
+      def initialize(options)
+        @options = options
+        @active_locals = {}
+      end
+
+      def cancel_local(port)
+        @active_locals.delete(port)
+      end
+
+      def local(lport, host, rport)
+        raise Errno::EADDRINUSE if @options[:reserved] && @options[:reserved][lport]
+        @active_locals[lport] = [lport, host, rport]
+      end
+    end
+
+    class MockSession
+      attr_reader :forward
+
+      def initialize(options={})
+        @forward = MockForward.new(options)
+      end
+
+      def close
+      end
+
+      def process(wait=nil)
+        true
+      end
+
+      def looping?
+        @looping
+      end
+
+      def loop
+        @looping = true
+        sleep 0.1 while yield
+        @looping = false
+      end
+    end
+end
\ No newline at end of file