connection session tests

git-svn-id: http://svn.jamisbuck.org/net-ssh/branches/v2@182 1d2a57f2-1ded-0310-ad52-83097a15a5de

7 files changed, 455 insertions, 12 deletions

diff --git a/lib/net/ssh/connection/session.rb b/lib/net/ssh/connection/session.rb
index a08a327..bb742ed 100644
--- a/lib/net/ssh/connection/session.rb
+++ b/lib/net/ssh/connection/session.rb
@@ -21,7 +21,7 @@ module Net; module SSH; module Connection
     attr_reader :pending_requests
     attr_reader :readers
     attr_reader :writers
-    attr_reader :channel_open_handler
+    attr_reader :channel_open_handlers
     attr_reader :options
 
     def initialize(transport, options={})
@@ -36,7 +36,7 @@ module Net; module SSH; module Connection
       @pending_requests = []
       @readers = [transport.socket]
       @writers = [transport.socket]
-      @channel_open_handler = {}
+      @channel_open_handlers = {}
       @on_global_request = {}
     end
 
@@ -107,7 +107,7 @@ module Net; module SSH; module Connection
     end
 
     def send_message(message)
-      transport.socket.enqueue_packet(message)
+      transport.enqueue_message(message)
     end
 
     def listen_to(io, &callback)
@@ -127,7 +127,7 @@ module Net; module SSH; module Connection
     end
 
     def on_open_channel(type, &block)
-      channel_open_handler[type] = block
+      channel_open_handlers[type] = block
     end
 
     def on_global_request(type, &block)
@@ -157,7 +157,7 @@ module Net; module SSH; module Connection
         result = callback ? callback.call(packet[:request_data], packet[:want_reply]) : false
 
         if result != :sent && result != true && result != false
-          raise "expected global request handler for #{packet[:request_type]} to return true or false"
+          raise "expected global request handler for `#{packet[:request_type]}' to return true, false, or :sent, but got #{result.inspect}"
         end
 
         if packet[:want_reply] && result != :sent
@@ -185,12 +185,13 @@ module Net; module SSH; module Connection
         channel = Channel.new(self, packet[:channel_type], local_id)
         channel.do_open_confirmation(packet[:remote_id], packet[:window_size], packet[:packet_size])
 
-        callback = channel_open_handler[packet[:channel_type]]
+        callback = channel_open_handlers[packet[:channel_type]]
 
         if callback
-          result = callback[self, channel, packet]
-          if Array === result && result.length == 2
-            failure = result
+          begin
+            callback[self, channel, packet]
+          rescue ChannelOpenFailed => err
+            failure = [err.code, err.reason]
           else
             channels[local_id] = channel
             msg = Buffer.from(:byte, CHANNEL_OPEN_CONFIRMATION, :long, channel.remote_id, :long, channel.local_id, :long, channel.local_maximum_window_size, :long, channel.local_maximum_packet_size)
diff --git a/lib/net/ssh/errors.rb b/lib/net/ssh/errors.rb
index 724ff67..d7c1ee2 100644
--- a/lib/net/ssh/errors.rb
+++ b/lib/net/ssh/errors.rb
@@ -5,6 +5,15 @@ module Net; module SSH
 
   class Disconnect < Exception; end
 
+  class ChannelOpenFailed < Exception
+    attr_reader :code, :reason
+
+    def initialize(code, reason)
+      @code, @reason = code, reason
+      super "#{reason} (#{code})"
+    end
+  end
+
   # Raised when the cached key for a particular host does not match the
   # key given by the host, which can be indicative of a man-in-the-middle
   # attack. When rescuing this exception, you can inspect the key fingerprint
diff --git a/lib/net/ssh/service/forward.rb b/lib/net/ssh/service/forward.rb
index 924fd57..9fedec4 100644
--- a/lib/net/ssh/service/forward.rb
+++ b/lib/net/ssh/service/forward.rb
@@ -149,13 +149,15 @@ module Net; module SSH; module Service
         remote = @remote_forwarded_ports[[connected_port, connected_address]]
 
         if remote.nil?
