Add call site indexing for faster searching

lib/brakeman/call_index.rb

+require 'set'
+
+class Brakeman::CallIndex
+  def initialize calls
+    @calls_by_method = Hash.new { |h,k| h[k] = [] }
+    @calls_by_target = Hash.new { |h,k| h[k] = [] }
+    @methods = Set.new
+    @targets = Set.new
+
+    index_calls calls
+  end
+
+  def find_calls options
+    target = options[:target] || options[:targets]
+    method = options[:method] || options[:methods]
+    nested = options[:nested]
+
+    #Find by narrowest category
+    if target and method and target.is_a? Array and method.is_a? Array
+      if target.length > method.length
+        calls = filter_by_target calls_by_methods(method), target
+      else
+        calls = calls_by_targets(target)
+        calls = filter_by_method calls, method
+      end
+
+    #Find by target, then by methods, if provided
+    elsif target
+      calls = calls_by_target target
+
+      if calls and method
+        calls = filter_by_method calls, method
+      end
+
+    #Find calls with no explicit target
+    #with either :target => nil or :target => false
+    elsif options.key? :target and not target and method
+      calls = calls_by_method method
+      calls = filter_by_target calls, nil
+
+    #Find calls by method
+    elsif method
+      calls = calls_by_method method
+    else
+      warn "Invalid arguments to CallCache#find_calls: #{options.inspect}"
+    end
+
+    return [] if calls.nil?
+
+    #Remove calls that are actually targets of other calls
+    #Unless those are explicitly desired
+    calls = filter_nested calls unless nested
+
+    calls
+  end
+
+  private
+
+  def index_calls calls
+    calls.each do |call|
+      @methods << call[:method].to_s
+      @targets << call[:target].to_s
+      @calls_by_method[call[:method]] << call
+      @calls_by_target[call[:target]] << call
+    end
+  end
+
+  def calls_by_target target
+    if target.is_a? Array
+      calls_by_targets target
+    elsif target.is_a? Regexp
+      targets = @targets.select do |t|
+        t.match target
+      end
+
+      if targets.empty?
+        []
+      elsif targets.length > 1
+        calls_by_targets targets
+      else
+        calls_by_target[targets.first]
+      end
+    else
+      @calls_by_target[target]
+    end
+  end
+
+  def calls_by_targets targets
+    calls = []
+
+    targets.each do |target|
+      calls.concat @calls_by_target[target] if @calls_by_target.key? target
+    end
+
+    calls
+  end
+
+  def calls_by_method method
+    if method.is_a? Array
+      calls_by_methods method
+    elsif method.is_a? Regexp
+      methods = @methods.select do |m|
+        m.match method
+      end
+
+      if methods.empty?
+        []
+      elsif methods.length > 1
+        calls_by_methods methods
+      else
+        calls_by_method[methods.first]
+      end
+    else
+      @calls_by_method[method]
+    end
+  end
+
+  def calls_by_methods methods
+    calls = []
+
+    methods.each do |method|
+      calls.concat @calls_by_method[method] if @calls_by_method.key? method
+    end
+
+    calls
+  end
+
+  def calls_with_no_target
+    @calls_by_target[nil]
+  end
+
+  def filter calls, key, value
+    if value.is_a? Array
+      values = Set.new value
+
+      calls.select do |call|
+        values.include? call[key]
+      end
+    elsif value.is_a? Regexp
+      calls.select do |call|
+        call[key].to_s.match value
+      end
+    else
+      calls.select do |call|
+        call[key] == value
+      end
+    end
+  end
+
+  def filter_by_method calls, method
+    filter calls, :method, method
+  end
+
+  def filter_by_target calls, target
+    filter calls, :target, target
+  end
+
+  def filter_nested calls
+    filter calls, :nested, false
+  end
+end

lib/brakeman/checks/base_check.rb

@@ -30,17 +30,18 @@ class Brakeman::BaseCheck < SexpProcessor
 
   #Add result to result list, which is used to check for duplicates
   def add_result result, location = nil
