Land #12960, add ttl to job results instantiated from an RPC request

lib/msf/base/simple/auxiliary.rb

@@ -69,7 +69,7 @@ def self.run_simple(omod, opts = {}, &block)
     end
 
     run_uuid = Rex::Text.rand_text_alphanumeric(24)
-    mod.framework.ready << run_uuid
+    mod.framework.job_state_tracker.waiting run_uuid
     ctx = [mod, run_uuid]
     if(mod.passive? or opts['RunAsJob'])
       mod.job_id = mod.framework.jobs.start_bg_job(
@@ -122,7 +122,7 @@ def self.check_simple(mod, opts)
 
 
     run_uuid = Rex::Text.rand_text_alphanumeric(24)
-    mod.framework.ready << run_uuid
+    mod.framework.job_state_tracker.waiting run_uuid
     ctx = [mod, run_uuid]
 
     if opts['RunAsJob']
@@ -169,15 +169,12 @@ def self.job_run_proc(ctx, &block)
       mod.setup
       mod.framework.events.on_module_run(mod)
       begin
-        mod.framework.running << run_uuid
-        mod.framework.ready.delete run_uuid
+        mod.framework.job_state_tracker.start run_uuid
         result = block.call(mod)
-        mod.framework.results[run_uuid] = {result: result}
+        mod.framework.job_state_tracker.completed(run_uuid, result)
       rescue ::Exception => e
-        mod.framework.results[run_uuid] = {error: e.to_s}
+        mod.framework.job_state_tracker.failed(run_uuid, e)
         raise
-      ensure
-        mod.framework.running.delete run_uuid
       end
     rescue Msf::Auxiliary::Complete
       mod.cleanup
@@ -194,7 +191,7 @@ def self.job_run_proc(ctx, &block)
       return
     rescue ::Interrupt => e
       mod.error = e
-      mod.print_error("Stopping running againest current target...")
+      mod.print_error("Stopping running against current target...")
       mod.cleanup
       mod.print_status("Control-C again to force quit all targets.")
       begin

lib/msf/base/simple/exploit.rb

@@ -190,7 +190,7 @@ def self.check_simple(mod, opts)
     mod.validate
 
     run_uuid = Rex::Text.rand_text_alphanumeric(24)
-    mod.framework.ready << run_uuid
+    mod.framework.job_state_tracker.waiting run_uuid
     ctx = [mod, run_uuid]
 
     if opts['RunAsJob']
@@ -220,15 +220,12 @@ def self.job_check_proc(ctx)
     run_uuid = ctx[1]
     begin
       mod.setup
-      mod.framework.running << run_uuid
-      mod.framework.ready.delete run_uuid
+      mod.framework.job_state_tracker.start run_uuid
       result = mod.has_check? ? mod.check : Msf::Exploit::CheckCode::Unsupported
-      mod.framework.results[run_uuid] = {result: result}
+      mod.framework.job_state_tracker.completed(run_uuid, result)
     rescue => e
-      mod.framework.results[run_uuid] = {error: e.to_s}
+      mod.framework.job_state_tracker.failed(run_uuid, e)
       mod.handle_exception e
-    ensure
-      mod.framework.running.delete run_uuid
     end
 
     return result

lib/msf/base/simple/framework.rb

@@ -1,6 +1,7 @@
 # -*- coding: binary -*-
 require 'msf/base/simple'
 require 'msf/base/simple/framework/module_paths'
+require 'msf/base/simple/job_state_tracker'
 
 module Msf
 module Simple
@@ -150,9 +151,7 @@ def self.simplify_module(instance, load_saved_config = true)
   #
   def init_simplified
     self.stats = Statistics.new(self)
-    self.ready = Set.new
-    self.running = Set.new
-    self.results = Hash.new
+    self.job_state_tracker = JobStateTracker.new
   end
 
   #
@@ -174,41 +173,29 @@ def save_config
   #
   attr_reader :stats
 
-  #
-  # Boolean indicating whether the cache is initialized yet
-  #
-  attr_reader :cache_initialized
 
   #
-  # Thread of the running rebuild operation
+  # JobStateTracker
   #
-  attr_reader :cache_thread
+  attr_reader :job_state_tracker
 
   #
-  # {Set<String>} of module run/check UUIDs waiting to be kicked off
-  #
-  attr_reader :ready
-
-  #
-  # {Hash<String,Hash>} of module run/check results, by UUID. Successful runs
-  # look like `{result: check_code}` and errors like `{error: message}`.
+  # Boolean indicating whether the cache is initialized yet
   #
-  attr_reader :results
+  attr_reader :cache_initialized
 
   #
-  # {Set<String>} of module run/check UUIDs currently in progress
+  # Thread of the running rebuild operation
   #
-  attr_reader :running
+  attr_reader :cache_thread
   attr_writer :cache_initialized # :nodoc:
   attr_writer :cache_thread # :nodoc:
 
 
 protected
 
-  attr_writer :ready # :nodoc:
-  attr_writer :results # :nodoc:
-  attr_writer :running # :nodoc:
   attr_writer :stats # :nodoc:
+  attr_writer :job_state_tracker # :nodoc:
 
 end
 

lib/msf/base/simple/job_state_tracker.rb

@@ -0,0 +1,84 @@
+require 'monitor'
+
+class JobStateTracker
+
+  include MonitorMixin
+
+  def initialize(result_ttl=nil)
+    self.ready = Set.new
+    self.running = Set.new
+    # Can be expanded upon later to allow the option of a MemCacheStore being backed by redis for example
+    self.results = ResultsMemoryStore.new(expires_in: result_ttl || 5.minutes)
+  end
+
+  def waiting(id)
+    ready << id
+  end
+
+  def start(id)
+    running << id
+    ready.delete(id)
+  end
+
+  def completed(id, result, ttl=nil)
+    begin
+      # ttl of nil means it will take the default expiry time
+      results.write(id, {result: result}, ttl)
+    ensure
+      running.delete(id)
+    end
+  end
+
+  def failed(id, error, ttl=nil)
+    begin
+      # ttl of nil means it will take the default expiry time
+      results.write(id, {error: error.to_s}, ttl)
+    ensure
+      running.delete(id)
+    end
+  end
+
+  def running?(id)
+    running.include? id
+  end
+
+  def waiting?(id)
+    ready.include? id
+  end
+
+  def finished?(id)
+    results.exist? id
+  end
+
+  def result(id)
+    results.fetch(id)
+  end
+
+  def delete(id)
+    results.delete(id)
+  end
+
+  def results_size
+    results.size
+  end
+
+  def waiting_size
+    ready.size
+  end
+
+  def running_size
+    running.size
+  end
+
+  alias :ack :delete
+
+  private
+
+  attr_accessor :ready, :running, :results
+
+  class ResultsMemoryStore < ActiveSupport::Cache::MemoryStore
+    def size
+      @data.size
+    end
+  end
+end

lib/msf/core/rpc/v10/rpc_base.rb

@@ -9,10 +9,10 @@ class RPC_Base
   #
   # return [void]
   def initialize(service)
-    self.service   = service
-    self.framework = service.framework
-    self.tokens    = service.tokens
-    self.users     = service.users
+    self.service           = service
+    self.framework         = service.framework
+    self.tokens            = service.tokens
+    self.users             = service.users
   end
 
   # Raises an Msf::RPC Exception.

lib/msf/core/rpc/v10/rpc_module.rb

@@ -412,9 +412,9 @@ def rpc_target_compatible_evasion_payloads(mname, target)
   #  rpc.call('module.running_stats')
   def rpc_running_stats
     {
-        "ready" => self.framework.ready.size,
-        "running" => self.framework.running.size,
-        "results" => self.framework.results.size,
+        "waiting" => self.framework.job_state_tracker.waiting_size,
+        "running" => self.framework.job_state_tracker.running_size,
+        "results" => self.framework.job_state_tracker.results_size,
     }
   end
 
@@ -523,7 +523,7 @@ def rpc_check(mtype, mname, opts)
   # TODO: expand these to take a list of UUIDs or stream with event data if
   # required for performance
   def rpc_results(uuid)
-    if r = self.framework.results[uuid]
+    if (r = self.framework.job_state_tracker.result(uuid))
       if r[:error]
         {"status" => "errored", "error" => r[:error]}
       else
@@ -537,17 +537,17 @@ def rpc_results(uuid)
           {"status" => "completed", "result" => r[:result]}
         end
       end
-    elsif self.framework.running.include? uuid
+    elsif self.framework.job_state_tracker.running? uuid
       {"status" => "running"}
-    elsif self.framework.ready.include? uuid
+    elsif self.framework.job_state_tracker.waiting uuid
       {"status" => "ready"}
     else
       error(404, "Results not found for module instance #{uuid}")
     end
   end
 
   def rpc_ack(uuid)
-    {"success" => !!self.framework.results.delete(uuid)}
+    {"success" => !!self.framework.job_state_tracker.ack(uuid)}
   end
 
   # Returns a list of executable format names.
@@ -740,7 +740,7 @@ def _run_exploit(mod, opts)
   end
 
   def _run_auxiliary(mod, opts)
-    uuid, job = Msf::Simple::Auxiliary.run_simple(mod, {
+    uuid, job = Msf::Simple::Auxiliary.run_simple(mod,{
       'Action'   => opts['ACTION'],
       'RunAsJob' => true,
       'Options'  => opts
@@ -752,7 +752,7 @@ def _run_auxiliary(mod, opts)
   end
 
   def _check_exploit(mod, opts)
-    uuid, job = Msf::Simple::Exploit.check_simple(mod, {
+    uuid, job = Msf::Simple::Exploit.check_simple(mod,{
         'RunAsJob' => true,
         'Options'  => opts
     })
@@ -763,7 +763,7 @@ def _check_exploit(mod, opts)
   end
 
   def _check_auxiliary(mod, opts)
-    uuid, job = Msf::Simple::Auxiliary.check_simple(mod, {
+    uuid, job = Msf::Simple::Auxiliary.check_simple(mod,{
         'Action'   => opts['ACTION'],
         'RunAsJob' => true,
         'Options'  => opts

lib/msf/core/rpc/v10/service.rb

@@ -15,7 +15,6 @@
 require 'msf/core/rpc/v10/rpc_job'
 require 'msf/core/rpc/v10/rpc_db'
 
-
 module Msf
 module RPC