Fix msfrpc hanging when updating saved command history

rapid7 · Sep 15, 2023 · d4793f3 · d4793f3
1 parent 1b29c48
commit d4793f3
Show file tree

Hide file tree

Showing 7 changed files with 49 additions and 52 deletions.
diff --git a/lib/metasploit/framework/spec/threads/suite.rb b/lib/metasploit/framework/spec/threads/suite.rb
@@ -96,6 +96,7 @@ def self.configure!
 
                       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
@@ -104,10 +105,10 @@ def self.configure!
 
                         caller = caller_by_thread_uuid[thread_uuid]
 
-                        error_lines << "Thread #{thread_uuid}'s status is #{thread.status.inspect} " \
+                        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."

diff --git a/lib/msf/base/sessions/command_shell.rb b/lib/msf/base/sessions/command_shell.rb
@@ -546,7 +546,7 @@ def cmd_irb(*args)
     if expressions.empty?
       print_status('Starting IRB shell...')
       print_status("You are in the \"self\" (session) object\n")
-      Rex::Ui::Text::Shell::HistoryManager.instance.with_context(name: :irb) do
+      framework.history_manager.with_context(name: :irb) do
         Rex::Ui::Text::IrbShell.new(self).run
       end
     else
@@ -585,7 +585,7 @@ def cmd_pry(*args)
     print_status('Starting Pry shell...')
     print_status("You are in the \"self\" (session) object\n")
     Pry.config.history_load = false
-    Rex::Ui::Text::Shell::HistoryManager.instance.with_context(history_file: Msf::Config.pry_history, name: :pry) do
+    framework.history_manager.with_context(history_file: Msf::Config.pry_history, name: :pry) do
       self.pry
     end
   end
@@ -746,7 +746,7 @@ def process_autoruns(datastore)
   # shell_write instead of operating on rstream directly.
   def _interact
     framework.events.on_session_interact(self)
-    Rex::Ui::Text::Shell::HistoryManager.instance.with_context(name: self.type.to_sym) {
+    framework.history_manager.with_context(name: self.type.to_sym) {
       _interact_stream
     }
   end

diff --git a/lib/msf/core/framework.rb b/lib/msf/core/framework.rb
@@ -62,6 +62,8 @@ def initialize(options={})
     # Allow specific module types to be loaded
     types = options[:module_types] || Msf::MODULE_TYPES
 
+    self.history_manager = Rex::Ui::Text::Shell::HistoryManager.new
+
     self.features = FeatureManager.instance
     self.features.load_config
 
@@ -190,9 +192,13 @@ def version
   #
   # The framework instance's feature manager. The feature manager is responsible
   # for configuring feature flags that can change characteristics of framework.
-  #
+  # @return [Msf::FeatureManager]
   attr_reader   :features
 
+  # The framework instance's history manager, responsible for managing command history
+  # in different contexts
+  # @return [Rex::Ui::Text::Shell::HistoryManager]
+  attr_reader :history_manager
 
   #
   # The framework instance's data service proxy
@@ -281,7 +287,8 @@ def eicar_corrupted?
   attr_writer   :db # :nodoc:
   attr_writer   :browser_profiles # :nodoc:
   attr_writer   :analyze # :nodoc:
-  attr_writer   :features # :nodoc:
+  attr_writer   :features  # :nodoc:
+  attr_writer   :history_manager  # :nodoc:
 
   private
 

diff --git a/lib/msf/ui/console/command_dispatcher/developer.rb b/lib/msf/ui/console/command_dispatcher/developer.rb
@@ -129,7 +129,7 @@ def cmd_irb(*args)
     if expressions.empty?
       print_status('Starting IRB shell...')
 
-      Rex::Ui::Text::Shell::HistoryManager.instance.with_context(name: :irb) do
+      framework.history_manager.with_context(name: :irb) do
         begin
           if active_module
             print_status("You are in #{active_module.fullname}\n")
@@ -192,7 +192,7 @@ def cmd_pry(*args)
     print_status('Starting Pry shell...')
 
     Pry.config.history_load = false
-    Rex::Ui::Text::Shell::HistoryManager.instance.with_context(history_file: Msf::Config.pry_history, name: :pry) do
+    framework.history_manager.with_context(history_file: Msf::Config.pry_history, name: :pry) do
       if active_module
         print_status("You are in the \"#{active_module.fullname}\" module object\n")
         active_module.pry
@@ -420,7 +420,7 @@ def cmd__historymanager(*args)
     end
 
     if opts.key?(:debug)
-      Rex::Ui::Text::Shell::HistoryManager.instance._debug = opts[:debug]
+      framework.history_manager._debug = opts[:debug]
       print_status("HistoryManager debugging is now #{opts[:debug] ? 'on' : 'off'}")
     end
 
@@ -430,7 +430,7 @@ def cmd__historymanager(*args)
         'Indent'  => 1,
         'Columns' => ['Id', 'File', 'Name']
       )