-    location ||= (@current_template && @current_template[:name]) || @current_class || @current_module || @current_set || result[1]
+    location ||= (@current_template && @current_template[:name]) || @current_class || @current_module || @current_set || result[:location][1]
     location = location[:name] if location.is_a? Hash
     location = location.to_sym
 
-    if result? result
-      line = result[-1].original_line || result[-1].line
-    else
+    if result.is_a? Hash
+      line = result[:call].original_line || result[:call].line
+    elsif sexp? result
       line = result.original_line || result.line
+    else
+      raise ArgumentError
     end
 
-
     @results << [line, location, result]
   end
 
@@ -135,17 +136,19 @@ class Brakeman::BaseCheck < SexpProcessor
   #This is to avoid reporting duplicates. Checks if the result has been
   #reported already from the same line number.
   def duplicate? result, location = nil
-    if result? result
-      line = result[-1].original_line || result[-1].line
-    else
+    if result.is_a? Hash
+      line = result[:call].original_line || result[:call].line
+    elsif sexp? result
       line = result.original_line || result.line
+    else
+      raise ArgumentError
     end
 
-
-    location ||= (@current_template && @current_template[:name]) || @current_class || @current_module || @current_set || result[1]
+    location ||= (@current_template && @current_template[:name]) || @current_class || @current_module || @current_set || result[:location][1]
 
     location = location[:name] if location.is_a? Hash
     location = location.to_sym
+
     @results.each do |r|
       if r[0] == line and r[1] == location
         if tracker.options[:combine_locations]

lib/brakeman/checks/check_cross_site_scripting.rb

@@ -176,7 +176,7 @@ class Brakeman::CheckCrossSiteScripting < Brakeman::BaseCheck
           confidence = CONFIDENCE[:low]
         end
 
-        warn :template => @current_template, 
+        warn :template => @current_template,
           :warning_type => "Cross Site Scripting", 
           :message => message,
           :line => exp.line,
@@ -274,7 +274,7 @@ class Brakeman::CheckLinkTo < Brakeman::CheckCrossSiteScripting
   def run_check
     #Ideally, I think this should also check to see if people are setting
     #:escape => false
-    methods = tracker.find_call [], :link_to 
+    methods = tracker.find_call :target => false, :method => :link_to 
 
     @models = tracker.models.keys
     @inspect_arguments = tracker.options[:check_arguments]
@@ -287,7 +287,7 @@ class Brakeman::CheckLinkTo < Brakeman::CheckCrossSiteScripting
   def process_result result
     #Have to make a copy of this, otherwise it will be changed to
     #an ignored method call by the code above.
-    call = result[-1] = result[-1].dup
+    call = result[:call] = result[:call].dup
 
     @matched = false
 

lib/brakeman/checks/check_evaluation.rb

@@ -8,7 +8,7 @@ class Brakeman::CheckEvaluation < Brakeman::BaseCheck
   #Process calls
   def run_check
     debug_info "Finding eval-like calls"
-    calls = tracker.find_call nil, [:eval, :instance_eval, :class_eval, :module_eval]
+    calls = tracker.find_call :method => [:eval, :instance_eval, :class_eval, :module_eval]
 
     debug_info "Processing eval-like calls"
     calls.each do |call|
@@ -18,11 +18,11 @@ class Brakeman::CheckEvaluation < Brakeman::BaseCheck
 
   #Warns if result includes user input
   def process_result result
-    if include_user_input? result[-1]
+    if include_user_input? result[:call]
       warn :result => result,
         :warning_type => "Dangerous Eval",
         :message => "User input in eval",
-        :code => result[-1],
+        :code => result[:call],
         :confidence => CONFIDENCE[:high]
     end
   end

lib/brakeman/checks/check_execute.rb

@@ -18,17 +18,17 @@ class Brakeman::CheckExecute < Brakeman::BaseCheck
     check_for_backticks tracker
 
     debug_info "Finding other system calls"
