diff --git a/lib/msf/core/module_manager/cache.rb b/lib/msf/core/module_manager/cache.rb index 9e991d16be..b9dc9dbe07 100644 --- a/lib/msf/core/module_manager/cache.rb +++ b/lib/msf/core/module_manager/cache.rb @@ -177,7 +177,7 @@ module Msf::ModuleManager::Cache reference_name = module_metadata.ref_name # Skip cached modules that are not in our allowed load paths - next if allowed_paths.select{|x| path.index(x) == 0}.empty? + next unless allowed_paths.any? { |x| path.start_with?(x) } parent_path = get_parent_path(path, type) diff --git a/lib/msf/core/modules/metadata/cache.rb b/lib/msf/core/modules/metadata/cache.rb index 2f88c114f2..aa10e2005e 100644 --- a/lib/msf/core/modules/metadata/cache.rb +++ b/lib/msf/core/modules/metadata/cache.rb @@ -78,6 +78,9 @@ class Cache end end end + if has_changes + rebuild_type_cache + end } if has_changes update_store @@ -89,8 +92,8 @@ class Cache def module_metadata(type) @mutex.synchronize do wait_for_load - # TODO: Should probably figure out a way to cache this - @module_metadata_cache.filter_map { |_, metadata| [metadata.ref_name, metadata] if metadata.type == type }.to_h + type_hash = @metadata_type_index[type] + type_hash ? type_hash.dup : {} end end @@ -129,7 +132,9 @@ class Cache module_metadata.ref_name.eql? module_name } - return old_cache_size != @module_metadata_cache.size + removed = old_cache_size != @module_metadata_cache.size + rebuild_type_cache if removed + removed end def wait_for_load @@ -141,29 +146,50 @@ class Cache # Remove all instances of modules pointing to the same path. This prevents stale data hanging # around when modules are incorrectly typed (eg: Auxiliary that should be Exploit) + had_type_mismatch_deletion = false @module_metadata_cache.delete_if {|_, module_metadata| - module_metadata.path.eql? metadata_obj.path && module_metadata.type != module_metadata.type + is_stale = module_metadata.path.eql?(metadata_obj.path) && module_metadata.type != metadata_obj.type + had_type_mismatch_deletion = true if is_stale + is_stale } - @module_metadata_cache[get_cache_key(module_instance)] = metadata_obj + cache_key = get_cache_key(module_instance) + @module_metadata_cache[cache_key] = metadata_obj + + if had_type_mismatch_deletion + # Type changed - full rebuild needed since we removed entries from other type buckets + rebuild_type_cache + else + # Common case - just update the single entry in the type index + type_hash = (@metadata_type_index[metadata_obj.type] ||= {}) + type_hash[metadata_obj.ref_name] = metadata_obj + end end def get_cache_key(module_instance) - key = '' - key << (module_instance.type.nil? ? '' : module_instance.type) - key << '_' - key << module_instance.class.refname - return key + "#{module_instance.type}_#{module_instance.class.refname}" + end + + # Rebuild the per-type index from the main cache. + def rebuild_type_cache + by_type = {} + @module_metadata_cache.each_value do |metadata| + type_hash = (by_type[metadata.type] ||= {}) + type_hash[metadata.ref_name] = metadata + end + @metadata_type_index = by_type end def initialize super @mutex = Mutex.new @module_metadata_cache = {} + @metadata_type_index = {} @store_loaded = false @console = Rex::Ui::Text::Output::Stdio.new @load_thread = Thread.new { init_store + rebuild_type_cache @store_loaded = true } end diff --git a/spec/lib/msf/core/modules/metadata/cache_spec.rb b/spec/lib/msf/core/modules/metadata/cache_spec.rb new file mode 100644 index 0000000000..87a2bcf155 --- /dev/null +++ b/spec/lib/msf/core/modules/metadata/cache_spec.rb @@ -0,0 +1,202 @@ +# frozen_string_literal: true + +require 'spec_helper' + +RSpec.describe Msf::Modules::Metadata::Cache do + # Build a testable Cache instance without triggering the Singleton constructor + # (which spawns a thread and loads the store from disk). + let(:cache) do + obj = described_class.send(:allocate) + obj.instance_variable_set(:@mutex, Mutex.new) + obj.instance_variable_set(:@module_metadata_cache, {}) + obj.instance_variable_set(:@metadata_type_index, {}) + obj.instance_variable_set(:@store_loaded, true) + obj.instance_variable_set(:@load_thread, Thread.new {}) + obj + end + + def make_metadata(type:, ref_name:, path: '/modules/test.rb') + Msf::Modules::Metadata::Obj.from_hash({ + 'name' => ref_name, + 'fullname' => "#{type}/#{ref_name}", + 'rank' => 300, + 'type' => type, + 'author' => ['rspec'], + 'description' => 'Test module', + 'references' => [], + 'mod_time' => '2024-01-01 00:00:00 +0000', + 'path' => path, + 'is_install_path' => false, + 'ref_name' => ref_name + }) + end + + def populate_cache(cache, *entries) + entries.each do |entry| + cache.instance_variable_get(:@module_metadata_cache)["#{entry.type}_#{entry.ref_name}"] = entry + end + cache.send(:rebuild_type_cache) + end + + # Fake module instance for refresh_metadata_instance_internal + def make_module_instance(type:, refname:, path: '/modules/test.rb') + mod = double('module_instance') + klass = double('module_class', refname: refname) + allow(mod).to receive(:type).and_return(type) + allow(mod).to receive(:class).and_return(klass) + allow(mod).to receive(:refname).and_return(refname) + allow(mod).to receive(:realname).and_return("#{type}/#{refname}") + allow(mod).to receive(:name).and_return(refname) + allow(mod).to receive(:aliases).and_return([]) + allow(mod).to receive(:disclosure_date).and_return(nil) + allow(mod).to receive(:rank).and_return(300) + allow(mod).to receive(:description).and_return('Test') + allow(mod).to receive(:author).and_return([]) + allow(mod).to receive(:references).and_return([]) + allow(mod).to receive(:post_auth?).and_return(false) + allow(mod).to receive(:default_cred?).and_return(false) + allow(mod).to receive(:platform_to_s).and_return('') + allow(mod).to receive(:platform).and_return(nil) + allow(mod).to receive(:arch_to_s).and_return('') + allow(mod).to receive(:datastore).and_return({}) + allow(mod).to receive(:file_path).and_return(path) + allow(mod).to receive(:has_check?).and_return(false) + allow(mod).to receive(:notes).and_return({}) + + # Stub specific respond_to? checks used by Obj#initialize + allow(mod).to receive(:respond_to?).with(:needs_cleanup).and_return(false) + allow(mod).to receive(:respond_to?).with(:actions).and_return(false) + allow(mod).to receive(:respond_to?).with(:autofilter_ports).and_return(false) + allow(mod).to receive(:respond_to?).with(:autofilter_services).and_return(false) + allow(mod).to receive(:respond_to?).with(:targets).and_return(false) + allow(mod).to receive(:respond_to?).with(:session_types).and_return(false) + allow(mod).to receive(:respond_to?).with(:payload_type).and_return(false) + mod + end + + describe '#module_metadata' do + it 'returns modules of the requested type' do + exploit = make_metadata(type: 'exploit', ref_name: 'test/vuln') + auxiliary = make_metadata(type: 'auxiliary', ref_name: 'scanner/test', path: '/modules/aux.rb') + populate_cache(cache, exploit, auxiliary) + + result = cache.module_metadata('exploit') + expect(result.keys).to eq(['test/vuln']) + expect(result['test/vuln']).to eq(exploit) + end + + it 'returns an empty hash for unknown types' do + exploit = make_metadata(type: 'exploit', ref_name: 'test/vuln') + populate_cache(cache, exploit) + + expect(cache.module_metadata('post')).to eq({}) + end + + it 'returns a copy that does not affect internal state' do + exploit = make_metadata(type: 'exploit', ref_name: 'test/vuln') + populate_cache(cache, exploit) + + result = cache.module_metadata('exploit') + result.delete('test/vuln') + + expect(cache.module_metadata('exploit').keys).to eq(['test/vuln']) + end + end + + describe '#rebuild_type_cache' do + it 'groups all entries by type' do + e1 = make_metadata(type: 'exploit', ref_name: 'test/a', path: '/modules/a.rb') + e2 = make_metadata(type: 'exploit', ref_name: 'test/b', path: '/modules/b.rb') + aux = make_metadata(type: 'auxiliary', ref_name: 'scan/c', path: '/modules/c.rb') + populate_cache(cache, e1, e2, aux) + + expect(cache.module_metadata('exploit').size).to eq(2) + expect(cache.module_metadata('auxiliary').size).to eq(1) + end + end + + describe '#refresh_metadata_instance_internal' do + it 'adds a new module to the type index' do + mod = make_module_instance(type: 'exploit', refname: 'test/new', path: '/modules/new.rb') + cache.send(:rebuild_type_cache) + cache.send(:refresh_metadata_instance_internal, mod) + + result = cache.module_metadata('exploit') + expect(result.keys).to eq(['test/new']) + end + + it 'updates an existing module in the type index' do + old = make_metadata(type: 'exploit', ref_name: 'test/mod', path: '/modules/mod.rb') + populate_cache(cache, old) + + mod = make_module_instance(type: 'exploit', refname: 'test/mod', path: '/modules/mod.rb') + cache.send(:refresh_metadata_instance_internal, mod) + + result = cache.module_metadata('exploit') + expect(result.size).to eq(1) + expect(result['test/mod']).not_to eq(old) + end + + context 'when a module changes type' do + it 'removes the old type entry and adds to the new type' do + # Module starts as auxiliary + old_aux = make_metadata(type: 'auxiliary', ref_name: 'test/mistyped', path: '/modules/mistyped.rb') + other_aux = make_metadata(type: 'auxiliary', ref_name: 'scan/other', path: '/modules/other.rb') + populate_cache(cache, old_aux, other_aux) + + expect(cache.module_metadata('auxiliary').size).to eq(2) + expect(cache.module_metadata('exploit')).to eq({}) + + # Now refresh it as an exploit (same path, different type) + mod = make_module_instance(type: 'exploit', refname: 'test/mistyped', path: '/modules/mistyped.rb') + cache.send(:refresh_metadata_instance_internal, mod) + + # Old auxiliary entry should be gone, other auxiliary should remain + aux_result = cache.module_metadata('auxiliary') + expect(aux_result.size).to eq(1) + expect(aux_result.keys).to eq(['scan/other']) + + # New exploit entry should exist + exploit_result = cache.module_metadata('exploit') + expect(exploit_result.size).to eq(1) + expect(exploit_result.keys).to eq(['test/mistyped']) + end + + it 'does not leave stale entries in the main cache' do + old = make_metadata(type: 'auxiliary', ref_name: 'test/stale', path: '/modules/stale.rb') + populate_cache(cache, old) + + mod = make_module_instance(type: 'exploit', refname: 'test/stale', path: '/modules/stale.rb') + cache.send(:refresh_metadata_instance_internal, mod) + + main_cache = cache.instance_variable_get(:@module_metadata_cache) + types = main_cache.values.map(&:type).uniq + expect(types).to eq(['exploit']) + end + end + end + + describe '#remove_from_cache' do + it 'removes the named module and returns true' do + mod = make_metadata(type: 'exploit', ref_name: 'test/remove', path: '/modules/remove.rb') + populate_cache(cache, mod) + + result = cache.send(:remove_from_cache, 'test/remove') + expect(result).to be true + expect(cache.instance_variable_get(:@module_metadata_cache)).to be_empty + expect(cache.module_metadata('exploit')).to eq({}) + end + + it 'returns false when the module does not exist' do + result = cache.send(:remove_from_cache, 'test/nonexistent') + expect(result).to be false + end + end + + describe '#get_cache_key' do + it 'returns type_refname' do + mod = make_module_instance(type: 'exploit', refname: 'test/key') + expect(cache.send(:get_cache_key, mod)).to eq('exploit_test/key') + end + end +end