CoCalc provides the best real-time collaborative environment for Jupyter Notebooks, LaTeX documents, and SageMath, scalable from individual users to large groups and classes!

1
require 'pathname'
2

3
# @note needs to use explicit nesting. so this file can be loaded directly without loading 'metasploit/framework' which
4
#   allows for faster loading of rake tasks.
5
module Metasploit
6
  module Framework
7
    module Spec
8
      module Threads
9
        module Suite
10
          #
11
          # CONSTANTS
12
          #
13

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

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

40
          #
41
          # Module Methods
42
          #
43

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

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

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

69
                  LOG_PATHNAME.parent.mkpath
70

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

78
                config.after(:suite) do
79
                  LOG_PATHNAME.parent.mkpath
80

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

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

91
                  if thread_count > EXPECTED_THREAD_COUNT_AROUND_SUITE
92
                    error_lines = []
93

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

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

101
                        # unmanaged thread, such as the main VM thread
102
                        unless thread_uuid
103
                          next
104
                        end
105

106
                        caller = caller_by_thread_uuid[thread_uuid]
107

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

117
                    raise RuntimeError,
118
                          "#{thread_count} #{'thread'.pluralize(thread_count)} exist(s) when only " \
119
                          "#{EXPECTED_THREAD_COUNT_AROUND_SUITE} " \
120
                          "#{'thread'.pluralize(EXPECTED_THREAD_COUNT_AROUND_SUITE)} expected after suite runs:\n" \
121
                          "#{error_lines.join}"
122
                  end
123
                end
124
              end
125

126
              @configured = true
127
            end
128

129
            @configured
130
          end
131

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

138
              parent_pathname = Pathname.new(__FILE__).parent
139
              threads_logger_pathname = parent_pathname.join('logger')
140
              load_pathname = parent_pathname.parent.parent.parent.parent.expand_path
141

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

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

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

159
            LOG_PATHNAME.each_line do |line|
160
              if in_suite
161
                if line.start_with?('after(:suite)')
162
                  break
163
                else
164
                  yield line
165
                end
166
              else
167
                if line.start_with?('before(:suite)')
168
                  in_suite = true
169
                end
170
              end
171
            end
172
          end
173

174
          # @note Ensure {LOG_PATHNAME} exists before calling.
175
          #
176
          # Yield each line for each Thread UUID gathered during the suite run.
177
          #
178
          # @yield [uuid, line]
179
          # @yieldparam uuid [String] the UUID of thread thread
180
          # @yieldparam line [String] a line in the `caller` for the given `uuid`
181
          # @yieldreturn [void]
182
          def self.each_thread_line
183
            in_thread_caller = false
184
            uuid = nil
185

186
            each_suite_line do |line|
187
              if in_thread_caller
188
                if line.start_with?('END Thread.new caller')
189
                  in_thread_caller = false
190
                  next
191
                else
192
                  yield uuid, line
193
                end
194
              else
195
                match = line.match(UUID_REGEXP)
196

197
                if match
198
                  in_thread_caller = true
199
                  uuid = match[:uuid]
200
                end
201
              end
202
            end
203
          end
204

205
          # The `caller` for each Thread UUID.
206
          #
207
          # @return [Hash{String => Array<String>}]
208
          def self.caller_by_thread_uuid
209
            lines_by_thread_uuid = Hash.new { |hash, uuid|
210
              hash[uuid] = []
211
            }
212

213
            each_thread_line do |uuid, line|
214
              lines_by_thread_uuid[uuid] << line
215
            end
216

217
            lines_by_thread_uuid
218
          end
219

220
          # @return
221
          def self.non_debugger_thread_list
222
            Thread.list.reject { |thread|
223
              # don't do `is_a? Debugger::DebugThread` because it requires Debugger::DebugThread to be loaded, which it
224
              # won't when not debugging.
225
              thread.class.name == 'Debugger::DebugThread' ||
226
                thread.class.name == 'Debase::DebugThread'
227
            }
228
          end
229
        end
230
      end
231
    end
232
  end
233
end
234

235