Handle SSH errors by closing the session

Author: sjanusz-r7

lib/msf/base/sessions/aws_instance_connect_command_shell_bind.rb

@@ -76,7 +76,7 @@ def desc
     end
 
     def bootstrap(datastore = {}, handler = nil)
-      @ssh_command_stream = Net::SSH::CommandStream.new(ssh_connection)
+      @ssh_command_stream = Net::SSH::CommandStream.new(ssh_connection, session: self)
 
       @ssh_command_stream.verify_channel
       # set remote_window_size to 32 which seems to help stability

lib/msf/base/sessions/ssh_command_shell_bind.rb

@@ -243,7 +243,7 @@ def bootstrap(datastore = {}, handler = nil)
       # shells accessed through SSH may respond to the echo command issued for verification as expected
       datastore['AutoVerifySession'] &= @platform.blank?
 
-      @rstream = Net::SSH::CommandStream.new(ssh_connection).lsock
+      @rstream = Net::SSH::CommandStream.new(ssh_connection, session: self).lsock
       super
 
       @info = "SSH #{username} @ #{@peer_info}"

lib/net/ssh/command_stream.rb

@@ -2,7 +2,7 @@
 
 class Net::SSH::CommandStream
 
-  attr_accessor :channel, :thread, :error, :ssh
+  attr_accessor :channel, :thread, :error, :ssh, :session
   attr_accessor :lsock, :rsock, :monitor
 
   module PeerInfo
@@ -13,7 +13,9 @@ module PeerInfo
 
   def shell_requested(channel, success)
     unless success
-      raise Net::SSH::ChannelRequestFailed, 'Shell/exec channel request failed'
+      error = Net::SSH::ChannelRequestFailed, 'Shell/exec channel request failed'
+      handle_error(error: error)
+      raise error
     end
 
     self.channel = channel
@@ -40,7 +42,8 @@ def shell_requested(channel, success)
     end
   end
 
-  def initialize(ssh, cmd = nil, pty: false, cleanup: false)
+  def initialize(ssh, cmd = nil, pty: false, cleanup: false, session: nil)
+    self.session = session
     self.lsock, self.rsock = Rex::Socket.tcp_socket_pair()
     self.lsock.extend(Rex::IO::Stream)
     self.lsock.extend(PeerInfo)
@@ -74,31 +77,42 @@ def initialize(ssh, cmd = nil, pty: false, cleanup: false)
       end
 
       channel.on_open_failed do |ch, code, desc|
-        raise Net::SSH::ChannelOpenFailed.new(code, 'Session channel open failed')
+        error = Net::SSH::ChannelOpenFailed.new(code, 'Session channel open failed')
+        handle_error(error: error)
+        raise error
       end
 
       self.monitor = Thread.new do
-        while(true)
-          next if not self.rsock.has_read_data?(1.0)
-          buff = self.rsock.read(16384)
-          break if not buff
-          verify_channel
-          self.channel.send_data(buff) if buff
+        begin
+          Kernel.loop do
+            next if not self.rsock.has_read_data?(1.0)
+
+            buff = self.rsock.read(16384)
+            break if not buff
+
+            verify_channel
+            self.channel.send_data(buff) if buff
+          end
+        rescue ::StandardError => e
+          handle_error(error: e)
         end
       end
 
-      while true
-        rssh.process(0.5) { true }
+      begin
+        Kernel.loop { rssh.process(0.5) { true } }
+      rescue ::StandardError => e
+        handle_error(error: e)
       end
 
       # Shut down the SSH session if requested
       if !rcmd.nil? && rcleanup
         rssh.close
       end
     end
+    self.thread.abort_on_exception = true
   rescue ::StandardError => e
     # XXX: This won't be set UNTIL there's a failure from a thread
-    self.error = e
+    handle_error(error: e)
   ensure
     self.monitor.kill if self.monitor
   end
@@ -119,6 +133,7 @@ def handle_error(error: nil)
   end
 
   def cleanup
+    self.session.alive = false if self.session
     self.monitor.kill
     self.lsock.close rescue nil
     self.rsock.close rescue nil