Skip to content

Commit 2edbc6a

Browse files
Land rapid7#19546, Improve database module cache performance
2 parents 90066b3 + 93e0ca7 commit 2edbc6a

File tree

3 files changed

+80
-20
lines changed

3 files changed

+80
-20
lines changed

lib/msf/core/db_manager/module_cache.rb

Lines changed: 68 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ def match_values(values)
2828
values.collect { |value| "%#{value}%" }
2929
end
3030

31-
def module_to_details_hash(m)
31+
def module_to_details_hash(m, with_mixins: true)
3232
res = {}
3333
bits = []
3434

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

9494

95-
m.class.mixins.each do |x|
96-
bits << [ :mixin, { :name => x.to_s } ]
95+
if with_mixins
96+
m.class.mixins.each do |x|
97+
bits << [ :mixin, { :name => x.to_s } ]
98+
end
9799
end
98100
end
99101

@@ -269,7 +271,6 @@ def update_all_module_details
269271
}
270272

271273
Mdm::Module::Detail.find_each do |md|
272-
273274
unless md.ready
274275
refresh << md
275276
next
@@ -291,6 +292,7 @@ def update_all_module_details
291292

292293
refresh.each { |md| md.destroy }
293294

295+
new_modules = []
294296
[
295297
['exploit', framework.exploits],
296298
['auxiliary', framework.auxiliary],
@@ -305,14 +307,12 @@ def update_all_module_details
305307
next if skip_reference_name_set.include? mn
306308
obj = mt[1].create(mn)
307309
next if not obj
308-
begin
309-
update_module_details(obj)
310-
rescue ::Exception => e
311-
elog("Error updating module details for #{obj.fullname}", error: e)
312-
end
310+
new_modules <<= obj
313311
end
314312
end
315313

314+
insert_all(new_modules)
315+
316316
self.framework.cache_initialized = true
317317
end
318318

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

334334
ApplicationRecord.connection_pool.with_connection do
335-
info = module_to_details_hash(module_instance)
335+
info = module_to_details_hash(module_instance, with_mixins: false)
336336
bits = info.delete(:bits) || []
337337
module_detail = Mdm::Module::Detail.create!(info)
338338

@@ -359,4 +359,62 @@ def update_module_details(module_instance)
359359
module_detail.save!
360360
end
361361
end
362+
363+
private
364+
365+
# Insert the Msf::Module array into the Mdm::Module::Detail database class
366+
#
367+
# @param [Array<Msf::Module>] modules
368+
def insert_all(modules)
369+
module_hashes = modules.filter_map do |mod|
370+
begin
371+
hash = module_to_details_hash(mod, with_mixins: false)
372+
# The insert_all API requires all hashes to have the same keys present, so explicitly set these potentially missing keys
373+
hash[:disclosure_date] ||= nil
374+
hash[:default_target] ||= nil
375+
hash[:default_action] ||= nil
376+
hash[:stance] ||= nil
377+
hash
378+
rescue ::Exception => e
379+
elog("Error updating module details for #{mod.fullname}", error: e)
380+
nil
381+
end
382+
end
383+
return if module_hashes.empty?
384+
385+
# 1) Bulk insert the module detail entries
386+
module_details = module_hashes.map { |mod_hash| mod_hash.except(:bits) }
387+
module_detail_ids = Mdm::Module::Detail.insert_all!(module_details, returning: %w[id]).map { |returning| returning['id'] }
388+
389+
# 2) Build the hashes for the associations
390+
associations = module_hashes.zip(module_detail_ids).each_with_object(Hash.new { |hash, key| hash[key] = [] }) do |(module_hash, detail_id), acc|
391+
module_hash[:bits].each do |args|
392+
otype, vals = args
393+
394+
case otype
395+
when :action
396+
acc[Mdm::Module::Action] << { detail_id: detail_id, name: vals[:name] }
397+
when :arch
398+
acc[Mdm::Module::Arch] << { detail_id: detail_id, name: vals[:name] }
399+
when :author
400+
acc[Mdm::Module::Author] << { detail_id: detail_id, name: vals[:name], email: vals[:email] }
401+
when :platform
402+
acc[Mdm::Module::Platform] << { detail_id: detail_id, name: vals[:name] }
403+
when :ref
404+
acc[Mdm::Module::Ref] << { detail_id: detail_id, name: vals[:name] }
405+
when :target
406+
acc[Mdm::Module::Target] << { detail_id: detail_id, index: vals[:index], name: vals[:name] }
407+
end
408+
end
409+
end
410+
411+
# 3) Insert all of the associations
412+
associations.each do |association_clazz, entries|
413+
next if entries.empty?
414+
415+
association_clazz.insert_all!(entries)
416+
end
417+
418+
nil
419+
end
362420
end

spec/support/shared/examples/msf/db_manager/module_cache.rb

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -812,7 +812,8 @@ def loader.load_error(module_path, error)
812812
allow(db_manager).to receive(
813813
:module_to_details_hash
814814
).with(
815-
module_instance
815+
module_instance,
816+
with_mixins: false
816817
).and_return(
817818
module_to_details_hash
818819
)

spec/support/shared/examples/msf/db_manager/update_all_module_details_refresh.rb

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -29,19 +29,20 @@
2929
update_all_module_details
3030
end
3131

32-
it 'should call update_module_details to create a new Mdm::Module::Detail from the module instance returned by create' do
33-
expect(db_manager).to receive(:update_module_details) do |module_instance|
34-
expect(module_instance).to be_a Msf::Module
35-
expect(module_instance.type).to eq module_detail.mtype
36-
expect(module_instance.refname).to eq module_detail.refname
37-
end
38-
32+
it 'should create a new Mdm::Module::Detail entry' do
3933
update_all_module_details
34+
35+
aggregate_failures do
36+
expect(Mdm::Module::Detail.count).to eq 1
37+
db_module_detail = Mdm::Module::Detail.first
38+
expect(db_module_detail.mtype).to eq(module_detail.mtype)
39+
expect(db_module_detail.refname).to eq(module_detail.refname)
40+
end
4041
end
4142

42-
context 'with exception raised by #update_module_details' do
43+
context 'with exception raised by #insert_all' do
4344
before(:example) do
44-
expect(db_manager).to receive(:update_module_details).and_raise(Exception)
45+
expect(db_manager).to receive(:module_to_details_hash).and_raise(Exception)
4546
end
4647

4748
it 'should log error' do

0 commit comments

Comments
 (0)