-    calls = tracker.find_call [:IO, :Open3, :Kernel, []], [:exec, :popen, :popen3, :syscall, :system]
+    calls = tracker.find_call :targets => [:IO, :Open3, :Kernel, nil], :methods => [:exec, :popen, :popen3, :syscall, :system]
 
     debug_info "Processing system calls"
     calls.each do |result|
-      process result
+      process_result result
     end
   end
 
   #Processes results from FindCall.
-  def process_result exp
-    call = exp[-1]
+  def process_result result
+    call = result[:call]
 
     args = process call[3]
 
@@ -39,8 +39,8 @@ class Brakeman::CheckExecute < Brakeman::BaseCheck
       failure = include_user_input?(args) || include_interp?(args)
     end
 
-    if failure and not duplicate? call, exp[1]
-      add_result call, exp[1]
+    if failure and not duplicate? result
+      add_result result
 
       if @string_interp
         confidence = CONFIDENCE[:med]
@@ -48,15 +48,13 @@ class Brakeman::CheckExecute < Brakeman::BaseCheck
         confidence = CONFIDENCE[:high]
       end
 
-      warn :result => exp,
+      warn :result => result,
         :warning_type => "Command Injection", 
         :message => "Possible command injection",
         :line => call.line,
         :code => call,
         :confidence => confidence
     end
-
-    exp
   end
 
   #Looks for calls using backticks such as

lib/brakeman/checks/check_file_access.rb

@@ -7,13 +7,13 @@ class Brakeman::CheckFileAccess < Brakeman::BaseCheck
 
   def run_check
     debug_info "Finding possible file access"
-    methods = tracker.find_call [:Dir, :File, :IO, :Kernel, :"Net::FTP", :"Net::HTTP", :PStore, :Pathname, :Shell, :YAML], [:[], :chdir, :chroot, :delete, :entries, :foreach, :glob, :install, :lchmod, :lchown, :link, :load, :load_file, :makedirs, :move, :new, :open, :read, :read_lines, :rename, :rmdir, :safe_unlink, :symlink, :syscopy, :sysopen, :truncate, :unlink]
+    methods = tracker.find_call :targets => [:Dir, :File, :IO, :Kernel, :"Net::FTP", :"Net::HTTP", :PStore, :Pathname, :Shell, :YAML], :methods => [:[], :chdir, :chroot, :delete, :entries, :foreach, :glob, :install, :lchmod, :lchown, :link, :load, :load_file, :makedirs, :move, :new, :open, :read, :read_lines, :rename, :rmdir, :safe_unlink, :symlink, :syscopy, :sysopen, :truncate, :unlink]
 
     debug_info "Finding calls to load()"
-    methods.concat tracker.find_call [], [:load]
+    methods.concat tracker.find_call :target => false, :method => :load
 
     debug_info "Finding calls using FileUtils"
-    methods.concat tracker.find_call(:FileUtils, nil)
+    methods.concat tracker.find_call :target => :FileUtils
 
     debug_info "Processing found calls"
     methods.each do |call|
@@ -22,13 +22,13 @@ class Brakeman::CheckFileAccess < Brakeman::BaseCheck
   end
 
   def process_result result
-    call = result[-1]
+    call = result[:call]
 
     file_name = call[3][1]
 
     if check = include_user_input?(file_name)
-      unless duplicate? call, result[1]
-        add_result call, result[1]
+      unless duplicate? result
+        add_result result
 
         if check == :params
           message = "Parameter"

lib/brakeman/checks/check_mail_to.rb

@@ -31,8 +31,8 @@ class Brakeman::CheckMailTo < Brakeman::BaseCheck
   def mail_to_javascript?
     debug_info "Checking calls to mail_to for javascript encoding"
 
-    tracker.find_call([], :mail_to).each do |result|
-      call = result[-1]
+    tracker.find_call(:target => false, :method => :mail_to).each do |result|
+      call = result[:call]
       args = call[-1]
 
       args.each do |arg|

lib/brakeman/checks/check_mass_assignment.rb

@@ -21,7 +21,7 @@ class Brakeman::CheckMassAssignment < Brakeman::BaseCheck
     @results = Set.new
 
     debug_info "Finding possible mass assignment calls on #{models.length} models"
