From 878c6892be0aa6dc606c3c239cb154ab5ef3797f Mon Sep 17 00:00:00 2001 From: Jamis Buck Date: Tue, 25 Mar 2008 22:42:55 -0600 Subject: initial gateway implementation and tests --- lib/net/ssh/gateway.rb | 203 +++++++++++++++++++++++++++++++++++++++++++++++++ test/gateway_test.rb | 116 ++++++++++++++++++++++++++++ 2 files changed, 319 insertions(+) create mode 100644 lib/net/ssh/gateway.rb create mode 100644 test/gateway_test.rb 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 -- cgit v1.2.1