Make channel open failures work via a callback to parallel how channel open successes work

5 files changed, 44 insertions, 6 deletions

diff --git a/CHANGELOG.rdoc b/CHANGELOG.rdoc
index debd223..6b869c9 100644
--- a/CHANGELOG.rdoc
+++ b/CHANGELOG.rdoc
@@ -1,3 +1,8 @@
+=== *unreleased*
+
+* Make channel open failure work with a callback so that failures can be handled similarly to successes [Jamis Buck]
+
+
 === 2.0 Preview Release 2 (1.99.1) / 22 Mar 2008
 
 * Partial support for ~/.ssh/config (and related) SSH configuration files [Daniel J. Berger, Jamis Buck]
diff --git a/lib/net/ssh/connection/channel.rb b/lib/net/ssh/connection/channel.rb
index 0bf61db..c52e833 100644
--- a/lib/net/ssh/connection/channel.rb
+++ b/lib/net/ssh/connection/channel.rb
@@ -124,7 +124,7 @@ module Net; module SSH; module Connection
       @properties = {}
 
       @pending_requests = []
-      @on_data = @on_extended_data = @on_process = @on_close = @on_eof = nil
+      @on_open_failed = @on_data = @on_extended_data = @on_process = @on_close = @on_eof = nil
       @on_request = {}
       @closing = @eof = false
     end
@@ -392,6 +392,21 @@ module Net; module SSH; module Connection
       old
     end
 
+    # Registers a callback to be invoked when the server was unable to open
+    # the requested channel. The channel itself will be passed to the block,
+    # along with the integer "reason code" for the failure, and a textual
+    # description of the failure from the server.
+    #
+    #   channel = session.open_channel do |ch|
+    #     # ..
+    #   end
+    #
+    #   channel.on_open_failed { |ch, code, desc| ... }
+    def on_open_failed(&block)
+      old, @on_open_failed = @on_open_failed, block
+      old
+    end
+
     # Registers a callback to be invoked when a channel request of the given
     # type is received. The callback will receive the channel as the first
     # argument, and the associated (unparsed) data as the second. The data
@@ -493,6 +508,18 @@ module Net; module SSH; module Connection
         @on_confirm_open.call(self) if @on_confirm_open
       end
 
+      # Invoked when the server failed to open the channel. If an #on_open_failed
+      # callback was specified, it will be invoked with the channel, reason code,
+      # and description as arguments. Otherwise, a ChannelOpenFailed exception
+      # will be raised.
+      def do_open_failed(reason_code, description)
+        if @on_open_failed
+          @on_open_failed.call(self, reason_code, description)
+        else
+          raise ChannelOpenFailed.new(reason_code, description)
+        end          
+      end
+
       # Invoked when the server sends a CHANNEL_WINDOW_ADJUST packet, and
       # causes the remote window size to be adjusted upwards by the given
       # number of bytes. This has the effect of allowing more data to be sent
diff --git a/lib/net/ssh/connection/session.rb b/lib/net/ssh/connection/session.rb
index ff69de5..b13d926 100644
--- a/lib/net/ssh/connection/session.rb
+++ b/lib/net/ssh/connection/session.rb
@@ -509,8 +509,8 @@ module Net; module SSH; module Connection
 
       def channel_open_failure(packet)
         error { "channel_open_failed: #{packet[:local_id]} #{packet[:reason_code]} #{packet[:description]}" }
-        channels.delete(packet[:local_id])
-        raise ChannelOpenFailed.new(packet[:reason_code], packet[:description])
+        channel = channels.delete(packet[:local_id])
+        channel.do_open_failed(packet[:reason_code], packet[:description])
       end
 
       def channel_window_adjust(packet)
diff --git a/lib/net/ssh/service/forward.rb b/lib/net/ssh/service/forward.rb
index c10c3cf..4310418 100644
--- a/lib/net/ssh/service/forward.rb
+++ b/lib/net/ssh/service/forward.rb
@@ -68,10 +68,15 @@ module Net; module SSH; module Service
         client = socket.accept
         debug { "received connection on #{bind_address}:#{local_port}" }
 
-        session.open_channel("direct-tcpip", :string, remote_host, :long, remote_port, :string, bind_address, :long, local_port) do |channel|
+        channel = session.open_channel("direct-tcpip", :string, remote_host, :long, remote_port, :string, bind_address, :long, local_port) do |channel|
           channel.info { "direct channel established" }
           prepare_client(client, channel, :local)
         end
+
+        channel.on_open_failed do |ch, code, description|
+          channel.error { "could not establish direct channel: #{description} (#{code})" }
+          client.close
+        end
       end
     end
 
diff --git a/test/connection/test_session.rb b/test/connection/test_session.rb
index f4c0265..394e18b 100644
--- a/test/connection/test_session.rb
+++ b/test/connection/test_session.rb
@@ -237,10 +237,11 @@ module Connection
       assert_equal "auth-agent", ch.type
     end
 
-    def test_channel_open_failure_should_remove_channel_raise_exception
+    def test_channel_open_failure_should_remove_channel_and_tell_channel_that_open_failed
       session.channels[1] = stub("channel")
+      session.channels[1].expects(:do_open_failed).with(1234, "some reason")
       transport.return(CHANNEL_OPEN_FAILURE, :long, 1, :long, 1234, :string, "some reason", :string, "lang tag")
-      assert_raises(Net::SSH::ChannelOpenFailed) { process_times(2) }
+      process_times(2)
       assert session.channels.empty?
     end