-      Rex::Ui::Text::Shell::HistoryManager.instance._contexts.each.with_index do |context, id|
+      framework.history_manager._contexts.each.with_index do |context, id|
         table << [id, context[:history_file], context[:name]]
       end
 

diff --git a/lib/rex/post/meterpreter/ui/console/command_dispatcher/core.rb b/lib/rex/post/meterpreter/ui/console/command_dispatcher/core.rb
@@ -599,7 +599,7 @@ def cmd_irb(*args)
     if expressions.empty?
       print_status('Starting IRB shell...')
       print_status("You are in the \"client\" (session) object\n")
-      Rex::Ui::Text::Shell::HistoryManager.instance.with_context(name: :irb) do
+      framework.history_manager.with_context(name: :irb) do
         Rex::Ui::Text::IrbShell.new(client).run
       end
     else
@@ -639,7 +639,7 @@ def cmd_pry(*args)
     print_status("You are in the \"client\" (session) object\n")
 
     Pry.config.history_load = false
-    Rex::Ui::Text::Shell::HistoryManager.instance.with_context(history_file: Msf::Config.pry_history, name: :pry) do
+    framework.history_manager.with_context(history_file: Msf::Config.pry_history, name: :pry) do
       client.pry
     end
   end

diff --git a/lib/rex/ui/text/shell.rb b/lib/rex/ui/text/shell.rb
@@ -130,7 +130,7 @@ def run(&block)
       # Pry is a development dependency, if not available suppressing history_load can be safely ignored.
     end
 
-    HistoryManager.instance.with_context(history_file: histfile, name: name) do
+    framework.history_manager.with_context(history_file: histfile, name: name) do
       self.hist_last_saved = Readline::HISTORY.length
 
       begin
@@ -177,7 +177,7 @@ def run(&block)
       end
     end
   ensure
-    HistoryManager.instance.flush
+    framework.history_manager.flush
     self.hist_last_saved = Readline::HISTORY.length
   end
 

diff --git a/lib/rex/ui/text/shell/history_manager.rb b/lib/rex/ui/text/shell/history_manager.rb
@@ -9,15 +9,13 @@ module Shell
 
 class HistoryManager
 
-  include Singleton
-
   MAX_HISTORY = 2000
 
   def initialize
     @contexts = []
-    @write_mutex = Mutex.new
-    @write_queue = {}
     @debug = false
+    @write_queue = ::Queue.new
+    @currently_processing = ::Queue.new
   end
 
   # Create a new history command context when executing the given block
@@ -40,7 +38,9 @@ def with_context(history_file: nil, name: nil, &block)
 
   # Flush the contents of the write queue to disk. Blocks synchronously.
   def flush
-    sleep 0.1 until @write_queue.empty?
+    until @write_queue.empty? && @currently_processing.empty?
+      sleep 0.1
+    end
 
     nil
   end
@@ -99,16 +99,9 @@ def load_history_file(history_file)
     return unless readline_available?
 
     clear_readline
-    commands = from_storage_queue(history_file)
-    if commands
-      commands.reverse.each do |c|
-        ::Readline::HISTORY << c
-      end
-    else
-      if File.exist?(history_file)
-        File.readlines(history_file).each do |e|
-          ::Readline::HISTORY << e.chomp
-        end
+    if File.exist?(history_file)
+      File.readlines(history_file).each do |e|
+        ::Readline::HISTORY << e.chomp
       end
     end
   end
@@ -140,32 +133,28 @@ def switch_context(new_context, old_context=nil)
   end
 
   def write_history_file(history_file, cmds)
-    entry_added = false
-    until entry_added
-      @write_mutex.synchronize do
-        if @write_queue[history_file].nil?
-          @write_queue[history_file] = cmds
-          entry_added = true
+    write_queue_ref = @write_queue
+    currently_processing_ref = @currently_processing
+    @write_thread ||= Rex::ThreadFactory.spawn("HistoryManagerWriter", false) do
+      while (event = write_queue_ref.pop)
+        begin
+          currently_processing_ref << event
+
+          history_file = event[:history_file]
+          cmds = event[:cmds]
+
+          File.open(history_file, 'wb+') do |f|
+            f.puts(cmds.reverse)
+          end
+        rescue => e
+          elog(e)
+        ensure
+          currently_processing_ref.pop
         end
       end
-      sleep 0.1 if !entry_added
-    end
-
-    Rex::ThreadFactory.spawn("#{history_file} Writer", false) do
-      File.open(history_file, 'wb+') do |f|
-        f.puts(cmds.reverse)
-      end
-
-      @write_mutex.synchronize do
-        @write_queue.delete(history_file)
-      end
     end
-  end
 
-  def from_storage_queue(history_file)
-    @write_mutex.synchronize do
-      @write_queue[history_file]
-    end
+    write_queue_ref << { type: :write, history_file: history_file, cmds: cmds }
   end
 end