Merge pull request #407 from presidentbeef/fix_multiple_block_related_issues

Fix multiple block related issues

lib/brakeman/processors/alias_processor.rb

@@ -220,13 +220,24 @@ class Brakeman::AliasProcessor < Brakeman::SexpProcessor
 
   #Process a method definition.
   def process_methdef exp
-    env.scope do
-      set_env_defaults
+    meth_env do
       exp.body = process_all! exp.body
     end
     exp
   end
 
+  def meth_env
+    begin
+      env.scope do
+        set_env_defaults
+        @meth_env = env.current
+        yield
+      end
+    ensure
+      @meth_env = nil
+    end
+  end
+
   #Process a method definition on self.
   def process_selfdef exp
     env.scope do
@@ -437,9 +448,11 @@ class Brakeman::AliasProcessor < Brakeman::SexpProcessor
       branch_scopes = []
       exps.each_with_index do |branch, i|
         scope do
+          @branch_env = env.current
           branch_index = 2 + i # s(:if, condition, then_branch, else_branch)
           exp[branch_index] = process_if_branch branch
           branch_scopes << env.current
+          @branch_env = nil
         end
       end
 
@@ -731,7 +744,17 @@ class Brakeman::AliasProcessor < Brakeman::SexpProcessor
     end
 
     if @ignore_ifs or not @inside_if
-      env[var] = value
+      if @meth_env and node_type? var, :ivar and env[var].nil?
+        @meth_env[var] = value
+      else
+        env[var] = value
+      end
+    elsif env.current[var]
+      env.current[var] = value
+    elsif @branch_env and @branch_env[var]
+      @branch_env[var] = value
+    elsif @branch_env and @meth_env and node_type? var, :ivar
+      @branch_env[var] = value
     else
       env.current[var] = value
     end
@@ -776,5 +799,4 @@ class Brakeman::AliasProcessor < Brakeman::SexpProcessor
       false
     end
   end
-
 end

lib/brakeman/processors/controller_alias_processor.rb

@@ -84,9 +84,7 @@ class Brakeman::ControllerAliasProcessor < Brakeman::AliasProcessor
     @current_method = meth_name
     @rendered = false if is_route
 
-    env.scope do
-      set_env_defaults
-
+    meth_env do
       if is_route
         before_filter_list(@current_method, @current_class).each do |f|
           process_before_filter f
@@ -124,7 +122,7 @@ class Brakeman::ControllerAliasProcessor < Brakeman::AliasProcessor
 
   #Check for +respond_to+
   def process_call_with_block exp
-    process_default exp
+    super
 
     if call? exp.block_call and exp.block_call.method == :respond_to
       @rendered = true

test/apps/rails4/app/controllers/friendly_controller.rb

 class FriendlyController
- 
+  some_helper_thing do
+    @user = User.current_user
+  end
+
   def find
     @user = User.friendly.find(params[:id])
     redirect_to @user
   end
-  
-end
\ No newline at end of file
+
+  def some_user_thing
+    redirect_to @user.url
+  end
+end

test/tests/alias_processor.rb

 class AliasProcessorTests < Test::Unit::TestCase
-  def assert_alias expected, original
+  def assert_alias expected, original, full = false
     original_sexp = RubyParser.new.parse original
     expected_sexp = RubyParser.new.parse expected
-
     processed_sexp = Brakeman::AliasProcessor.new.process_safely original_sexp
-    result = processed_sexp.last
 
-    assert_equal expected_sexp, result
+    if full
+      assert_equal expected_sexp, processed_sexp
+    else
+      assert_equal expected_sexp, processed_sexp.last
+    end
+  end
+
+  def assert_output input, output
+    assert_alias output, input, true
   end
 
   def test_addition
@@ -377,4 +383,146 @@ class AliasProcessorTests < Test::Unit::TestCase
     y
     RUBY
   end
+
+  def test_block_with_local
+    assert_output <<-INPUT, <<-OUTPUT
+      def a
+        if b
+          c = nil
+          ds.each do |d|
+            e = T.new
+            c = e.map
+          end
+
+          r("f" + c.name)
+        else
+          g
+        end
+      end
+    INPUT
+      def a
+        if b
+          c = nil
+          ds.each do |d|
+            e = T.new
+            c = T.new.map
+          end
+
+          r("f" + T.new.map.name)
+        else
+          g
+        end
+      end
+    OUTPUT
+  end
+
+  def test_block_in_class_scope
+    # Make sure blocks in class do not mess up instance variable scope
+    # for subsequent methods
+    assert_output <<-INPUT, <<-OUTPUT
+      class A
+        x do
+          @a = 1
+        end
+
+        def b
+          @a
+        end
+      end
+    INPUT
+      class A
+        x do
+          @a = 1
+        end
+
+        def b
+          @a
+        end
+      end
+    OUTPUT
+  end
+
+  def test_instance_method_scope_in_block
+    # Make sure instance variables set inside blocks are set at the method
+    # scope
+    assert_output <<-INPUT, <<-OUTPUT
+      class A
+        def b
+          x do
+            @a = 1
+          end
+
+          @a
+        end
+      end
+    INPUT
+      class A
+        def b
+          x do
+            @a = 1
+          end
+
+          1
+        end
+      end
+    OUTPUT
+  end
+
+  def test_instance_method_scope_in_if_with_blocks
+    # Make sure instance variables set inside if expressions are set at the
+    # method scope after being combined
+    assert_output <<-INPUT, <<-OUTPUT
+      class A
+        def b
+          if something
+            x do
+              @a = 1
+            end
+          else
+            y do
+              @a = 2
+            end
+          end
+
+          @a
+        end
+      end
+    INPUT
+      class A
+        def b
+          if something
+            x do
+              @a = 1
+            end
+          else
+            y do
+              @a = 2
+            end
+          end
+
+          (1 or 2)
+        end
+      end
+    OUTPUT
+  end
+
+  def test_branch_env_is_closed_after_if_statement
+    assert_output <<-'INPUT', <<-'OUTPUT'
+      def a
+        if b
+          return unless c # this was causing problems
+          @d = D.find(1)
+          @d
+        end
+      end
+    INPUT
+      def a
+        if b
+          return unless c
+          @d = D.find(1)
+          D.find(1)
+       end
+      end
+    OUTPUT
+  end
 end

test/tests/rails4.rb

@@ -110,4 +110,15 @@ class Rails4Tests < Test::Unit::TestCase
       :confidence => 0,
       :relative_path => "app/controllers/application_controller.rb"
   end
+
+  def test_redirect_with_instance_variable_from_block
+    assert_no_warning :type => :warning,
+      :warning_code => 18,
+      :fingerprint => "e024f0cf67432409ec4afc80216fb2f6c9929fbbd32c2421e8867cd254f22d04",
+      :warning_type => "Redirect",
+      :line => 12,
+      :message => /^Possible\ unprotected\ redirect/,
+      :confidence => 0,
+      :relative_path => "app/controllers/friendly_controller.rb"
+  end
 end