Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Improve database module cache performance #19546

Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
78 changes: 68 additions & 10 deletions lib/msf/core/db_manager/module_cache.rb
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ def match_values(values)
values.collect { |value| "%#{value}%" }
end

def module_to_details_hash(m)
def module_to_details_hash(m, with_mixins: true)
res = {}
bits = []

Expand Down Expand Up @@ -92,8 +92,10 @@ def module_to_details_hash(m)
res[:stance] = m.stance.to_s.index("aggressive") ? "aggressive" : "passive"


m.class.mixins.each do |x|
bits << [ :mixin, { :name => x.to_s } ]
if with_mixins
m.class.mixins.each do |x|
bits << [ :mixin, { :name => x.to_s } ]
end
end
end

Expand Down Expand Up @@ -269,7 +271,6 @@ def update_all_module_details
}

Mdm::Module::Detail.find_each do |md|

unless md.ready
refresh << md
next
Expand All @@ -291,6 +292,7 @@ def update_all_module_details

refresh.each { |md| md.destroy }

new_modules = []
[
['exploit', framework.exploits],
['auxiliary', framework.auxiliary],
Expand All @@ -305,14 +307,12 @@ def update_all_module_details
next if skip_reference_name_set.include? mn
obj = mt[1].create(mn)
next if not obj
begin
update_module_details(obj)
rescue ::Exception => e
elog("Error updating module details for #{obj.fullname}", error: e)
end
new_modules <<= obj
end
end

insert_all(new_modules)

self.framework.cache_initialized = true
end

Expand All @@ -332,7 +332,7 @@ def update_module_details(module_instance)
return if not self.migrated

ApplicationRecord.connection_pool.with_connection do
info = module_to_details_hash(module_instance)
info = module_to_details_hash(module_instance, with_mixins: false)
bits = info.delete(:bits) || []
module_detail = Mdm::Module::Detail.create!(info)

Expand All @@ -359,4 +359,62 @@ def update_module_details(module_instance)
module_detail.save!
end
end

private

# Insert the Msf::Module array into the Mdm::Module::Detail database class
#
# @param [Array<Msf::Module>] modules
def insert_all(modules)
module_hashes = modules.filter_map do |mod|
begin
hash = module_to_details_hash(mod, with_mixins: false)
# The insert_all API requires all hashes to have the same keys present, so explicitly set these potentially missing keys
hash[:disclosure_date] ||= nil
hash[:default_target] ||= nil
hash[:default_action] ||= nil
hash[:stance] ||= nil
hash
rescue ::Exception => e
elog("Error updating module details for #{mod.fullname}", error: e)
nil
end
end
return if module_hashes.empty?

# 1) Bulk insert the module detail entries
module_details = module_hashes.map { |mod_hash| mod_hash.except(:bits) }
module_detail_ids = Mdm::Module::Detail.insert_all!(module_details, returning: %w[id]).map { |returning| returning['id'] }

# 2) Build the hashes for the associations
associations = module_hashes.zip(module_detail_ids).each_with_object(Hash.new { |hash, key| hash[key] = [] }) do |(module_hash, detail_id), acc|
module_hash[:bits].each do |args|
otype, vals = args

case otype
when :action
acc[Mdm::Module::Action] << { detail_id: detail_id, name: vals[:name] }
when :arch
acc[Mdm::Module::Arch] << { detail_id: detail_id, name: vals[:name] }
when :author
acc[Mdm::Module::Author] << { detail_id: detail_id, name: vals[:name], email: vals[:email] }
when :platform
acc[Mdm::Module::Platform] << { detail_id: detail_id, name: vals[:name] }
when :ref
acc[Mdm::Module::Ref] << { detail_id: detail_id, name: vals[:name] }
when :target
acc[Mdm::Module::Target] << { detail_id: detail_id, index: vals[:index], name: vals[:name] }
end
end
end

# 3) Insert all of the associations
associations.each do |association_clazz, entries|
next if entries.empty?

association_clazz.insert_all!(entries)
end

nil
end
end
3 changes: 2 additions & 1 deletion spec/support/shared/examples/msf/db_manager/module_cache.rb
Original file line number Diff line number Diff line change
Expand Up @@ -812,7 +812,8 @@ def loader.load_error(module_path, error)
allow(db_manager).to receive(
:module_to_details_hash
).with(
module_instance
module_instance,
with_mixins: false
).and_return(
module_to_details_hash
)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -29,19 +29,20 @@
update_all_module_details
end

it 'should call update_module_details to create a new Mdm::Module::Detail from the module instance returned by create' do
expect(db_manager).to receive(:update_module_details) do |module_instance|
expect(module_instance).to be_a Msf::Module
expect(module_instance.type).to eq module_detail.mtype
expect(module_instance.refname).to eq module_detail.refname
end

it 'should create a new Mdm::Module::Detail entry' do
update_all_module_details

aggregate_failures do
expect(Mdm::Module::Detail.count).to eq 1
db_module_detail = Mdm::Module::Detail.first
expect(db_module_detail.mtype).to eq(module_detail.mtype)
expect(db_module_detail.refname).to eq(module_detail.refname)
end
end

context 'with exception raised by #update_module_details' do
context 'with exception raised by #insert_all' do
before(:example) do
expect(db_manager).to receive(:update_module_details).and_raise(Exception)
expect(db_manager).to receive(:module_to_details_hash).and_raise(Exception)
end

it 'should log error' do
Expand Down
Loading