-          raise Net::SSH::Exception, "unknown request from remote forwarded connection on #{connected_address}:#{connected_port}"
+          raise Net::SSH::ChannelOpenFailed.new(1, "unknown request from remote forwarded connection on #{connected_address}:#{connected_port}")
         end
 
+        client = TCPSocket.new(remote.host, remote.port)
         trace { "connected #{connected_address}:#{connected_port} originator #{originator_address}:#{originator_port}" }
 
-        client = TCPSocket.new(remote.host, remote.port)
         prepare_client(client, channel, :remote)
+      rescue SocketError => err
+        raise Net::SSH::ChannelOpenFailed.new(2, "could not connect to remote host (#{remote.host}:#{remote.port}): #{err.message}")
       end
 
       def auth_agent_channel(session, channel, packet)
diff --git a/lib/net/ssh/transport/session.rb b/lib/net/ssh/transport/session.rb
index 9f0a610..8b96b42 100644
--- a/lib/net/ssh/transport/session.rb
+++ b/lib/net/ssh/transport/session.rb
@@ -138,6 +138,10 @@ module Net; module SSH; module Transport
       socket.send_packet(message)
     end
 
+    def enqueue_message(message)
+      socket.enqueue_packet(message)
+    end
+
     def configure_client(options={})
       socket.client.set(options)
     end
diff --git a/test/common.rb b/test/common.rb
index 1823d40..b58ae80 100644
--- a/test/common.rb
+++ b/test/common.rb
@@ -28,7 +28,7 @@ class MockTransport < Net::SSH::Transport::Session
 
   attr_reader :client_options
   attr_reader :server_options
-  attr_reader :hints
+  attr_reader :hints, :queue
 
   def initialize(options={})
     self.logger = options[:logger]
@@ -52,6 +52,10 @@ class MockTransport < Net::SSH::Transport::Session
     end
   end
 
+  def poll_message
+    @queue.shift
+  end
+
   def next_message
     @queue.shift or raise "expected a message from the server but nothing was ready to send"
   end
