Merge branch 'landing/4207' into upstream-master

Land rapid7#4207

* Ensure that `rake spec` doesn't create too many threads

Author: trevrosen

lib/metasploit/framework.rb

@@ -35,6 +35,7 @@ module Framework
     extend ActiveSupport::Autoload
 
     autoload :Spec
+    autoload :ThreadFactoryProvider
 
     # Returns the root of the metasploit-framework project.  Use in place of
     # `Rails.root`.

lib/metasploit/framework/thread_factory_provider.rb

@@ -0,0 +1,26 @@
+# Wraps {Msf::Framework} so that {Msf::Framework#threads} is only created on the first call to {#spawn} by
+# {Rex::ThreadFactory#spawn}, which allows the threads used by {Msf::ThreadManager} to be created lazily.
+#
+# @example Setting Rex::ThreadFactory.provider and spawning threads
+#   Rex::ThreadFactory.provider = Metasploit::Framework::ThreadFactoryProvider.new(framework: framework)
+#   # framework.threads created here
+#   Rex::ThreadFactory.spawn("name", false) { ... }
+#
+class Metasploit::Framework::ThreadFactoryProvider < Metasploit::Model::Base
+  #
+  # Attributes
+  #
+
+  # @!attribute framework
+  #   The framework managing the spawned threads.
+  #
+  #   @return [Msf::Framework]
+  attr_accessor :framework
+
+  # Spawns a thread monitored by {Msf::ThreadManager} in {Msf::Framework#threads}.
+  #
+  # (see Msf::ThreadManager#spawn)
+  def spawn(name, critical, *args, &block)
+    framework.threads.spawn(name, critical, *args, &block)
+  end
+end

lib/msf/core/framework.rb

@@ -1,4 +1,15 @@
 # -*- coding: binary -*-
+
+#
+# Standard Library
+#
+
+require 'monitor'
+
+#
+# Project
+#
+
 require 'metasploit/framework/version'
 require 'msf/core'
 require 'msf/util'
@@ -12,6 +23,7 @@ module Msf
 #
 ###
 class Framework
+  include MonitorMixin
 
   #
   # Versioning information
@@ -66,22 +78,22 @@ module Offspring
   #
   # Creates an instance of the framework context.
   #
-  def initialize(opts={})
+  def initialize(options={})
+    self.options = options
+    # call super to initialize MonitorMixin.  #synchronize won't work without this.
+    super()
 
     # Allow specific module types to be loaded
-    types = opts[:module_types] || Msf::MODULE_TYPES
+    types = options[:module_types] || Msf::MODULE_TYPES
 
-    self.threads   = ThreadManager.new(self)
     self.events    = EventDispatcher.new(self)
     self.modules   = ModuleManager.new(self,types)
-    self.sessions  = SessionManager.new(self)
     self.datastore = DataStore.new
     self.jobs      = Rex::JobContainer.new
     self.plugins   = PluginManager.new(self)
-    self.db        = DBManager.new(self, opts)
 
     # Configure the thread factory
-    Rex::ThreadFactory.provider = self.threads
+    Rex::ThreadFactory.provider = Metasploit::Framework::ThreadFactoryProvider.new(framework: self)
 
     subscriber = FrameworkEventSubscriber.new(self)
     events.add_exploit_subscriber(subscriber)
@@ -155,11 +167,6 @@ def version
   #
   attr_reader   :modules
   #
-  # Session manager that tracks sessions associated with this framework
-  # instance over the course of their lifetime.
-  #
-  attr_reader   :sessions
-  #
   # The global framework datastore that can be used by modules.
   #
   attr_reader   :datastore
@@ -180,28 +187,62 @@ def version
   # unloading of plugins.
   #
   attr_reader   :plugins
-  #
+
   # The framework instance's db manager. The db manager
   # maintains the database db and handles db events
   #
