Mirror of metasploit
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

thread_manager.rb 4.9KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187
  1. # -*- coding: binary -*-
  2. require 'msf/core/plugin'
  3. =begin
  4. require 'active_record'
  5. #
  6. # This monkeypatch can help to diagnose errors involving connection pool
  7. # exhaustion and other strange ActiveRecord including errors like:
  8. #
  9. # DEPRECATION WARNING: Database connections will not be closed automatically, please close your
  10. # database connection at the end of the thread by calling `close` on your
  11. # connection. For example: ActiveRecord::Base.connection.close
  12. #
  13. # and
  14. #
  15. # ActiveRecord::StatementInvalid NoMethodError: undefined method `fields' for nil:NilClass: SELECT "workspaces".* FROM "workspaces" WHERE "workspaces"."id" = 24 LIMIT 1
  16. #
  17. #
  18. # Based on this code: https://gist.github.com/1364551 linked here:
  19. # http://bibwild.wordpress.com/2011/11/14/multi-threading-in-rails-activerecord-3-0-3-1/
  20. module ActiveRecord
  21. class Base
  22. class << self
  23. def connection
  24. unless connection_pool.active_connection?
  25. $stdout.puts("AR::B.connection implicit checkout")
  26. $stdout.puts(caller.join("\n"))
  27. raise ImplicitConnectionForbiddenError.new("Implicit ActiveRecord checkout attempted!")
  28. end
  29. retrieve_connection
  30. end
  31. end
  32. end
  33. class ImplicitConnectionForbiddenError < ActiveRecord::ConnectionTimeoutError ; end
  34. end
  35. =end
  36. module Msf
  37. ###
  38. #
  39. # This class manages the threads spawned by the framework object, this provides some additional
  40. # features over standard ruby threads.
  41. #
  42. ###
  43. class ThreadManager < Array
  44. include Framework::Offspring
  45. attr_accessor :monitor
  46. #
  47. # Initializes the thread manager.
  48. #
  49. def initialize(framework)
  50. self.framework = framework
  51. self.monitor = spawn_monitor
  52. end
  53. #
  54. # Spawns a monitor thread for removing dead threads
  55. #
  56. def spawn_monitor
  57. ::Thread.new do
  58. begin
  59. ::Thread.current[:tm_name] = "Thread Monitor"
  60. ::Thread.current[:tm_crit] = true
  61. while true
  62. ::IO.select(nil, nil, nil, 1.0)
  63. self.each_index do |i|
  64. state = self[i].alive? rescue false
  65. self[i] = nil if not state
  66. end
  67. self.delete(nil)
  68. end
  69. rescue ::Exception => e
  70. elog("thread monitor: #{e} #{e.backtrace} source:#{self[:tm_call].inspect}")
  71. end
  72. end
  73. end
  74. #
  75. # Spawns a new thread
  76. #
  77. def spawn(name, crit, *args, &block)
  78. t = nil
  79. if block
  80. t = ::Thread.new(name, crit, caller, block, *args) do |*argv|
  81. ::Thread.current[:tm_name] = argv.shift.to_s
  82. ::Thread.current[:tm_crit] = argv.shift
  83. ::Thread.current[:tm_call] = argv.shift
  84. ::Thread.current[:tm_time] = Time.now
  85. begin
  86. argv.shift.call(*argv)
  87. rescue ::Exception => e
  88. elog(
  89. "thread exception: #{::Thread.current[:tm_name]} critical=#{::Thread.current[:tm_crit]} " \
  90. "error: #{e.class} #{e}\n" \
  91. " source:\n" \
  92. " #{::Thread.current[:tm_call].join "\n "}"
  93. )
  94. elog("Call Stack\n#{e.backtrace.join("\n")}")
  95. raise e
  96. ensure
  97. if framework.db and framework.db.active
  98. # NOTE: despite the Deprecation Warning's advice, this should *NOT*
  99. # be ActiveRecord::Base.connection.close which causes unrelated
  100. # threads to raise ActiveRecord::StatementInvalid exceptions at
  101. # some point in the future, presumably due to the pool manager
  102. # believing that the connection is still usable and handing it out
  103. # to another thread.
  104. ::ActiveRecord::Base.connection_pool.release_connection
  105. end
  106. end
  107. end
  108. else
  109. t = ::Thread.new(name, crit, caller, *args) do |*argv|
  110. ::Thread.current[:tm_name] = argv.shift
  111. ::Thread.current[:tm_crit] = argv.shift
  112. ::Thread.current[:tm_call] = argv.shift
  113. ::Thread.current[:tm_time] = Time.now
  114. # Calling spawn without a block means we cannot force a database
  115. # connection release when the thread completes, so doing so can
  116. # potentially use up all database resources and starve all subsequent
  117. # threads that make use of the database. Log a warning so we can track
  118. # down this kind of usage.
  119. dlog("Thread spawned without a block!")
  120. dlog("Call stack: \n#{::Thread.current[:tm_call].join("\n")}")
  121. end
  122. end
  123. self << t
  124. t
  125. end
  126. #
  127. # Registers an existing thread
  128. #
  129. def register(t, name, crit)
  130. t[:tm_name] = name
  131. t[:tm_crit] = crit
  132. t[:tm_call] = caller
  133. t[:tm_time] = Time.now
  134. self << t
  135. t
  136. end
  137. #
  138. # Updates an existing thread
  139. #
  140. def update(ut, name, crit)
  141. ti = nil
  142. self.each_index do |i|
  143. tt = self[i]
  144. next if not tt
  145. if ut.__id__ == tt.__id__
  146. ti = i
  147. break
  148. end
  149. end
  150. t = self[ti]
  151. if not t
  152. raise RuntimeError, "Thread not found"
  153. end
  154. t[:tm_name] = name
  155. t[:tm_crit] = crit
  156. t
  157. end
  158. #
  159. # Kills a thread by index
  160. #
  161. def kill(idx)
  162. self[idx].kill rescue false
  163. end
  164. end
  165. end