diff --git a/lib/postgres/postgres-pr/message.rb b/lib/postgres/postgres-pr/message.rb index 95f1c15c0dd4b..2607b0a992b7a 100644 --- a/lib/postgres/postgres-pr/message.rb +++ b/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 diff --git a/lib/rex/post/mysql/ui/console/command_dispatcher/client.rb b/lib/rex/post/mysql/ui/console/command_dispatcher/client.rb index 81cbd0b2033c3..5e42f3ecf09a5 100644 --- a/lib/rex/post/mysql/ui/console/command_dispatcher/client.rb +++ b/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 diff --git a/lib/rex/post/smb/ui/console.rb b/lib/rex/post/smb/ui/console.rb index 6cf64d332ac42..71a086a1533a8 100644 --- a/lib/rex/post/smb/ui/console.rb +++ b/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) diff --git a/lib/rex/post/sql/ui/console/command_dispatcher/client.rb b/lib/rex/post/sql/ui/console/command_dispatcher/client.rb index 5fffcaa26f16e..aca92d3b670e0 100644 --- a/lib/rex/post/sql/ui/console/command_dispatcher/client.rb +++ b/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