From 2033cd9f3e290c8a420dbbe3e7a9b64487451727 Mon Sep 17 00:00:00 2001
From: Justin Collins <justin@presidentbeef.com>
Date: Tue, 19 Jan 2021 23:10:56 -0800
Subject: [PATCH] Do dataflow on all of 'library' code

Previously, it was only on method bodies. This probably won't be a huge
change, since most code is inside methods.

As a result of this change, now all classes will have instance variable
values propagated from their initializers, not just 'libraries'.
---
 lib/brakeman/processor.rb                           |  3 ++-
 lib/brakeman/processors/alias_processor.rb          | 11 +++++++++++
 lib/brakeman/processors/library_processor.rb        |  9 ---------
 .../config/initializers/deeper_dir/constants.rb     |  1 +
 test/apps/rails6/lib/run_stuff.rb                   |  8 ++++++++
 test/tests/rails6.rb                                | 13 +++++++++++++
 6 files changed, 35 insertions(+), 10 deletions(-)
 create mode 100644 test/apps/rails6/config/initializers/deeper_dir/constants.rb

diff --git a/lib/brakeman/processor.rb b/lib/brakeman/processor.rb
index fd4fa75f77..c3c48d6627 100644
--- a/lib/brakeman/processor.rb
+++ b/lib/brakeman/processor.rb
@@ -95,7 +95,8 @@ def process_initializer file_name, src
 
     #Process source for a library file
     def process_lib src, file_name
-      LibraryProcessor.new(@tracker).process_library src, file_name
+      result = LibraryProcessor.new(@tracker).process_library src, file_name
+      AliasProcessor.new(@tracker, file_name).process result if result
     end
   end
 end
diff --git a/lib/brakeman/processors/alias_processor.rb b/lib/brakeman/processors/alias_processor.rb
index 51b9b83377..5513586812 100644
--- a/lib/brakeman/processors/alias_processor.rb
+++ b/lib/brakeman/processors/alias_processor.rb
@@ -31,6 +31,7 @@ def initialize tracker = nil, current_file = nil
     @helper_method_info = Hash.new({})
     @or_depth_limit = (tracker && tracker.options[:branch_limit]) || 5 #arbitrary default
     @meth_env = nil
+    @initializer_envs = {}
     @current_file = current_file
     set_env_defaults
   end
@@ -432,7 +433,17 @@ def process_block exp
   #Process a method definition.
   def process_defn exp
     meth_env do
+      unless exp.method_name == :initialize
+        env.current.merge! @initializer_envs.fetch(@current_class, {}).to_h
+      end
+
       exp.body = process_all! exp.body
+
+      if @current_class and exp.method_name == :initialize
+        # Call Env#current to get just a hash of
+        # instance variables and their values
+        @initializer_envs[@current_class] = only_ivars.current
+      end
     end
     exp
   end
diff --git a/lib/brakeman/processors/library_processor.rb b/lib/brakeman/processors/library_processor.rb
index 51568c8241..fd8744c200 100644
--- a/lib/brakeman/processors/library_processor.rb
+++ b/lib/brakeman/processors/library_processor.rb
@@ -30,15 +30,6 @@ def process_module exp
   end
 
   def process_defn exp
-    if exp.method_name == :initialize
-      @alias_processor.process_safely exp.body_list
-      @initializer_env = @alias_processor.only_ivars
-    elsif node_type? exp, :defn
-      exp = @alias_processor.process_safely exp, @initializer_env
-    else
-      exp = @alias_processor.process exp
-    end
-
     if @current_class
       exp.body = process_all! exp.body
       @current_class.add_method :public, exp.method_name, exp, @current_file
diff --git a/test/apps/rails6/config/initializers/deeper_dir/constants.rb b/test/apps/rails6/config/initializers/deeper_dir/constants.rb
new file mode 100644
index 0000000000..3e043e1df0
--- /dev/null
+++ b/test/apps/rails6/config/initializers/deeper_dir/constants.rb
@@ -0,0 +1 @@
+SOME_CONSTANT = '1'
diff --git a/test/apps/rails6/lib/run_stuff.rb b/test/apps/rails6/lib/run_stuff.rb
index 6d71b211f8..8cba7b7b80 100644
--- a/test/apps/rails6/lib/run_stuff.rb
+++ b/test/apps/rails6/lib/run_stuff.rb
@@ -4,4 +4,12 @@ def run
       `cat #{temp_file.path}`
     end
   end
+
+  RUN_THINGS = {
+    SOME_CONSTANT => "ASafeString"
+  }
+
+  def use_group_things
+    RUN_THINGS[params[:key]].constantize.new
+  end
 end
diff --git a/test/tests/rails6.rb b/test/tests/rails6.rb
index cecb27a21b..faf2c6b15f 100644
--- a/test/tests/rails6.rb
+++ b/test/tests/rails6.rb
@@ -251,6 +251,19 @@ def test_remote_code_execution_not_query_parameters
       :user_input => s(:call, s(:call, s(:const, :User), :first), :some_method_thing)
   end
 
+  def test_constants_in_libraries
+    assert_no_warning :type => :warning,
+      :warning_code => 24,
+      :fingerprint => "bf29c627bcf090125caf10c73e33cd3f7b66f5f1a2d5ba03b418ed51e3e15b72",
+      :warning_type => "Remote Code Execution",
+      :line => 13,
+      :message => /^Unsafe\ reflection\ method\ `constantize`\ c/,
+      :confidence => 1,
+      :relative_path => "lib/run_stuff.rb",
+      :code => s(:call, s(:call, s(:const, :RUN_THINGS), :[], s(:call, s(:params), :[], s(:lit, :key))), :constantize),
+      :user_input => s(:call, s(:params), :[], s(:lit, :key))
+  end
+
   def test_safe_yaml_load_option
     assert_no_warning :type => :warning,
       :warning_code => 25,