Land rapid7#9220, Module cache improvements

.gitignore

@@ -93,3 +93,7 @@ docker-compose.local*
 # Ignore python bytecode
 *.pyc
 rspec.failures
+
+
+#Ignore any base disk store files
+db/modules_metadata_base.pstore

lib/metasploit/framework/spec/constants.rb

@@ -22,6 +22,7 @@ module Metasploit::Framework::Spec::Constants
     Error
     External
     Loader
+    Metadata
     MetasploitClassCompatibilityError
     Namespace
     VersionCompatibilityError

lib/msf/core/framework.rb

@@ -233,24 +233,8 @@ def threads?
     }
   end
 
+  # TODO: Anything still using this should be ported to use metadata::cache search
   def search(match, logger: nil)
-    # Check if the database is usable
-    use_db = true
-    if self.db
-      if !(self.db.migrated && self.db.modules_cached)
-        logger.print_warning("Module database cache not built yet, using slow search") if logger
-        use_db = false
-      end
-    else
-      logger.print_warning("Database not connected, using slow search") if logger
-      use_db = false
-    end
-
-    # Used the database for search
-    if use_db
-      return self.db.search_modules(match)
-    end
-
     # Do an in-place search
     matches = []
     [ self.exploits, self.auxiliary, self.post, self.payloads, self.nops, self.encoders ].each do |mset|

lib/msf/core/module/full_name.rb

@@ -21,7 +21,7 @@ module ClassMethods
     #
 
     def fullname
-      type + '/' + refname
+      "#{type}/#{refname}"
     end
 
     def promptname

lib/msf/core/module_manager/cache.rb

@@ -3,6 +3,7 @@
 # Gems
 #
 require 'active_support/concern'
+require 'msf/core/modules/metadata/cache'
 
 # Concerns the module cache maintained by the {Msf::ModuleManager}.
 module Msf::ModuleManager::Cache
@@ -98,7 +99,7 @@ def load_cached_module(type, reference_name)
   end
 
   # @overload refresh_cache_from_module_files
-  #   Rebuilds database and in-memory cache for all modules.
+  #   Rebuilds module metadata store and in-memory cache for all modules.
   #
   #   @return [void]
   # @overload refresh_cache_from_module_files(module_class_or_instance)
@@ -107,14 +108,21 @@ def load_cached_module(type, reference_name)
   #   @param (see Msf::DBManager#update_module_details)
   #   @return [void]
   def refresh_cache_from_module_files(module_class_or_instance = nil)
-    if framework_migrated?
-      if module_class_or_instance
-        framework.db.update_module_details(module_class_or_instance)
-      else
-        framework.db.update_all_module_details
-      end
-      refresh_cache_from_database(self.module_paths)
+    if module_class_or_instance
+      Msf::Modules::Metadata::Cache.instance.refresh_metadata_instance(module_class_or_instance)
+    else
+      module_sets =
+          [
+              ['exploit', @framework.exploits],
+              ['auxiliary', @framework.auxiliary],
+              ['post', @framework.post],
+              ['payload', @framework.payloads],
+              ['encoder', @framework.encoders],
+              ['nop', @framework.nops]
+          ]
+      Msf::Modules::Metadata::Cache.instance.refresh_metadata(module_sets)
     end
+    refresh_cache_from_database(self.module_paths)
   end
 
   # Refreshes the in-memory cache from the database cache.
@@ -126,19 +134,11 @@ def refresh_cache_from_database(allowed_paths=[""])
 
   protected
 
-  # Returns whether the framework migrations have been run already.
-  #
-  # @return [true] if migrations have been run
-  # @return [false] otherwise
-  def framework_migrated?
-    framework.db && framework.db.migrated
-  end
-
   # @!attribute [rw] module_info_by_path
-  #   @return (see #module_info_by_path_from_database!)
+  #   @return (see #module_info_by_path_from_store!)
   attr_accessor :module_info_by_path
 
-  # Return a module info from Mdm::Module::Details in database.
+  # Return a module info from Msf::Modules::Metadata::Obj.
   #
   # @note Also sets module_set(module_type)[module_reference_name] to Msf::SymbolicModule if it is not already set.
   #
@@ -148,41 +148,35 @@ def framework_migrated?
   def module_info_by_path_from_database!(allowed_paths=[""])
     self.module_info_by_path = {}
 