-    calls = tracker.find_call models, [:new,
+    calls = tracker.find_call :targets => models, :methods => [:new,
       :attributes=, 
       :update_attribute, 
       :update_attributes, 
@@ -31,13 +31,13 @@ class Brakeman::CheckMassAssignment < Brakeman::BaseCheck
 
     debug_info "Processing possible mass assignment calls"
     calls.each do |result|
-      process result
+      process_result result
     end
   end
 
   #All results should be Model.new(...) or Model.attributes=() calls
   def process_result res
-    call = res[-1]
+    call = res[:call]
 
     check = check_call call
 

lib/brakeman/checks/check_quote_table_name.rb

@@ -32,6 +32,6 @@ class Brakeman::CheckQuoteTableName < Brakeman::BaseCheck
   def uses_quote_table_name?
     debug_info "Finding calls to quote_table_name()"
 
-    not tracker.find_call([], :quote_table_name).empty?
+    not tracker.find_call(:target => false, :method => :quote_table_name).empty?
   end
 end

lib/brakeman/checks/check_redirect.rb

@@ -12,13 +12,13 @@ class Brakeman::CheckRedirect < Brakeman::BaseCheck
   def run_check
     debug_info "Finding calls to redirect_to()"
 
-    @tracker.find_call(nil, :redirect_to).each do |c|
-      process c
+    @tracker.find_call(:target => false, :method => :redirect_to).each do |res|
+      process_result res
     end
   end
 
-  def process_result exp
-    call = exp[-1]
+  def process_result result
+    call = result[:call]
 
     method = call[2]
 
@@ -29,15 +29,13 @@ class Brakeman::CheckRedirect < Brakeman::BaseCheck
         confidence = CONFIDENCE[:low]
       end
 
-      warn :result => exp,
+      warn :result => result,
         :warning_type => "Redirect",
         :message => "Possible unprotected redirect",
         :line => call.line,
         :code => call,
         :confidence => confidence
     end
-
-    exp
   end
 
   #Custom check for user input. First looks to see if the user input

lib/brakeman/checks/check_send_file.rb

@@ -8,7 +8,7 @@ class Brakeman::CheckSendFile < Brakeman::CheckFileAccess
   def run_check
     debug_info "Finding all calls to send_file()"
 
-    methods = tracker.find_call nil, :send_file
+    methods = tracker.find_call :target => false, :method => :send_file
 
     methods.each do |call|
       process_result call

lib/brakeman/checks/check_sql.rb

@@ -16,23 +16,31 @@ class Brakeman::CheckSQL < Brakeman::BaseCheck
     @rails_version = tracker.config[:rails_version]
 
     debug_info "Finding possible SQL calls on models"
-    calls = tracker.find_model_find tracker.models.keys
+    if tracker.options[:rails3]
+      calls = tracker.find_call :targets => tracker.models.keys,
+        :methods => /^(find.*|first|last|all|where|order|group|having)$/,
+        :nested => true
+    else
+      calls = tracker.find_call :targets => tracker.models.keys,
+        :methods => /^(find.*|first|last|all)$/,
+        :nested => true
+    end
 
     debug_info "Finding possible SQL calls with no target"
-    calls.concat tracker.find_call([], /^(find.*|last|first|all|count|sum|average|minumum|maximum|count_by_sql)$/)
+    calls.concat tracker.find_call(:target => nil, :method => /^(find.*|last|first|all|count|sum|average|minumum|maximum|count_by_sql)$/)
 
     debug_info "Finding possible SQL calls using constantized()"
-    calls.concat tracker.find_model_find(nil).select { |result| constantize_call? result }
+    calls.concat tracker.find_call(:method => /^(find.*|last|first|all|count|sum|average|minumum|maximum|count_by_sql)$/).select { |result| constantize_call? result }
 
     debug_info "Processing possible SQL calls"
     calls.each do |c|
-      process c
+      process_result c
     end
   end
 
   #Process result from FindCall.
-  def process_result exp
-    call = exp[-1]
+  def process_result result
+    call = result[:call]
 
     args = process call[3]
 
@@ -44,8 +52,8 @@ class Brakeman::CheckSQL < Brakeman::BaseCheck
       failed = (args.length > 1 and check_arguments args[-1])
     end
 
-    if failed and not exp[-1].original_line and not duplicate? exp
-      add_result exp
+    if failed and not call.original_line and not duplicate? result
+      add_result result
 
       if include_user_input? args[-1]
         confidence = CONFIDENCE[:high]
@@ -53,7 +61,7 @@ class Brakeman::CheckSQL < Brakeman::BaseCheck
         confidence = CONFIDENCE[:med]
       end
 
-      warn :result => exp,
+      warn :result => result,
         :warning_type => "SQL Injection",
         :message => "Possible SQL injection",
         :confidence => confidence
@@ -66,12 +74,11 @@ class Brakeman::CheckSQL < Brakeman::BaseCheck
         confidence = CONFIDENCE[:low]
       end
 
-      warn :result => exp, 
+      warn :result => result, 
         :warning_type => "SQL Injection", 
         :message => "Upgrade to Rails >= 2.1.2 to escape :limit and :offset. Possible SQL injection",
         :confidence => confidence
     end
-    exp
   end
 
   private

lib/brakeman/checks/check_strip_tags.rb

@@ -26,6 +26,6 @@ class Brakeman::CheckStripTags < Brakeman::BaseCheck
   def uses_strip_tags?
     debug_info "Finding calls to strip_tags()"
 
-    not tracker.find_call([], :strip_tags).empty?
+    not tracker.find_call(:target => false, :method => :strip_tags).empty?
   end
 end

lib/brakeman/checks/check_translate_bug.rb

@@ -37,6 +37,6 @@ class Brakeman::CheckTranslateBug < Brakeman::BaseCheck
   def uses_translate?
     debug_info "Finding calls to translate() or t()"
 
-    not tracker.find_call([], [:t, :translate]).empty?
+    not tracker.find_call(:target => nil, :methods => [:t, :translate]).empty?
   end
 end

lib/brakeman/checks/check_without_protection.rb

@@ -24,7 +24,7 @@ class Brakeman::CheckWithoutProtection < Brakeman::BaseCheck
     @results = Set.new
 
     debug_info "Finding all mass assignments"
-    calls = tracker.find_call models, [:new,
+    calls = tracker.find_call :targets => models, :methods => [:new,
       :attributes=, 
       :update_attribute, 
       :update_attributes, 
@@ -34,13 +34,13 @@ class Brakeman::CheckWithoutProtection < Brakeman::BaseCheck
 
     debug_info "Processing all mass assignments"
     calls.each do |result|
-      process result
+      process_result result
     end
   end
 
   #All results should be Model.new(...) or Model.attributes=() calls
   def process_result res
-    call = res[-1]
+    call = res[:call]
     last_arg = call[3][-1]
 
     if hash? last_arg and not @results.include? call
@@ -66,7 +66,5 @@ class Brakeman::CheckWithoutProtection < Brakeman::BaseCheck
         end
       end
     end
-
-    res
   end
 end

lib/brakeman/processors/lib/find_all_calls.rb

+require 'brakeman/processors/base_processor'
+
+class Brakeman::FindAllCalls < Brakeman::BaseProcessor
+  attr_reader :calls
+
+  def initialize tracker
+    super
+    @current_class = nil
+    @current_method = nil
+    @in_target = false
+    @calls = []
+  end
+
+  #Process the given source. Provide either class and method being searched
+  #or the template. These names are used when reporting results.
+  def process_source exp, klass = nil, method = nil, template = nil
+    @current_class = klass
+    @current_method = method
+    @current_template = template
+    process exp
+  end
+
+  #Process body of method
+  def process_methdef exp
+    process exp[3]
+  end
+
+  #Process body of method
+  def process_selfdef exp
+    process exp[4]
+  end
+
+  #Process body of block
+  def process_rlist exp
+    exp[1..-1].each do |e|
+      process e
+    end
+
+    exp
+  end
+
+  def process_call exp
+    target = get_target exp[1]
+    
+    if call? target
+      already_in_target = @in_target
+      @in_target = true
+      process target
+      @in_target = already_in_target
+    end
+
+    method = exp[2]
+    process exp[3]
+    
+    if @current_template
+      @calls << { :target => target, :method => method, :call => exp, :nested => @in_target, :location => [:template, @current_template]}
+    else
+      @calls << { :target => target, :method => method, :call => exp, :nested => @in_target, :location => [:class, @current_class, @current_method] }
+    end
+
+    exp
+  end
+
+  #Process an assignment like a call
+  def process_attrasgn exp
+    process_call exp
+  end
+
+  private
+
+  #Gets the target of a call as a Symbol
+  #if possible
+  def get_target exp
+    if sexp? exp
+      case exp.node_type
+      when :ivar, :lvar, :const
+        exp[1]
+      when :true, :false
+        exp[0]
+      when :lit
+        exp[1]
+      when :colon2
+        class_name exp
+      else
+        exp
+      end
+    else
+      exp
+    end
+  end
+end

lib/brakeman/scanner.rb

@@ -65,6 +65,8 @@ class Brakeman::Scanner
     process_models
     warn "Processing controllers..."
     process_controllers
+    warn "Indexing call sites..."
+    index_call_sites
     tracker
   end
 
@@ -257,6 +259,10 @@ class Brakeman::Scanner
       end
     end
   end
+
+  def index_call_sites
+    tracker.index_call_sites
+  end
 end
 
 #This is from Rails 3 version of the Erubis handler

lib/brakeman/tracker.rb

 require 'set'
+require 'brakeman/call_index'
 require 'brakeman/checks'
 require 'brakeman/report'
 require 'brakeman/processors/lib/find_call'
+require 'brakeman/processors/lib/find_all_calls'
 require 'brakeman/processors/lib/find_model_call'
 
 #The Tracker keeps track of all the processed information.
@@ -41,6 +43,7 @@ class Brakeman::Tracker
     @checks = nil
     @processed = nil
     @template_cache = Set.new
+    @call_index = nil
   end
 
   #Add an error to the list. If no backtrace is given,
@@ -95,20 +98,9 @@ class Brakeman::Tracker
   end
 
   #Find a method call.
-  #
-  #See FindCall for details on arguments.
-  def find_call target, method
-    finder = Brakeman::FindCall.new target, method, self
-
-    self.each_method do |definition, set_name, method_name|
-      finder.process_source definition, set_name, method_name
-    end
-
-    self.each_template do |name, template|
-      finder.process_source template[:src], nil, nil, template
-    end
-
-    finder.matches
+  def find_call options
+    index_calls unless @call_index
+    @call_index.find_calls options
   end
 
   #Finds method call on models.
@@ -143,4 +135,18 @@ class Brakeman::Tracker
   def report
     Brakeman::Report.new(self)
   end
+
+  def index_call_sites
+    finder = Brakeman::FindAllCalls.new self
+
+    self.each_method do |definition, set_name, method_name|
+      finder.process_source definition, set_name, method_name
+    end
+
+    self.each_template do |name, template|
+      finder.process_source template[:src], nil, nil, template
+    end
+
+    @call_index = Brakeman::CallIndex.new finder.calls
+  end
 end

lib/brakeman/warning.rb

@@ -17,13 +17,13 @@ class Brakeman::Warning
 
     result = options[:result]
     if result
-      if result.length == 3 #template result
-        @template ||= result[1]
-        @code ||= result[2]
+      if result[:location][0] == :template #template result
+        @template ||= result[:location][1]
+        @code ||= result[:call]
       else
-        @class ||= result[1]
-        @method ||= result[2]
-        @code ||= result[3]
+        @class ||= result[:location][1]
+        @method ||= result[:location][2]
+        @code ||= result[:call]
       end
     end