Skip to content

Commit

Permalink
Improves handling of dying SMB and SQL sessions
Browse files Browse the repository at this point in the history
  • Loading branch information
cgranleese-r7 committed Mar 20, 2024
1 parent d88185b commit 33ad5f6
Show file tree
Hide file tree
Showing 4 changed files with 44 additions and 2 deletions.
6 changes: 5 additions & 1 deletion lib/postgres/postgres-pr/message.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
8 changes: 8 additions & 0 deletions lib/rex/post/mysql/ui/console/command_dispatcher/client.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 4 additions & 0 deletions lib/rex/post/smb/ui/console.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
28 changes: 27 additions & 1 deletion lib/rex/post/sql/ui/console/command_dispatcher/client.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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]}"
Expand All @@ -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
Expand Down

0 comments on commit 33ad5f6

Please sign in to comment.