First pass at changes needed for module metadata caching

Author: clee-r7

db/modules_metadata_base.pstore

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

lib/metasploit/framework/spec/constants.rb

@@ -91,8 +91,6 @@ def initialize(options={})
     events.add_general_subscriber(subscriber)
     events.add_db_subscriber(subscriber)
     events.add_ui_subscriber(subscriber)
-
-    Msf::Modules::Metadata::Cache.instance.init(self)
   end
 
   def inspect

lib/msf/core/framework.rb

@@ -21,7 +21,7 @@ module ClassMethods
     #
 
     def fullname
-      type + '/' + (refname.nil? ? "" : refname)
+      type + '/' + (refname.nil? ? '' : refname)
     end
 
     def shortname

lib/msf/core/module/full_name.rb

@@ -98,7 +98,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 metadat store and in-memory cache for all modules.
   #
   #   @return [void]
   # @overload refresh_cache_from_module_files(module_class_or_instance)
@@ -107,14 +107,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 +133,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 +147,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/module_manager/cache.rb

@@ -205,8 +205,6 @@ def add_module(klass, reference_name, info = {})
 
     self[reference_name] = klass
 
-    Msf::Modules::Metadata::Cache.instance.cache_module_metadata(klass)
-
     klass
   end
 

lib/msf/core/module_set.rb

@@ -1,3 +1,7 @@
+# -*- coding: binary -*-
+
+# Namespace for module metadata related data and operations
 module Msf::Modules::Metadata
 
-end
+end
+

lib/msf/core/modules/metadata.rb

@@ -2,156 +2,120 @@
 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.
 #
-# To prevent excessive startup times module definitions are not parsed for metadata until startup
-# is complete. Once startup is complete to prevent CPU spin loading is then done gradually and
-# only when an operation using the cache is called is CPU use maximized.
-#
 module Msf
 module Modules
 module Metadata
 
 class Cache
   include Singleton
-  include ::Msf::UiEventSubscriber
   include Msf::Modules::Metadata::Search
+  include Msf::Modules::Metadata::Store
 
   #
-  # Init registers this class as a listener to be notified when the console is done loading,
-  # this acts a hint to this class to start trickle loading the metadata
+  # Refreshes cached module metadata as well as updating the store
   #
-  def init(framework)
-    register_ui_listener(framework)
-    @framework = framework
+  def refresh_metadata_instance(module_instance)
+    refresh_metadata_instance_internal(module_instance)
+    update_store
   end
 
   #
-  # Parses module metadata into an in memory cache.
-  #
-  # @param klass_or_instance - An instantiated or class instance of a module.
+  #  Returns the module data cache, but first ensures all the metadata is loaded
   #
-  def cache_module_metadata(klass_or_instance)
-    # Only use cache if db is not in use for now
-    return  if @framework.db.active
-
-    if klass_or_instance.is_a?(Class)
-      if @module_load_complete
-        add_module_metadata_from_class(klass_or_instance)
-      else
-        @module_definitions << klass_or_instance
-      end
-    else
-      add_module_metadata_from_instance(klass_or_instance)
-    end
+  def get_metadata
+    wait_for_load
+    @module_metadata_cache.values
   end
 
-  #########
-  protected
-  #########
-
   #
-  # Notify the thread responsible for loading metadata that it can start loading.
+  # Checks for modules loaded that are not a part of the cache and updates the underlying store
+  # if there are changes.
   #
-  def on_ui_start(rev)
-    return if @module_load_complete
-    @startup_called = true
+  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 the module data cache, but first ensures all the metadata is loaded
+  # Returns  a hash(type->set) which references modules that have not changed.
   #
-  def get_module_metadata_cache
-    wait_for_load
-    return @module_metadata_cache
+  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 register_ui_listener(framework)
-    begin
-      framework.events.add_ui_subscriber(self)
-    rescue Exception => e
-      elog('Unable to register metadata cache service as UI listener')
-    end
-  end
-
   def wait_for_load
-    if (!@module_load_complete)
-      @trickle_load = false
-      @console.print_warning('Waiting to finish parsing module metadata')
-      @module_load_thread.join
-    end
+    @load_thread.join unless @store_loaded
   end
 
-  def add_module_metadata_from_class(module_class_definition)
-    begin
-      instance = module_class_definition.new
-      add_module_metadata_from_instance(instance)
-    rescue Exception => e
-      elog("Error adding module metadata: #{e.message}")
-    end
-  end
-
-  def add_module_metadata_from_instance(module_instance)
-    module_metadata = Obj.new(module_instance)
-    @module_metadata_cache[get_cache_key(module_instance)] = module_metadata
+  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.name
+    key << module_instance.refname
     return key
   end
 
-  #
-  # This method is used by the @module_load_thread
-  #
-  def load_module_definitions
-    loop do
-      if @startup_called
-        break;
-      end
-
-      sleep 0.3
-    end
-
-    count = 0
-    @module_definitions.each {|module_definition|
-      add_module_metadata_from_class(module_definition)
-      count = count + 1
-      if (@trickle_load && count > @trickle_load_batch)
-        sleep @trickle_load_interval
-        count = 0
-      end
-    }
-
-    @module_load_complete = true
-    GC.start(full_mark: true, immediate_sweep: true)
-  end
-
   def initialize
-    @module_load_complete = false
-    @startup_called = false;
-    @trickle_load = true
-    @trickle_load_batch = 200
-    @trickle_load_interval = 0.5
     @module_metadata_cache = {}
-    @module_definitions = []
-    @module_load_thread = Thread.new {
-      load_module_definitions
-    }
-
+    @store_loaded = false
     @console = Rex::Ui::Text::Output::Stdio.new
+    @load_thread = Thread.new  {
+      init_store
+      @store_loaded = true
+    }
   end
 end