Add tests for instance variable scoping

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

@@ -88,4 +88,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