Improves handling of dying SMB and SQL sessions

lib/postgres/postgres-pr/message.rb

@@ -46,7 +46,11 @@ def self.read(stream, startup=false)
     type = stream.read_exactly_n_bytes(1) unless startup
     length = stream.read_exactly_n_bytes(4).to_s.unpack('N').first  # FIXME: length should be signed, not unsigned
 
-    raise ParseError if (length.nil? || length < 4)
+    if length.nil?
+      raise EOFError
+    elsif length < 4
+      raise ParseError
+    end
 
     # If we didn't read any bytes and startup was not set, then type will be nil, so don't continue.
     unless startup

lib/rex/post/mysql/ui/console/command_dispatcher/client.rb

@@ -39,6 +39,14 @@ def normalise_sql_result(result)
             # meaning we don't have that in the Result object.
             { rows: result.entries, columns: result.fields.each.map(&:name), errors: [] }
           end
+
+          def handle_error(e)
+            case e
+            when Mysql::ClientError::ServerLost
+              _close_session
+            end
+            super
+          end
         end
       end
     end

lib/rex/post/smb/ui/console.rb

@@ -97,6 +97,10 @@ def run_command(dispatcher, method, arguments)
             log_error(e.message)
           rescue ::Errno::EPIPE, ::OpenSSL::SSL::SSLError, ::IOError
             session.kill
+          rescue ::RubySMB::Error::CommunicationError => e
+            log_error("Error running command #{method}: #{e.class} #{e}")
+            elog(e)
+            session.alive = false
           rescue ::StandardError => e
             log_error("Error running command #{method}: #{e.class} #{e}")
             elog(e)

lib/rex/post/sql/ui/console/command_dispatcher/client.rb

@@ -120,7 +120,7 @@ def cmd_query_help
               def run_query(query)
                 begin
                   result = client.query(query)
-                rescue ::RuntimeError, ::StandardError => e
+                rescue => e
                   elog("Running query '#{query}' failed on session #{self.inspect}", error: e)
                   return { status: :error, result: { errors: [e] } }
                 end
@@ -177,6 +177,9 @@ def cmd_query(*args)
                   print_line(formatted_result.to_s)
                 when :error
                   print_error "Query has failed. Reasons: #{result[:result][:errors].join(', ')}"
+                  result[:result][:errors].each do |error|
+                    handle_error(error)
+                  end
                 else
                   elog "Unknown query status: #{result[:status]}"
                   print_error "Unknown query status: #{result[:status]}"
@@ -188,6 +191,29 @@ def process_query(query: '')
 
                 query.lines.each.map { |line| line.chomp.chomp('\\').strip }.reject(&:empty?).compact.join(' ')
               end
+
+              # Handles special cased error for each protocol
+              # that are return when a session has died resulting in a session
+              # needing to be closed
+              #
+              # @param [Object] e
+              def handle_error(e)
+                case e
+                when EOFError
+                  _close_session
+                end
+              end
+
+              private
+
+              # Sets session.alive to close sessions and handle setting session.client.interacting
+              # if currently in the context of the `query_interactive` command
+              #
+              # @return [FalseClass]
+              def _close_session
+                session.alive = false
+                session.client.interacting = false if session.client && session.client.respond_to?(:interacting)
+              end
             end
           end
         end