Skip to content

Commit

Permalink
Improve database module cache performance
Browse files Browse the repository at this point in the history
  • Loading branch information
adfoster-r7 committed Oct 10, 2024
1 parent 953f6c1 commit 93e0ca7
Show file tree
Hide file tree
Showing 3 changed files with 80 additions and 20 deletions.
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

0 comments on commit 93e0ca7

Please sign in to comment.