Merge pull request #21231 from bcoles/msf-module-cache

Module metadata: Fix stale module detection and add per-type metadata index

Author: adfoster-r7

lib/msf/core/module_manager/cache.rb

@@ -177,7 +177,7 @@ def module_info_by_path_from_database!(allowed_paths=[""])
       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)
 

lib/msf/core/modules/metadata/cache.rb

@@ -78,6 +78,9 @@ def refresh_metadata(module_sets)
           end
         end
       end
+      if has_changes
+        rebuild_type_cache
+      end
     }
     if has_changes
       update_store
@@ -89,8 +92,8 @@ def refresh_metadata(module_sets)
   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 @@ def remove_from_cache(module_name)
       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 @@ def refresh_metadata_instance_internal(module_instance)
 
     # 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

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