From 93e0ca7cd5cebf50157c9f4271beec6466bd0bde Mon Sep 17 00:00:00 2001 From: adfoster-r7 Date: Wed, 9 Oct 2024 13:17:48 +0100 Subject: [PATCH] Improve database module cache performance --- lib/msf/core/db_manager/module_cache.rb | 78 ++++++++++++++++--- .../examples/msf/db_manager/module_cache.rb | 3 +- .../update_all_module_details_refresh.rb | 19 ++--- 3 files changed, 80 insertions(+), 20 deletions(-) diff --git a/lib/msf/core/db_manager/module_cache.rb b/lib/msf/core/db_manager/module_cache.rb index 1c26e3025144..c42d15128016 100644 --- a/lib/msf/core/db_manager/module_cache.rb +++ b/lib/msf/core/db_manager/module_cache.rb @@ -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 = [] @@ -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 @@ -269,7 +271,6 @@ def update_all_module_details } Mdm::Module::Detail.find_each do |md| - unless md.ready refresh << md next @@ -291,6 +292,7 @@ def update_all_module_details refresh.each { |md| md.destroy } + new_modules = [] [ ['exploit', framework.exploits], ['auxiliary', framework.auxiliary], @@ -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 @@ -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) @@ -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] 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 diff --git a/spec/support/shared/examples/msf/db_manager/module_cache.rb b/spec/support/shared/examples/msf/db_manager/module_cache.rb index be886a56d5dc..7cc1f4bcb82e 100644 --- a/spec/support/shared/examples/msf/db_manager/module_cache.rb +++ b/spec/support/shared/examples/msf/db_manager/module_cache.rb @@ -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 ) diff --git a/spec/support/shared/examples/msf/db_manager/update_all_module_details_refresh.rb b/spec/support/shared/examples/msf/db_manager/update_all_module_details_refresh.rb index 88a99cb63125..17dbae67d08a 100644 --- a/spec/support/shared/examples/msf/db_manager/update_all_module_details_refresh.rb +++ b/spec/support/shared/examples/msf/db_manager/update_all_module_details_refresh.rb @@ -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