-    if framework_migrated?
-      allowed_paths = allowed_paths.map{|x| x + "/"}
-
-      ActiveRecord::Base.connection_pool.with_connection do
-        # TODO record module parent_path in Mdm::Module::Detail so it does not need to be derived from file.
-        # Use find_each so Mdm::Module::Details are returned in batches, which will
-        # handle the growing number of modules better than all.each.
-        Mdm::Module::Detail.find_each do |module_detail|
-          path = module_detail.file
-          type = module_detail.mtype
-          reference_name = module_detail.refname
-
-          # Skip cached modules that are not in our allowed load paths
-          next if allowed_paths.select{|x| path.index(x) == 0}.empty?
-
-          # The load path is assumed to be the next level above the type directory
-          type_dir = File.join('', Mdm::Module::Detail::DIRECTORY_BY_TYPE[type], '')
-          parent_path = path.split(type_dir)[0..-2].join(type_dir) # TODO: rewrite
-
-          module_info_by_path[path] = {
-              :reference_name => reference_name,
-              :type => type,
-              :parent_path => parent_path,
-              :modification_time => module_detail.mtime
-          }
-
-          typed_module_set = module_set(type)
-
-          # Don't want to trigger as {Msf::ModuleSet#create} so check for
-          # key instead of using ||= which would call {Msf::ModuleSet#[]}
-          # which would potentially call {Msf::ModuleSet#create}.
-          unless typed_module_set.has_key? reference_name
-            typed_module_set[reference_name] = Msf::SymbolicModule
-          end
-        end
+    allowed_paths = allowed_paths.map{|x| x + "/"}
+
+    metadata = Msf::Modules::Metadata::Cache.instance.get_metadata
+    metadata.each do |module_metadata|
+      path = module_metadata.path
+      type = module_metadata.type
+      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?
+
+      # The load path is assumed to be the next level above the type directory
+      type_dir = File.join('', Mdm::Module::Detail::DIRECTORY_BY_TYPE[type], '')
+      parent_path = path.split(type_dir)[0..-2].join(type_dir) # TODO: rewrite
+
+      module_info_by_path[path] = {
+          :reference_name => reference_name,
+          :type => type,
+          :parent_path => parent_path,
+          :modification_time => module_metadata.mod_time
+      }
+
+      typed_module_set = module_set(type)
+
+      # Don't want to trigger as {Msf::ModuleSet#create} so check for
+      # key instead of using ||= which would call {Msf::ModuleSet#[]}
+      # which would potentially call {Msf::ModuleSet#create}.
+      unless typed_module_set.has_key? reference_name
+        typed_module_set[reference_name] = Msf::SymbolicModule
       end
     end
 

lib/msf/core/modules/metadata.rb

@@ -0,0 +1,8 @@
+# -*- coding: binary -*-
+require 'msf/core/modules'
+
+# Namespace for module metadata related data and operations
+module Msf::Modules::Metadata
+
+end
+

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

