Implemented keepalive_maxcount

4 files changed, 40 insertions, 6 deletions

diff --git a/lib/net/ssh.rb b/lib/net/ssh.rb
index d9a1f57..d4eb307 100644
--- a/lib/net/ssh.rb
+++ b/lib/net/ssh.rb
@@ -63,7 +63,7 @@ module Net
     VALID_OPTIONS = [
       :auth_methods, :bind_address, :compression, :compression_level, :config,
       :encryption, :forward_agent, :hmac, :host_key,
-      :keepalive, :keepalive_interval, :kex, :keys, :key_data,
+      :keepalive, :keepalive_interval, :keepalive_maxcount, :kex, :keys, :key_data,
       :languages, :logger, :paranoid, :password, :port, :proxy,
       :rekey_blocks_limit,:rekey_limit, :rekey_packet_limit, :timeout, :verbose,
       :global_known_hosts_file, :user_known_hosts_file, :host_key_alias,
@@ -130,6 +130,7 @@ module Net
     #   the keepalive_interval seconds. Defaults to +false+.
     #   :keepalive_interval => the interval seconds for keepalive.
     #   Defaults to +300+ seconds.
+    #   :keepalive_countmax => the maximun number of keepalive packet miss allowed.
     # * :kex => the key exchange algorithm (or algorithms) to use
     # * :keys => an array of file names of private keys to use for publickey
     #   and hostbased authentication
diff --git a/lib/net/ssh/connection/session.rb b/lib/net/ssh/connection/session.rb
index 9be74e8..12631a1 100644
--- a/lib/net/ssh/connection/session.rb
+++ b/lib/net/ssh/connection/session.rb
@@ -80,6 +80,7 @@ module Net; module SSH; module Connection
       @max_win_size = (options.has_key?(:max_win_size) ? options[:max_win_size] : 0x20000)
 
       @last_keepalive_sent_at = nil
+      @unresponded_keepalive_count = 0
     end
 
     # Retrieves a custom property from this instance. This can be used to
@@ -612,12 +613,23 @@ module Net; module SSH; module Connection
         Time.now - @last_keepalive_sent_at >= keepalive_interval
       end
 
+      def keepalive_maxcount
+        (options[:keepalive_maxcount] || 3).to_i
+      end
+
       def send_keepalive_as_needed(readers, writers)
         return unless readers.nil? && writers.nil?
         return unless should_send_keepalive?
-        info { "sending keepalive" }
-        msg = Net::SSH::Buffer.from(:byte, Packet::IGNORE, :string, "keepalive")
-        send_message(msg)
+        info { "sending keepalive #{@unresponded_keepalive_count}" }
+
+        @unresponded_keepalive_count += 1
+        send_global_request("keepalive@openssh.com") { |success, response|
+          @unresponded_keepalive_count = 0
+        }
+        if keepalive_maxcount > 0 && @unresponded_keepalive_count > keepalive_maxcount
+          error { "Timeout, server #{host} not responding. Missed #{@unresponded_keepalive_count} timeouts." }
+          raise Net::SSH::Timeout, "Timeout, server #{host} not responding."
+        end
         @last_keepalive_sent_at = Time.now
       end
 
diff --git a/lib/net/ssh/errors.rb b/lib/net/ssh/errors.rb
index 8f5a136..910c961 100644
--- a/lib/net/ssh/errors.rb
+++ b/lib/net/ssh/errors.rb
@@ -14,6 +14,10 @@ module Net; module SSH
   # unexpectedly.
   class Disconnect < Exception; end
 
+  # This exception is raised when the remote host has disconnected/
+  # timeouted unexpectedly.
+  class Timeout < Disconnect; end
+
   # This exception is primarily used internally, but if you have a channel
   # request handler (see Net::SSH::Connection::Channel#on_request) that you
   # want to fail in such a way that the server knows it failed, you can
diff --git a/test/connection/test_session.rb b/test/connection/test_session.rb
index 8af6cd3..268afb9 100644
--- a/test/connection/test_session.rb
+++ b/test/connection/test_session.rb
@@ -366,12 +366,29 @@ module Connection
     def test_process_should_call_enqueue_message_if_io_select_timed_out
       timeout = Net::SSH::Connection::Session::DEFAULT_IO_SELECT_TIMEOUT
       options = { :keepalive => true }
-      expected_packet = P(:byte, Net::SSH::Packet::IGNORE, :string, "keepalive")
+      expected_packet = P(:byte, Net::SSH::Packet::GLOBAL_REQUEST, :string, "keepalive@openssh.com", :bool, true)
       IO.stubs(:select).with([socket],[],nil,timeout).returns(nil)
-      transport.expects(:enqueue_message).with{ |msg| msg.content == expected_packet.content  }
+      transport.expects(:enqueue_message).with{ |msg| msg.content == expected_packet.content }
       session(options).process
     end
 
+    def test_process_should_raise_if_keepalives_not_answered
+      timeout = Net::SSH::Connection::Session::DEFAULT_IO_SELECT_TIMEOUT
+      options = { :keepalive => true, :keepalive_interval => 300, :keepalive_maxcount => 3 }
+      expected_packet = P(:byte, Net::SSH::Packet::GLOBAL_REQUEST, :string, "keepalive@openssh.com", :bool, true)
+      [1,2,3].each do |i|
+        Time.stubs(:now).returns(i*300)
+        IO.stubs(:select).with([socket],[],nil,timeout).returns(nil)
+        transport.expects(:enqueue_message).with{ |msg| msg.content == expected_packet.content  }
+        session(options).process
+      end
+
+      Time.stubs(:now).returns(4*300)
+      IO.stubs(:select).with([socket],[],nil,timeout).returns(nil)
+      transport.expects(:enqueue_message).with{ |msg| msg.content == expected_packet.content }
+      assert_raises(Net::SSH::Timeout) { session(options).process }
+    end
+
     def test_process_should_not_call_enqueue_message_unless_io_select_timed_out
       timeout = Net::SSH::Connection::Session::DEFAULT_IO_SELECT_TIMEOUT
       options = { :keepalive => true }