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

@@ -17,6 +17,23 @@ def name
             'MySQL Client'
           end
 
+          # @param [Object] query
+          # @return [Hash{Symbol->Symbol | Hash{Symbol->Array<Exception>}}]
+          def run_query(query)
+            begin
+              result = client.query(query)
+            rescue Mysql::ClientError::ServerLost => e
+              elog("Running query '#{query}' failed on session #{self.inspect}", error: e)
+              session.alive = false
+              return { status: :error, result: { errors: [e] } }
+            rescue ::RuntimeError, ::StandardError => e
+              elog("Running query '#{query}' failed on session #{self.inspect}", error: e)
+              return { status: :error, result: { errors: [e] } }
+            end
+
+            super
+          end
+
           # @return [Object]
           def cmd_query_help
             print_line 'Usage: query'

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::EncryptionError => 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,6 +120,10 @@ def cmd_query_help
               def run_query(query)
                 begin
                   result = client.query(query)
+                rescue EOFError => e
+                  elog("Running query '#{query}' failed on session #{self.inspect}", error: e)
+                  session.alive = false
+                  return { status: :error, result: { errors: [e] } }
                 rescue ::RuntimeError, ::StandardError => e
                   elog("Running query '#{query}' failed on session #{self.inspect}", error: e)
                   return { status: :error, result: { errors: [e] } }