diff --git a/test/connection/test_session.rb b/test/connection/test_session.rb
new file mode 100644
index 0000000..b6921a0
--- /dev/null
+++ b/test/connection/test_session.rb
@@ -0,0 +1,416 @@
+$LOAD_PATH.unshift("#{File.dirname(__FILE__)}/..").uniq!
+require 'common'
+require 'net/ssh/connection/session'
+
+module Connection
+
+  class TestSession < Test::Unit::TestCase
+    include Net::SSH::Connection::Constants
+
+    def test_constructor_should_set_defaults
+      assert session.channels.empty?
+      assert session.listeners.empty?
+      assert session.pending_requests.empty?
+      assert_equal [socket], session.readers
+      assert_equal [socket], session.writers
+    end
+
+    def test_on_open_channel_should_register_block_with_given_channel_type
+      flag = false
+      session.on_open_channel("testing") { flag = true }
+      assert_not_nil session.channel_open_handlers["testing"]
+      session.channel_open_handlers["testing"].call
+      assert flag, "callback should have been invoked"
+    end
+
+    def test_forward_should_create_and_cache_instance_of_forward_service
+      assert_instance_of Net::SSH::Service::Forward, session.forward
+      assert_equal session.forward.object_id, session.forward.object_id
+    end
+
+    def test_listen_to_should_add_argument_to_readers_list
+      io = stub("io")
+      session.listen_to(io)
+      assert session.readers.include?(io)
+      assert !session.writers.include?(io)
+      assert !session.listeners.key?(io)
+    end
+
+    def test_listen_to_should_add_argument_to_writers_list_if_it_responds_to_pending_write
+      io = stub("io", :pending_write? => true)
+      session.listen_to(io)
+      assert session.readers.include?(io)
+      assert session.writers.include?(io)
+      assert !session.listeners.key?(io)
+    end
+
+    def test_listen_to_should_add_argument_to_listeners_list_if_block_is_given
+      io = stub("io", :pending_write? => true)
+      flag = false
+      session.listen_to(io) { flag = true }
+      assert !flag, "callback should not be invoked immediately"
+      assert session.readers.include?(io)
+      assert session.writers.include?(io)
+      assert session.listeners.key?(io)
+      session.listeners[io].call
+      assert flag, "callback should have been invoked"
+    end
+
+    def test_stop_listening_to_should_remove_argument_from_lists
+      io = stub("io", :pending_write? => true)
+
+      session.listen_to(io) { }
+      assert session.readers.include?(io)
+      assert session.writers.include?(io)
+      assert session.listeners.key?(io)
+
+      session.stop_listening_to(io)
+      assert !session.readers.include?(io)
+      assert !session.writers.include?(io)
+      assert !session.listeners.key?(io)
+    end
+
+    def test_send_message_should_enqueue_message_at_transport_layer
+      packet = P(:byte, REQUEST_SUCCESS)
+      session.send_message(packet)
+      assert_equal packet.to_s, socket.write_buffer
+    end
+
+    def test_open_channel_defaults_should_use_session_channel
+      flag = false
+      channel = session.open_channel { flag = true }
+      assert !flag, "callback should not be invoked immediately"
+      channel.do_open_confirmation(1,2,3)
+      assert flag, "callback should have been invoked"
+      assert_equal "session", channel.type
+      assert_equal 0, channel.local_id
+      assert_equal channel, session.channels[channel.local_id]
+
+      packet = P(:byte, CHANNEL_OPEN, :string, "session", :long, channel.local_id,
+        :long, channel.local_maximum_window_size, :long, channel.local_maximum_packet_size)
+      assert_equal packet.to_s, socket.write_buffer
+    end
+
+    def test_open_channel_with_type_should_use_type
+      channel = session.open_channel("direct-tcpip")
+      assert_equal "direct-tcpip", channel.type
+      packet = P(:byte, CHANNEL_OPEN, :string, "direct-tcpip", :long, channel.local_id,
+        :long, channel.local_maximum_window_size, :long, channel.local_maximum_packet_size)
+      assert_equal packet.to_s, socket.write_buffer
+    end
+
+    def test_open_channel_with_extras_should_append_extras_to_packet
+      channel = session.open_channel("direct-tcpip", :string, "other.host", :long, 1234)
+      packet = P(:byte, CHANNEL_OPEN, :string, "direct-tcpip", :long, channel.local_id,
+        :long, channel.local_maximum_window_size, :long, channel.local_maximum_packet_size,
+        :string, "other.host", :long, 1234)
+      assert_equal packet.to_s, socket.write_buffer
+    end
+
+    def test_send_global_request_without_callback_should_not_expect_reply
+      packet = P(:byte, GLOBAL_REQUEST, :string, "testing", :bool, false)
+      session.send_global_request("testing")
+      assert_equal packet.to_s, socket.write_buffer
+      assert session.pending_requests.empty?
+    end
+
+    def test_send_global_request_with_callback_should_expect_reply
+      packet = P(:byte, GLOBAL_REQUEST, :string, "testing", :bool, true)
+      proc = Proc.new {}
+      session.send_global_request("testing", &proc)
+      assert_equal packet.to_s, socket.write_buffer
+      assert_equal [proc], session.pending_requests
+    end
+
+    def test_send_global_request_with_extras_should_append_extras_to_packet
+      packet = P(:byte, GLOBAL_REQUEST, :string, "testing", :bool, false, :string, "other.host", :long, 1234)
+      session.send_global_request("testing", :string, "other.host", :long, 1234)
+      assert_equal packet.to_s, socket.write_buffer
+    end
+
+    def test_process_should_exit_immediately_if_block_is_false
+      session.channels[0] = stub("channel", :closing? => false)
+      session.channels[0].expects(:process).never
+      process_times(0)
+    end
+
+    def test_process_should_exit_after_processing_if_block_is_true_then_false
+      session.channels[0] = stub("channel", :closing? => false)
+      session.channels[0].expects(:process)
+      IO.expects(:select).never
+      process_times(2)
+    end
+
+    def test_process_should_not_process_channels_that_are_closing
+      session.channels[0] = stub("channel", :closing? => true)
+      session.channels[0].expects(:process).never
+      IO.expects(:select).never
+      process_times(2)
+    end
+
+    def test_global_request_packets_should_be_silently_handled_if_no_handler_exists_for_them
+      transport.return(GLOBAL_REQUEST, :string, "testing", :bool, false)
+      process_times(2)
+      assert transport.queue.empty?
+      assert !socket.pending_write?
+    end
+
+    def test_global_request_packets_should_be_auto_replied_to_even_if_no_handler_exists
+      transport.return(GLOBAL_REQUEST, :string, "testing", :bool, true)
+      process_times(2)
+      assert_equal P(:byte, REQUEST_FAILURE).to_s, socket.write_buffer
+    end
+
+    def test_global_request_handler_should_not_trigger_auto_reply_if_no_reply_is_wanted
+      flag = false
+      session.on_global_request("testing") { flag = true }
+      assert !flag, "callback should not be invoked yet"
+      transport.return(GLOBAL_REQUEST, :string, "testing", :bool, false)
+      process_times(2)
+      assert transport.queue.empty?
+      assert !socket.pending_write?
+      assert flag, "callback should have been invoked"
+    end
+
+    def test_global_request_handler_returning_true_should_trigger_success_auto_reply
+      flag = false
+      session.on_global_request("testing") { flag = true }
+      transport.return(GLOBAL_REQUEST, :string, "testing", :bool, true)
+      process_times(2)
+      assert_equal P(:byte, REQUEST_SUCCESS).to_s, socket.write_buffer
+      assert flag
+    end
+
+    def test_global_request_handler_returning_false_should_trigger_failure_auto_reply
+      flag = false
+      session.on_global_request("testing") { flag = true; false }
+      transport.return(GLOBAL_REQUEST, :string, "testing", :bool, true)
+      process_times(2)
+      assert_equal P(:byte, REQUEST_FAILURE).to_s, socket.write_buffer
+      assert flag
+    end
+
+    def test_global_request_handler_returning_sent_should_not_trigger_auto_reply
+      flag = false
+      session.on_global_request("testing") { flag = true; :sent }
+      transport.return(GLOBAL_REQUEST, :string, "testing", :bool, true)
+      process_times(2)
+      assert !socket.pending_write?
+      assert flag
+    end
+
+    def test_global_request_handler_returning_other_value_should_raise_error
+      session.on_global_request("testing") { "bug" }
+      transport.return(GLOBAL_REQUEST, :string, "testing", :bool, true)
+      assert_raises(RuntimeError) { process_times(2) }
+    end
+
+    def test_request_success_packets_should_invoke_next_pending_request_with_true
+      result = nil
+      session.pending_requests << Proc.new { |*args| result = args }
+      transport.return(REQUEST_SUCCESS)
+      process_times(2)
+      assert_equal [true, P(:byte, REQUEST_SUCCESS)], result
+      assert session.pending_requests.empty?
+    end
+
+    def test_request_failure_packets_should_invoke_next_pending_request_with_false
+      result = nil
+      session.pending_requests << Proc.new { |*args| result = args }
+      transport.return(REQUEST_FAILURE)
+      process_times(2)
+      assert_equal [false, P(:byte, REQUEST_FAILURE)], result
+      assert session.pending_requests.empty?
+    end
+
+    def test_channel_open_packet_without_corresponding_channel_open_handler_should_result_in_channel_open_failure
+      transport.return(CHANNEL_OPEN, :string, "auth-agent", :long, 14, :long, 0x20000, :long, 0x10000)
+      process_times(2)
+      assert_equal P(:byte, CHANNEL_OPEN_FAILURE, :long, 14, :long, 3, :string, "unknown channel type auth-agent", :string, "").to_s, socket.write_buffer
+    end
+
+    def test_channel_open_packet_with_corresponding_handler_should_result_in_channel_open_failure_when_handler_returns_an_error
+      transport.return(CHANNEL_OPEN, :string, "auth-agent", :long, 14, :long, 0x20000, :long, 0x10000)
+      session.on_open_channel "auth-agent" do |s, ch, p|
+        raise Net::SSH::ChannelOpenFailed.new(1234, "we iz in ur channelz!")
+      end
+      process_times(2)
+      assert_equal P(:byte, CHANNEL_OPEN_FAILURE, :long, 14, :long, 1234, :string, "we iz in ur channelz!", :string, "").to_s, socket.write_buffer
+    end
+
+    def test_channel_open_packet_with_corresponding_handler_should_result_in_channel_open_confirmation_when_handler_succeeds
+      transport.return(CHANNEL_OPEN, :string, "auth-agent", :long, 14, :long, 0x20001, :long, 0x10001)
+      result = nil
+      session.on_open_channel("auth-agent") { |*args| result = args }
+      process_times(2)
+      assert_equal P(:byte, CHANNEL_OPEN_CONFIRMATION, :long, 14, :long, 0, :long, 0x20000, :long, 0x10000).to_s, socket.write_buffer
+      assert_not_nil(ch = session.channels[0])
+      assert_equal [session, ch, P(:byte, CHANNEL_OPEN, :string, "auth-agent", :long, 14, :long, 0x20001, :long, 0x10001)], result
+      assert_equal 0, ch.local_id
+      assert_equal 14, ch.remote_id
+      assert_equal 0x20001, ch.remote_maximum_window_size
+      assert_equal 0x10001, ch.remote_maximum_packet_size
+      assert_equal 0x20000, ch.local_maximum_window_size
+      assert_equal 0x10000, ch.local_maximum_packet_size
+      assert_equal "auth-agent", ch.type
+    end
+
+    def test_channel_open_confirmation_packet_should_be_routed_to_corresponding_channel
+      channel_at(14).expects(:do_open_confirmation).with(1234, 0x20001, 0x10001)
+      transport.return(CHANNEL_OPEN_CONFIRMATION, :long, 14, :long, 1234, :long, 0x20001, :long, 0x10001)
+      process_times(2)
+    end
+
+    def test_channel_window_adjust_packet_should_be_routed_to_corresponding_channel
+      channel_at(14).expects(:do_window_adjust).with(5000)
+      transport.return(CHANNEL_WINDOW_ADJUST, :long, 14, :long, 5000)
+      process_times(2)
+    end
+
+    def test_channel_request_packet_should_be_routed_to_corresponding_channel
+      channel_at(14).expects(:do_request).with("testing", false, Net::SSH::Buffer.new)
+      transport.return(CHANNEL_REQUEST, :long, 14, :string, "testing", :bool, false)
+      process_times(2)
+    end
+
+    def test_channel_data_packet_should_be_routed_to_corresponding_channel
+      channel_at(14).expects(:do_data).with("bring it on down")
+      transport.return(CHANNEL_DATA, :long, 14, :string, "bring it on down")
+      process_times(2)
+    end
+
+    def test_channel_extended_data_packet_should_be_routed_to_corresponding_channel
+      channel_at(14).expects(:do_extended_data).with(1, "bring it on down")
+      transport.return(CHANNEL_EXTENDED_DATA, :long, 14, :long, 1, :string, "bring it on down")
+      process_times(2)
+    end
+
+    def test_channel_eof_packet_should_be_routed_to_corresponding_channel
+      channel_at(14).expects(:do_eof).with()
+      transport.return(CHANNEL_EOF, :long, 14)
+      process_times(2)
+    end
+
+    def test_channel_success_packet_should_be_routed_to_corresponding_channel
+      channel_at(14).expects(:do_success).with()
+      transport.return(CHANNEL_SUCCESS, :long, 14)
+      process_times(2)
+    end
+
+    def test_channel_failure_packet_should_be_routed_to_corresponding_channel
+      channel_at(14).expects(:do_failure).with()
+      transport.return(CHANNEL_FAILURE, :long, 14)
+      process_times(2)
+    end
+
+    def test_channel_close_packet_should_be_routed_to_corresponding_channel_and_channel_should_be_closed_and_removed
+      channel_at(14).expects(:do_close).with()
+      session.channels[14].expects(:close).with()
+      transport.return(CHANNEL_CLOSE, :long, 14)
+      process_times(2)
+      assert session.channels.empty?
+    end
+
+    def test_multiple_pending_dispatches_should_be_dispatched_together
+      channel_at(14).expects(:do_eof).with()
+      session.channels[14].expects(:do_success).with()
+      transport.return(CHANNEL_SUCCESS, :long, 14)
+      transport.return(CHANNEL_EOF, :long, 14)
+      process_times(2)
+    end
+
+    def test_writers_without_pending_writes_should_not_be_considered_for_select
+      IO.expects(:select).with([socket],[],nil,nil).returns([[],[],[]])
+      session.process
+    end
+
+    def test_writers_with_pending_writes_should_be_considered_for_select
+      socket.enqueue("laksdjflasdkf")
+      IO.expects(:select).with([socket],[socket],nil,nil).returns([[],[],[]])
+      session.process
+    end
+
+    def test_ready_readers_should_be_filled
+      socket.expects(:recv).returns("this is some data")
+      IO.expects(:select).with([socket],[],nil,nil).returns([[socket],[],[]])
+      session.process
+      assert_equal [socket], session.readers
+    end
+
+    def test_ready_readers_that_cant_be_filled_should_be_removed
+      socket.expects(:recv).returns("")
+      socket.expects(:close)
+      IO.expects(:select).with([socket],[],nil,nil).returns([[socket],[],[]])
+      session.process
+      assert_equal [], session.readers
+      assert_equal [], session.writers
+    end
+
+    def test_ready_readers_that_are_registered_with_a_block_should_call_block_instead_of_fill
+      io = stub("io")
+      flag = false
+      session.listen_to(io) { flag = true }
+      IO.expects(:select).with([socket,io],[],nil,nil).returns([[io],[],[]])
+      session.process
+      assert flag, "callback should have been invoked"
+    end
+
+    def test_ready_writers_should_call_send_pending
+      socket.enqueue("laksdjflasdkf")
+      socket.expects(:send).with("laksdjflasdkf", 0).returns(13)
+      IO.expects(:select).with([socket],[socket],nil,nil).returns([[],[socket],[]])
+      session.process
+    end
+
+    def test_process_should_call_rekey_as_needed
+      transport.expects(:rekey_as_needed)
+      IO.expects(:select).with([socket],[],nil,nil).returns([[],[],[]])
+      session.process
+    end
+
+    def test_loop_should_call_process_until_process_returns_false
+      IO.stubs(:select).with([socket],[],nil,nil).returns([[],[],[]])
+      session.expects(:process).with(nil).times(4).returns(true,true,true,false).yields
+      n = 0
+      session.loop { n += 1 }
+      assert_equal 4, n
+    end
+
+    private
+
+      module MockSocket
+        # so that we can easily test the contents that were enqueued, without
+        # worrying about all the packet stream overhead
+        def enqueue_packet(message)
+          enqueue(message.to_s)
+        end
+      end
+
+      def socket
+        @socket ||= begin
+          socket ||= Object.new
+          socket.extend(Net::SSH::Transport::PacketStream)
+          socket.extend(MockSocket)
+          socket
+        end
+      end
+
+      def channel_at(local_id)
+        session.channels[local_id] = stub("channel", :process => true, :closing? => false)
+      end
+
+      def transport(options={})
+        @transport ||= MockTransport.new(options.merge(:socket => socket))
+      end
+
+      def session(options={})
+        @session ||= Net::SSH::Connection::Session.new(transport, options)
+      end
+
+      def process_times(n)
+        i = 0
+        session.process { (i += 1) < n }
+      end
+  end
+
+end
\ No newline at end of file
diff --git a/test/transport/test_session.rb b/test/transport/test_session.rb
index 097e57f..76b4158 100644
--- a/test/transport/test_session.rb
+++ b/test/transport/test_session.rb
@@ -257,6 +257,13 @@ module Transport
       session.send_message(packet)
     end
 
+    def test_enqueue_message_should_delegate_to_socket
+      session!
+      packet = P(:byte, SERVICE_ACCEPT, :string, "test")
+      socket.expects(:enqueue_packet).with(packet)
+      session.enqueue_message(packet)
+    end
+
     def test_configure_client_should_pass_options_to_socket_client_state
       session.configure_client :compression => :standard
       assert_equal :standard, socket.client.compression