@@ -0,0 +1,124 @@
+require 'singleton'
+require 'msf/events'
+require 'rex/ui/text/output/stdio'
+require 'msf/core/constants'
+require 'msf/core/modules/metadata'
+require 'msf/core/modules/metadata/obj'
+require 'msf/core/modules/metadata/search'
+require 'msf/core/modules/metadata/store'
+
+#
+# Core service class that provides storage of module metadata as well as operations on the metadata.
+# Note that operations on this metadata are included as separate modules.
+#
+module Msf
+module Modules
+module Metadata
+
+class Cache
+  include Singleton
+  include Msf::Modules::Metadata::Search
+  include Msf::Modules::Metadata::Store
+
+  #
+  # Refreshes cached module metadata as well as updating the store
+  #
+  def refresh_metadata_instance(module_instance)
+    refresh_metadata_instance_internal(module_instance)
+    update_store
+  end
+
+  #
+  #  Returns the module data cache, but first ensures all the metadata is loaded
+  #
+  def get_metadata
+    wait_for_load
+    @module_metadata_cache.values
+  end
+
+  #
+  # Checks for modules loaded that are not a part of the cache and updates the underlying store
+  # if there are changes.
+  #
+  def refresh_metadata(module_sets)
+    unchanged_module_references = get_unchanged_module_references
+    has_changes = false
+    module_sets.each do |mt|
+      unchanged_reference_name_set = unchanged_module_references[mt[0]]
+
+      mt[1].keys.sort.each do |mn|
+        next if unchanged_reference_name_set.include? mn
+        module_instance = mt[1].create(mn)
+        next if not module_instance
+        begin
+          refresh_metadata_instance_internal(module_instance)
+          has_changes = true
+        rescue Exception => e
+          elog("Error updating module details for #{module_instance.fullname}: #{$!.class} #{$!} : #{e.message}")
+        end
+      end
+    end
+
+    update_store if has_changes
+  end
+
+  #
+  # Returns  a hash(type->set) which references modules that have not changed.
+  #
+  def get_unchanged_module_references
+    skip_reference_name_set_by_module_type = Hash.new { |hash, module_type|
+      hash[module_type] = Set.new
+    }
+
+    @module_metadata_cache.each_value do |module_metadata|
+
+      unless module_metadata.path and ::File.exist?(module_metadata.path)
+        next
+      end
+
+      if ::File.mtime(module_metadata.path).to_i != module_metadata.mod_time.to_i
+        next
+      end
+
+      skip_reference_name_set = skip_reference_name_set_by_module_type[module_metadata.type]
+      skip_reference_name_set.add(module_metadata.ref_name)
+    end
+
+    return skip_reference_name_set_by_module_type
+  end
+
+  #######
+  private
+  #######
+
+  def wait_for_load
+    @load_thread.join unless @store_loaded
+  end
+
+  def refresh_metadata_instance_internal(module_instance)
+    metadata_obj = Obj.new(module_instance)
+    @module_metadata_cache[get_cache_key(module_instance)] = metadata_obj
+  end
+
+  def get_cache_key(module_instance)
+    key = ''
+    key << (module_instance.type.nil? ? '' : module_instance.type)
+    key << '_'
+    key << module_instance.refname
+    return key
+  end
+
+  def initialize
+    @module_metadata_cache = {}
+    @store_loaded = false
+    @console = Rex::Ui::Text::Output::Stdio.new
+    @load_thread = Thread.new  {
+      init_store
+      @store_loaded = true
+    }
+  end
+end
+
+end
+end
+end

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

@@ -0,0 +1,71 @@
+require 'msf/core/modules/metadata'
+
+#
+# Simple object for storing a modules metadata.
+#
+module Msf
+module Modules
+module Metadata
+
+class Obj
+  attr_reader :name
+  attr_reader :full_name
+  attr_reader :rank
+  attr_reader :disclosure_date
+  attr_reader :type
+  attr_reader :author
+  attr_reader :description
+  attr_reader :references
+  attr_reader :is_server
+  attr_reader :is_client
+  attr_reader :platform
+  attr_reader :arch
+  attr_reader :rport
+  attr_reader :targets
+  attr_reader :mod_time
+  attr_reader :is_install_path
+  attr_reader :ref_name
+
+  def initialize(module_instance)
+    @name = module_instance.name
+    @full_name = module_instance.fullname
+    @disclosure_date = module_instance.disclosure_date
+    @rank = module_instance.rank.to_i
+    @type = module_instance.type
+    @description = module_instance.description.to_s.strip
+    @author = module_instance.author.map{|x| x.to_s}
+    @references = module_instance.references.map{|x| [x.ctx_id, x.ctx_val].join("-") }
+    @is_server = (module_instance.respond_to?(:stance) and module_instance.stance == "aggressive")
+    @is_client = (module_instance.respond_to?(:stance) and module_instance.stance == "passive")
+    @platform = module_instance.platform_to_s
+    @arch = module_instance.arch_to_s
+    @rport = module_instance.datastore['RPORT'].to_s
+    @path = module_instance.file_path
+    @mod_time = ::File.mtime(@path) rescue Time.now
+    @ref_name = module_instance.refname
+    install_path = Msf::Config.install_root.to_s
+    if (@path.to_s.include? (install_path))
+      @path = @path.sub(install_path, '')
+      @is_install_path = true
+    end
+
+    if module_instance.respond_to?(:targets) and module_instance.targets
+      @targets = module_instance.targets.map{|x| x.name}
+    end
+  end
+
+  def update_mod_time(mod_time)
+    @mod_time = mod_time
+  end
+
+  def path
+    if @is_install_path
+      return ::File.join(Msf::Config.install_root, @path)
+    end
+
+    @path
+  end
+end
+end
+end
+end