2 files changed, 27 insertions, 22 deletions

diff --git a/lib/net/ssh/connection/keepalive.rb b/lib/net/ssh/connection/keepalive.rb
index c66a9db..242071d 100644
--- a/lib/net/ssh/connection/keepalive.rb
+++ b/lib/net/ssh/connection/keepalive.rb
@@ -1,51 +1,53 @@
 require 'net/ssh/loggable'
-
 module Net; module SSH; module Connection
 
-module Keepalive
+class Keepalive
   include Loggable
 
-  # Default IO.select timeout threshold
-  DEFAULT_IO_SELECT_TIMEOUT = 300
-
-  def initialize_keepalive
+  def initialize(session)
     @last_keepalive_sent_at = nil
     @unresponded_keepalive_count = 0
+    @session = session
+    self.logger = session.logger
+  end
+
+  def options
+    @session.options
   end
 
-  def keepalive_enabled?
+  def enabled?
     options[:keepalive]
   end
 
-  def keepalive_interval
-    options[:keepalive_interval] || DEFAULT_IO_SELECT_TIMEOUT
+  def interval
+    options[:keepalive_interval] || Session::DEFAULT_IO_SELECT_TIMEOUT
   end
 
-  def should_send_keepalive?
-    return false unless keepalive_enabled?
+  def should_send?
+    return false unless enabled?
     return true unless @last_keepalive_sent_at
-    Time.now - @last_keepalive_sent_at >= keepalive_interval
+    Time.now - @last_keepalive_sent_at >= interval
   end
 
   def keepalive_maxcount
     (options[:keepalive_maxcount] || 3).to_i
   end
 
-  def send_keepalive_as_needed(readers, writers)
+  def send_as_needed(readers, writers)
     return unless readers.nil? && writers.nil?
-    return unless should_send_keepalive?
+    return unless should_send?
     info { "sending keepalive #{@unresponded_keepalive_count}" }
 
     @unresponded_keepalive_count += 1
-    send_global_request("keepalive@openssh.com") { |success, response|
+    @session.send_global_request("keepalive@openssh.com") { |success, response|
       debug { "keepalive response successful. Missed #{@unresponded_keepalive_count-1} keepalives" }
       @unresponded_keepalive_count = 0
     }
     @last_keepalive_sent_at = Time.now
     if keepalive_maxcount > 0 && @unresponded_keepalive_count > keepalive_maxcount
-      error { "Timeout, server #{host} not responding. Missed #{@unresponded_keepalive_count-1} timeouts." }
+      error { "Timeout, server #{@session.host} not responding. Missed #{@unresponded_keepalive_count-1} timeouts." }
       @unresponded_keepalive_count = 0
-      raise Net::SSH::Timeout, "Timeout, server #{host} not responding."
+      raise Net::SSH::Timeout, "Timeout, server #{@session.host} not responding."
     end
   end
 end
diff --git a/lib/net/ssh/connection/session.rb b/lib/net/ssh/connection/session.rb
index 3137514..aebd9e2 100644
--- a/lib/net/ssh/connection/session.rb
+++ b/lib/net/ssh/connection/session.rb
@@ -24,7 +24,10 @@ module Net; module SSH; module Connection
   #     ssh.exec! "/etc/init.d/some_process start"
   #   end
   class Session
-    include Constants, Loggable, Keepalive
+    include Constants, Loggable
+
+    # Default IO.select timeout threshold
+    DEFAULT_IO_SELECT_TIMEOUT = 300
 
     # The underlying transport layer abstraction (see Net::SSH::Transport::Session).
     attr_reader :transport
@@ -77,7 +80,7 @@ module Net; module SSH; module Connection
       @max_pkt_size = (options.has_key?(:max_pkt_size) ? options[:max_pkt_size] : 0x8000)
       @max_win_size = (options.has_key?(:max_win_size) ? options[:max_win_size] : 0x20000)
 
-      initialize_keepalive
+      @keepalive = Keepalive.new(self)
     end
 
     # Retrieves a custom property from this instance. This can be used to
@@ -242,7 +245,7 @@ module Net; module SSH; module Connection
         writer.send_pending
       end
 
-      send_keepalive_as_needed(readers, writers)
+      @keepalive.send_as_needed(readers, writers)
       transport.rekey_as_needed
 
       return true
@@ -596,8 +599,8 @@ module Net; module SSH; module Connection
 
       def io_select_wait(wait)
         return wait if wait
-        return wait unless keepalive_enabled?
-        keepalive_interval
+        return wait unless @keepalive.enabled?
+        @keepalive.interval
       end
 
       MAP = Constants.constants.inject({}) do |memo, name|