metasploit-gs/lib/metasploit/framework/spec/threads/suite.rb

require 'pathname'

# @note needs to use explicit nesting. so this file can be loaded directly without loading 'metasploit/framework' which
#   allows for faster loading of rake tasks.
module Metasploit
  module Framework
    module Spec
      module Threads
        module Suite
          #
          # CONSTANTS
          #

          # Number of allowed threads when threads are counted in `after(:suite)` or `before(:suite)`
          #
          # Known threads:
          #   1. Main Ruby thread
          #   2. Active Record connection pool thread
          #   3. Framework thread manager, a monitor thread for removing dead threads
          #      https://github.com/rapid7/metasploit-framework/blame/04e8752b9b74cbaad7cb0ea6129c90e3172580a2/lib/msf/core/thread_manager.rb#L66-L89
          #   4. Ruby's Timeout library thread, an automatically created monitor thread when using `Thread.timeout(1) { }`
          #      https://github.com/ruby/timeout/blob/bd25f4b138b86ef076e6d9d7374b159fffe5e4e9/lib/timeout.rb#L129-L137
          #   5. REMOTE_DB thread, if enabled
          #
          # Intermittent threads that are non-deterministically left behind, which should be fixed in the future:
          #   1. metadata cache hydration
          #      https://github.com/rapid7/metasploit-framework/blob/115946cd06faccac654e956e8ba9cf72ff328201/lib/msf/core/modules/metadata/cache.rb#L150-L153
          #   2. session manager
          #      https://github.com/rapid7/metasploit-framework/blob/115946cd06faccac654e956e8ba9cf72ff328201/lib/msf/core/session_manager.rb#L153-L168
          #
          EXPECTED_THREAD_COUNT_AROUND_SUITE = ENV['REMOTE_DB'] ? 7 : 6

          # `caller` for all Thread.new calls
          LOG_PATHNAME = Pathname.new('log/metasploit/framework/spec/threads/suite.log')
          # Regular expression for extracting the UUID out of {LOG_PATHNAME} for each Thread.new caller block
          UUID_REGEXP = /BEGIN Thread.new caller \((?<uuid>.*)\)/
          # Name of thread local variable that Thread UUID is stored
          UUID_THREAD_LOCAL_VARIABLE = "metasploit/framework/spec/threads/logger/uuid"

          #
          # Module Methods
          #

          # Configures `before(:suite)` and `after(:suite)` callback to detect thread leaks.
          #
          # @return [void]
          def self.configure!
            unless @configured
              RSpec.configure do |config|
                config.before(:suite) do
                  thread_count = Metasploit::Framework::Spec::Threads::Suite.non_debugger_thread_list.count

                  # check with if first so that error message can be constructed lazily
                  if thread_count > EXPECTED_THREAD_COUNT_AROUND_SUITE
                    # LOG_PATHNAME may not exist if suite run without `rake spec`
                    if LOG_PATHNAME.exist?
                      log = LOG_PATHNAME.read()
                    else
                      log "Run `rake spec` to log where Thread.new is called."
                    end

                    raise RuntimeError,
                          "#{thread_count} #{'thread'.pluralize(thread_count)} exist(s) when " \
                          "only #{EXPECTED_THREAD_COUNT_AROUND_SUITE} " \
                          "#{'thread'.pluralize(EXPECTED_THREAD_COUNT_AROUND_SUITE)} expected before suite runs:\n" \
                          "#{log}"
                  end

                  LOG_PATHNAME.parent.mkpath

                  LOG_PATHNAME.open('a') do |f|
                    # separator so after(:suite) can differentiate between threads created before(:suite) and during the
                    # suites
                    f.puts 'before(:suite)'
                  end
                end

                config.after(:suite) do
                  LOG_PATHNAME.parent.mkpath

                  LOG_PATHNAME.open('a') do |f|
                    # separator so that a flip flop can be used when reading the file below.  Also useful if it turns
                    # out any threads are being created after this callback, which could be the case if another
                    # after(:suite) accidentally created threads by creating an Msf::Simple::Framework instance.
                    f.puts 'after(:suite)'
                  end

                  thread_list = Metasploit::Framework::Spec::Threads::Suite.non_debugger_thread_list
                  thread_count = thread_list.count

                  if thread_count > EXPECTED_THREAD_COUNT_AROUND_SUITE
                    error_lines = []

                    if LOG_PATHNAME.exist?
                      caller_by_thread_uuid = Metasploit::Framework::Spec::Threads::Suite.caller_by_thread_uuid

                      thread_list.each do |thread|
                        thread_uuid = thread[Metasploit::Framework::Spec::Threads::Suite::UUID_THREAD_LOCAL_VARIABLE]
                        thread_name = thread[:tm_name]

                        # unmanaged thread, such as the main VM thread
                        unless thread_uuid
                          next
                        end

                        caller = caller_by_thread_uuid[thread_uuid]

                        error_lines << "Thread #{thread_uuid}'s (name=#{thread_name} status is #{thread.status.inspect} " \
                                       "and was started here:\n"
                        error_lines.concat(caller)
                        error_lines << "The thread backtrace was:\n#{thread.backtrace ? thread.backtrace.join("\n") : 'nil (no backtrace)'}\n"
                      end
                    else
                      error_lines << "Run `rake spec` to log where Thread.new is called."
                    end

                    raise RuntimeError,
                          "#{thread_count} #{'thread'.pluralize(thread_count)} exist(s) when only " \
                          "#{EXPECTED_THREAD_COUNT_AROUND_SUITE} " \
                          "#{'thread'.pluralize(EXPECTED_THREAD_COUNT_AROUND_SUITE)} expected after suite runs:\n" \
                          "#{error_lines.join}"
                  end
                end
              end

              @configured = true
            end

            @configured
          end

          def self.define_task
            Rake::Task.define_task('metasploit:framework:spec:threads:suite') do
              if Metasploit::Framework::Spec::Threads::Suite::LOG_PATHNAME.exist?
                Metasploit::Framework::Spec::Threads::Suite::LOG_PATHNAME.delete
              end

              parent_pathname = Pathname.new(__FILE__).parent
              threads_logger_pathname = parent_pathname.join('logger')
              load_pathname = parent_pathname.parent.parent.parent.parent.expand_path

              # Must append to RUBYOPT or Rubymine debugger will not work
              ENV['RUBYOPT'] = "#{ENV['RUBYOPT']} -I#{load_pathname} -r#{threads_logger_pathname}"
            end

            Rake::Task.define_task(spec: 'metasploit:framework:spec:threads:suite')
          end

          # @note Ensure {LOG_PATHNAME} exists before calling.
          #
          # Yields each line of {LOG_PATHNAME} that happened during the suite run.
          #
          # @yield [line]
          # @yieldparam line [String] a line in the {LOG_PATHNAME} between `before(:suite)` and `after(:suite)`
          # @yieldreturn [void]
          def self.each_suite_line
            in_suite = false

            LOG_PATHNAME.each_line do |line|
              if in_suite
                if line.start_with?('after(:suite)')
                  break
                else
                  yield line
                end
              else
                if line.start_with?('before(:suite)')
                  in_suite = true
                end
              end
            end
          end

          # @note Ensure {LOG_PATHNAME} exists before calling.
          #
          # Yield each line for each Thread UUID gathered during the suite run.
          #
          # @yield [uuid, line]
          # @yieldparam uuid [String] the UUID of thread thread
          # @yieldparam line [String] a line in the `caller` for the given `uuid`
          # @yieldreturn [void]
          def self.each_thread_line
            in_thread_caller = false
            uuid = nil

            each_suite_line do |line|
              if in_thread_caller
                if line.start_with?('END Thread.new caller')
                  in_thread_caller = false
                  next
                else
                  yield uuid, line
                end
              else
                match = line.match(UUID_REGEXP)

                if match
                  in_thread_caller = true
                  uuid = match[:uuid]
                end
              end
            end
          end

          # The `caller` for each Thread UUID.
          #
          # @return [Hash{String => Array<String>}]
          def self.caller_by_thread_uuid
            lines_by_thread_uuid = Hash.new { |hash, uuid|
              hash[uuid] = []
            }

            each_thread_line do |uuid, line|
              lines_by_thread_uuid[uuid] << line
            end

            lines_by_thread_uuid
          end

          # @return
          def self.non_debugger_thread_list
            Thread.list.reject { |thread|
              # don't do `is_a? Debugger::DebugThread` because it requires Debugger::DebugThread to be loaded, which it
              # won't when not debugging.
              thread.class.name == 'Debugger::DebugThread' ||
                thread.class.name == 'Debase::DebugThread'
            }
          end
        end
      end
    end
  end
end