-  attr_reader   :db
+  # @return [Msf::DBManager]
+  def db
+    synchronize {
+      @db ||= Msf::DBManager.new(self, options)
+    }
+  end
+
+  # Session manager that tracks sessions associated with this framework
+  # instance over the course of their lifetime.
   #
+  # @return [Msf::SessionManager]
+  def sessions
+    synchronize {
+      @sessions ||= Msf::SessionManager.new(self)
+    }
+  end
+
   # The framework instance's thread manager. The thread manager
   # provides a cleaner way to manage spawned threads
   #
-  attr_reader   :threads
+  # @return [Msf::ThreadManager]
+  def threads
+    synchronize {
+      @threads ||= Msf::ThreadManager.new(self)
+    }
+  end
+
+  # Whether {#threads} has been initialized
+  #
+  # @return [true] if {#threads} has been initialized
+  # @return [false] otherwise
+  def threads?
+    synchronize {
+      instance_variable_defined? :@threads
+    }
+  end
 
 protected
 
+  # @!attribute options
+  #   Options passed to {#initialize}
+  #
+  #   @return [Hash]
+  attr_accessor :options
+
   attr_writer   :events # :nodoc:
   attr_writer   :modules # :nodoc:
-  attr_writer   :sessions # :nodoc:
   attr_writer   :datastore # :nodoc:
   attr_writer   :auxmgr # :nodoc:
   attr_writer   :jobs # :nodoc:
   attr_writer   :plugins # :nodoc:
   attr_writer   :db # :nodoc:
-  attr_writer   :threads # :nodoc:
 end
 
 class FrameworkEventSubscriber

spec/lib/msf/core/framework_spec.rb

@@ -4,7 +4,17 @@
 require 'msf/core/framework'
 
 describe Msf::Framework do
-  include_context 'Msf::Framework#threads cleaner'
+  context '#initialize' do
+    subject(:framework) {
+      described_class.new
+    }
+
+    it 'creates no threads' do
+      expect {
+        framework
+      }.not_to change { Thread.list.count }
+    end
+  end
 
   describe "#version" do
     CURRENT_VERSION = "4.10.1-dev"

spec/msfcli_spec.rb

@@ -143,7 +143,6 @@ def get_stdout(&block)
   #
   context "#dump_module_list" do
     include_context 'Metasploit::Framework::Spec::Constants cleaner'
-    include_context 'Msf::Framework#threads cleaner'
 
     let(:framework) {
       msfcli.framework
@@ -490,7 +489,6 @@ def get_stdout(&block)
 
   context "#init_modules" do
     include_context 'Metasploit::Framework::Spec::Constants cleaner'
-    include_context 'Msf::Framework#threads cleaner'
 
     let(:args) {
       [
@@ -645,7 +643,6 @@ def get_stdout(&block)
 
   context "#engage_mode" do
     include_context 'Metasploit::Framework::Spec::Constants cleaner'
-    include_context 'Msf::Framework#threads cleaner'
 
     subject(:engage_mode) {
       msfcli.engage_mode(modules)

spec/support/shared/contexts/msf/framework/threads/cleaner.rb

@@ -1,5 +1,13 @@
 shared_context 'Msf::Framework#threads cleaner' do
-  after(:each) do
+  after(:each) do |example|
+    unless framework.threads?
+      fail RuntimeError.new(
+               "framework.threads was never initialized. There are no threads to clean up. " \
+               "Remove `include_context Msf::Framework#threads cleaner` from context around " \
+               "'#{example.metadata.full_description}'"
+           )
+    end
+
     # explicitly kill threads so that they don't exhaust connection pool
     thread_manager = framework.threads
 

spec/support/shared/contexts/msf/simple/framework.rb

@@ -3,8 +3,6 @@
 require 'metasploit/framework'
 
 shared_context 'Msf::Simple::Framework' do
-  include_context 'Msf::Framework#threads cleaner'
-
   let(:dummy_pathname) do
     Rails.root.join('spec', 'dummy')
   end